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

Enable location support for Vertex AI (#3707)

Browse Source
This commit is contained in:
Chris Weaver 2025-01-19 09:41:35 -08:00 committed by GitHub
parent 342bb9f685
commit bf30dab9c4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -275,17 +275,22 @@ class DefaultMultiLLM(LLM):
# addtional kwargs (and some kwargs MUST be passed in rather than set as # addtional kwargs (and some kwargs MUST be passed in rather than set as
# env variables) # env variables)
if custom_config: if custom_config:
# Specifically pass in "vertex_credentials" as a model_kwarg to the # Specifically pass in "vertex_credentials" / "vertex_location" as a
# 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_credentials_key = "vertex_credentials"
vertex_credentials = custom_config.get(vertex_credentials_key) vertex_location_key = "vertex_location"
if vertex_credentials and model_provider == "vertex_ai": for k, v in custom_config.items():
model_kwargs[vertex_credentials_key] = vertex_credentials if model_provider == "vertex_ai":
else: if k == vertex_credentials_key:
# standard case model_kwargs[k] = v
for k, v in custom_config.items(): continue
os.environ[k] = v elif k == vertex_location_key:
model_kwargs[k] = v
continue
# for all values, set them as env variables
os.environ[k] = v
if extra_headers: if extra_headers:
model_kwargs.update({"extra_headers": extra_headers}) model_kwargs.update({"extra_headers": extra_headers})