Small cleanup

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"))
+        LITELLM_EXTRA_HEADERS = json.loads(_LITELLM_EXTRA_HEADERS_RAW)
-    except Exception as e:
+    except Exception:
-        import logging
+        # need to import here to avoid circular imports
-        logger = logging.getLogger(__name__)
+        from danswer.utils.logger import setup_logger
+
+        logger = setup_logger()
         logger.error(
             "Failed to parse LITELLM_EXTRA_HEADERS, must be a valid JSON object"
         )

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})

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

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)