highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-05-28 12:39:54 +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
# env variables)
if custom_config:
# Specifically pass in "vertex_credentials" as a model_kwarg to the
# completion call for vertex AI. More details here:
# Specifically pass in "vertex_credentials" / "vertex_location" as a
# model_kwarg to the completion call for vertex AI. More details here:
# https://docs.litellm.ai/docs/providers/vertex
vertex_credentials_key = "vertex_credentials"
vertex_credentials = custom_config.get(vertex_credentials_key)
if vertex_credentials and model_provider == "vertex_ai":
model_kwargs[vertex_credentials_key] = vertex_credentials
else:
# standard case
for k, v in custom_config.items():
os.environ[k] = v
vertex_location_key = "vertex_location"
for k, v in custom_config.items():
if model_provider == "vertex_ai":
if k == vertex_credentials_key:
model_kwargs[k] = v
continue
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:
model_kwargs.update({"extra_headers": extra_headers})