diff --git a/backend/onyx/llm/utils.py b/backend/onyx/llm/utils.py index 5dccfd256b1e..9674eb4a4ae8 100644 --- a/backend/onyx/llm/utils.py +++ b/backend/onyx/llm/utils.py @@ -419,7 +419,7 @@ def _strip_colon_from_model_name(model_name: str) -> str: return ":".join(model_name.split(":")[:-1]) if ":" in model_name else model_name -def _find_model_obj(model_map: dict, provider: str, model_name: str) -> dict | None: +def find_model_obj(model_map: dict, provider: str, model_name: str) -> dict | None: stripped_model_name = _strip_extra_provider_from_model_name(model_name) model_names = [ @@ -537,7 +537,7 @@ def get_llm_max_tokens( return GEN_AI_MAX_TOKENS try: - model_obj = _find_model_obj( + model_obj = find_model_obj( model_map, model_provider, model_name, @@ -646,7 +646,7 @@ def get_max_input_tokens_from_llm_provider( def model_supports_image_input(model_name: str, model_provider: str) -> bool: model_map = get_model_map() try: - model_obj = _find_model_obj( + model_obj = find_model_obj( model_map, model_provider, model_name, diff --git a/backend/onyx/tools/utils.py b/backend/onyx/tools/utils.py index e5e56822d22c..2debc0c8c6c8 100644 --- a/backend/onyx/tools/utils.py +++ b/backend/onyx/tools/utils.py @@ -6,14 +6,15 @@ from onyx.configs.app_configs import AZURE_DALLE_API_KEY from onyx.db.connector import check_connectors_exist from onyx.db.document import check_docs_exist from onyx.db.models import LLMProvider -from onyx.llm.utils import get_model_map, _find_model_obj +from onyx.llm.utils import find_model_obj +from onyx.llm.utils import get_model_map from onyx.natural_language_processing.utils import BaseTokenizer from onyx.tools.tool import Tool def explicit_tool_calling_supported(model_provider: str, model_name: str) -> bool: model_map = get_model_map() - model_obj = _find_model_obj( + model_obj = find_model_obj( model_map=model_map, provider=model_provider, model_name=model_name,