highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-06-19 20:40:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: Fix duplicate kwarg issue when calling litellm.main.completion (#4533)

Browse Source 
* Fix duplicate kwarg issue

* Change how vertex_credentials are passed

* Modify temporary dict instead

* Change string to a global constant

* Add extra condition to if-check during population of map
This commit is contained in:
Raunak Bhagat 2025-04-16 19:29:53 -07:00 committed by GitHub
parent 2a9fd9342e
commit fe94bdf936
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
backend/onyx/llm/chat_llm.py
View File

@ -51,6 +51,7 @@ litellm.drop_params = True
litellm.telemetry = False litellm.telemetry = False
_LLM_PROMPT_LONG_TERM_LOG_CATEGORY = "llm_prompt" _LLM_PROMPT_LONG_TERM_LOG_CATEGORY = "llm_prompt"
VERTEX_CREDENTIALS_KWARG = "vertex_credentials"
class LLMTimeoutError(Exception): class LLMTimeoutError(Exception):
@ -303,11 +304,10 @@ class DefaultMultiLLM(LLM):
# Specifically pass in "vertex_credentials" / "vertex_location" as a # Specifically pass in "vertex_credentials" / "vertex_location" as a
# model_kwarg to the completion call for vertex AI. More details here: # model_kwarg to the completion call for vertex AI. More details here:
# https://docs.litellm.ai/docs/providers/vertex # https://docs.litellm.ai/docs/providers/vertex
vertex_credentials_key = "vertex_credentials"
vertex_location_key = "vertex_location" vertex_location_key = "vertex_location"
for k, v in custom_config.items(): for k, v in custom_config.items():
if model_provider == "vertex_ai": if model_provider == "vertex_ai":
if k == vertex_credentials_key: if k == VERTEX_CREDENTIALS_KWARG:
model_kwargs[k] = v model_kwargs[k] = v
continue continue
elif k == vertex_location_key: elif k == vertex_location_key:
@ -412,6 +412,13 @@ class DefaultMultiLLM(LLM):
processed_prompt = _prompt_to_dict(prompt) processed_prompt = _prompt_to_dict(prompt)
self._record_call(processed_prompt) self._record_call(processed_prompt)
final_model_kwargs = {**self._model_kwargs}
if (
VERTEX_CREDENTIALS_KWARG not in final_model_kwargs
and self.config.credentials_file
):
final_model_kwargs[VERTEX_CREDENTIALS_KWARG] = self.config.credentials_file
try: try:
return litellm.completion( return litellm.completion(
mock_response=MOCK_LLM_RESPONSE, mock_response=MOCK_LLM_RESPONSE,
@ -457,14 +464,7 @@ class DefaultMultiLLM(LLM):
if structured_response_format if structured_response_format
else {} else {}
), ),
**( **final_model_kwargs,
{
"vertex_credentials": self.config.credentials_file,
}
if self.config.model_provider == "vertex_ai"
else {}
),
**self._model_kwargs,
) )
except Exception as e: except Exception as e:
self._record_error(processed_prompt, e) self._record_error(processed_prompt, e)