Support LITELLM_EXTRA_BODY env variable (#3119)

* Support LITELLM_EXTRA_BODY env variable

* Remove unused param

* Add comment
This commit is contained in:
Chris Weaver 2024-11-12 15:17:44 -08:00 committed by GitHub
parent 942e47db29
commit 021d0cf314
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 0 deletions

View File

@ -119,3 +119,14 @@ if _LITELLM_PASS_THROUGH_HEADERS_RAW:
logger.error(
"Failed to parse LITELLM_PASS_THROUGH_HEADERS, must be a valid JSON object"
)
# 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:
try:
LITELLM_EXTRA_BODY = json.loads(_LITELLM_EXTRA_BODY_RAW)
except Exception:
pass

View File

@ -26,6 +26,7 @@ from danswer.configs.app_configs import LOG_ALL_MODEL_INTERACTIONS
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 @@ class DefaultMultiLLM(LLM):
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 @@ class DefaultMultiLLM(LLM):
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