From 8d57fb4845603a791d88c75dbbbc05f2675bfa54 Mon Sep 17 00:00:00 2001 From: Weves Date: Tue, 12 Nov 2024 14:20:21 -0800 Subject: [PATCH 1/3] Support LITELLM_EXTRA_BODY env variable --- backend/danswer/configs/model_configs.py | 9 +++++++++ backend/danswer/llm/chat_llm.py | 5 +++++ 2 files changed, 14 insertions(+) diff --git a/backend/danswer/configs/model_configs.py b/backend/danswer/configs/model_configs.py index c9668cd8136..3f0a9749572 100644 --- a/backend/danswer/configs/model_configs.py +++ b/backend/danswer/configs/model_configs.py @@ -119,3 +119,12 @@ logger.error( "Failed to parse LITELLM_PASS_THROUGH_HEADERS, must be a valid JSON object" ) + + +LITELLM_EXTRA_BODY: dict | None = None +_LITELLM_EXTRA_BODY_RAW = os.environ.get("LITELLM_EXTRA_BODY") +if _LITELLM_EXTRA_BODY_RAW: + try: + LITELLM_EXTRA_BODY = json.loads(_LITELLM_EXTRA_BODY_RAW) + except Exception: + pass diff --git a/backend/danswer/llm/chat_llm.py b/backend/danswer/llm/chat_llm.py index 6bd4b51d40b..0854c384692 100644 --- a/backend/danswer/llm/chat_llm.py +++ b/backend/danswer/llm/chat_llm.py @@ -26,6 +26,7 @@ from danswer.configs.app_configs import LOG_DANSWER_MODEL_INTERACTIONS from danswer.configs.model_configs import DISABLE_LITELLM_STREAMING from danswer.configs.model_configs import GEN_AI_TEMPERATURE +from danswer.configs.model_configs import LITELLM_EXTRA_BODY from danswer.llm.interfaces import LLM from danswer.llm.interfaces import LLMConfig from danswer.llm.interfaces import ToolChoiceOptions @@ -213,6 +214,7 @@ def __init__( temperature: float = GEN_AI_TEMPERATURE, custom_config: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, + extra_body: dict | None = LITELLM_EXTRA_BODY, ): self._timeout = timeout self._model_provider = model_provider @@ -246,6 +248,8 @@ def __init__( model_kwargs: dict[str, Any] = {} if extra_headers: model_kwargs.update({"extra_headers": extra_headers}) + if extra_body: + model_kwargs.update({"extra_body": extra_body}) self._model_kwargs = model_kwargs @@ -283,6 +287,7 @@ def _completion( tool_choice: ToolChoiceOptions | None, stream: bool, structured_response_format: dict | None = None, + extra_body: dict | None = None, ) -> litellm.ModelResponse | litellm.CustomStreamWrapper: if isinstance(prompt, list): prompt = [ From 96fb512c41412c0169efe8bc783522e24aa6415e Mon Sep 17 00:00:00 2001 From: Weves Date: Tue, 12 Nov 2024 14:21:41 -0800 Subject: [PATCH 2/3] Remove unused param --- backend/danswer/llm/chat_llm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backend/danswer/llm/chat_llm.py b/backend/danswer/llm/chat_llm.py index 0854c384692..8529bac37ba 100644 --- a/backend/danswer/llm/chat_llm.py +++ b/backend/danswer/llm/chat_llm.py @@ -287,7 +287,6 @@ def _completion( tool_choice: ToolChoiceOptions | None, stream: bool, structured_response_format: dict | None = None, - extra_body: dict | None = None, ) -> litellm.ModelResponse | litellm.CustomStreamWrapper: if isinstance(prompt, list): prompt = [ From 06c722bde2b2413bdf9927f794034c8c70fd84c6 Mon Sep 17 00:00:00 2001 From: Weves Date: Tue, 12 Nov 2024 14:23:02 -0800 Subject: [PATCH 3/3] Add comment --- backend/danswer/configs/model_configs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backend/danswer/configs/model_configs.py b/backend/danswer/configs/model_configs.py index 3f0a9749572..0618bf5f684 100644 --- a/backend/danswer/configs/model_configs.py +++ b/backend/danswer/configs/model_configs.py @@ -121,6 +121,8 @@ ) +# if specified, will merge the specified JSON with the existing body of the +# request before sending it to the LLM LITELLM_EXTRA_BODY: dict | None = None _LITELLM_EXTRA_BODY_RAW = os.environ.get("LITELLM_EXTRA_BODY") if _LITELLM_EXTRA_BODY_RAW: