diff --git a/backend/danswer/configs/model_configs.py b/backend/danswer/configs/model_configs.py index 19afc352e8cf..151b41811c52 100644 --- a/backend/danswer/configs/model_configs.py +++ b/backend/danswer/configs/model_configs.py @@ -101,12 +101,15 @@ DISABLE_LITELLM_STREAMING = ( # extra headers to pass to LiteLLM LITELLM_EXTRA_HEADERS = None -if os.environ.get("LITELLM_EXTRA_HEADERS"): +_LITELLM_EXTRA_HEADERS_RAW = os.environ.get("LITELLM_EXTRA_HEADERS") +if _LITELLM_EXTRA_HEADERS_RAW: try: - LITELLM_EXTRA_HEADERS = json.loads(os.environ.get("LITELLM_EXTRA_HEADERS")) - except Exception as e: - import logging - logger = logging.getLogger(__name__) + LITELLM_EXTRA_HEADERS = json.loads(_LITELLM_EXTRA_HEADERS_RAW) + except Exception: + # need to import here to avoid circular imports + from danswer.utils.logger import setup_logger + + logger = setup_logger() logger.error( "Failed to parse LITELLM_EXTRA_HEADERS, must be a valid JSON object" ) diff --git a/backend/danswer/llm/chat_llm.py b/backend/danswer/llm/chat_llm.py index 357d1ae93086..42d4214f761d 100644 --- a/backend/danswer/llm/chat_llm.py +++ b/backend/danswer/llm/chat_llm.py @@ -1,6 +1,7 @@ import abc import os from collections.abc import Iterator +from typing import Any import litellm # type:ignore from langchain.chat_models.base import BaseChatModel @@ -89,7 +90,7 @@ class DefaultMultiLLM(LangChainChatLLM): """Uses Litellm library to allow easy configuration to use a multitude of LLMs See https://python.langchain.com/docs/integrations/chat/litellm""" - DEFAULT_MODEL_PARAMS = { + DEFAULT_MODEL_PARAMS: dict[str, Any] = { "frequency_penalty": 0, "presence_penalty": 0, } @@ -124,7 +125,9 @@ class DefaultMultiLLM(LangChainChatLLM): for k, v in custom_config.items(): os.environ[k] = v - model_kwargs = DefaultMultiLLM.DEFAULT_MODEL_PARAMS if model_provider == "openai" else {} + model_kwargs = ( + DefaultMultiLLM.DEFAULT_MODEL_PARAMS if model_provider == "openai" else {} + ) if extra_headers: model_kwargs.update({"extra_headers": extra_headers}) diff --git a/backend/danswer/llm/factory.py b/backend/danswer/llm/factory.py index 02bc2425201a..3f4172c51615 100644 --- a/backend/danswer/llm/factory.py +++ b/backend/danswer/llm/factory.py @@ -1,6 +1,7 @@ from danswer.configs.app_configs import DISABLE_GENERATIVE_AI from danswer.configs.chat_configs import QA_TIMEOUT -from danswer.configs.model_configs import GEN_AI_TEMPERATURE, LITELLM_EXTRA_HEADERS +from danswer.configs.model_configs import GEN_AI_TEMPERATURE +from danswer.configs.model_configs import LITELLM_EXTRA_HEADERS from danswer.db.engine import get_session_context_manager from danswer.db.llm import fetch_default_provider from danswer.db.llm import fetch_provider diff --git a/backend/danswer/server/query_and_chat/token_budget.py b/backend/danswer/server/query_and_chat/token_budget.py index 1d1238c52770..49a84f5b0ec7 100644 --- a/backend/danswer/server/query_and_chat/token_budget.py +++ b/backend/danswer/server/query_and_chat/token_budget.py @@ -21,7 +21,13 @@ TIME_PERIOD_HOURS_DEFAULT = 12 def is_under_token_budget(db_session: Session) -> bool: - settings_json = cast(str, get_dynamic_config_store().load(TOKEN_BUDGET_SETTINGS)) + try: + settings_json = cast( + str, get_dynamic_config_store().load(TOKEN_BUDGET_SETTINGS) + ) + except Exception: + return True + settings = json.loads(settings_json) is_enabled = settings.get(ENABLE_TOKEN_BUDGET, False)