Skip to content

Commit

Permalink
fix: vertex ai credentials parsing
Browse files Browse the repository at this point in the history
- Add validation for vertex_credentials in custom_config when using Vertex AI
- Parse service account JSON from custom_config for litellm completion
- Add error handling for missing or invalid credentials
  • Loading branch information
sanjuruk committed Nov 27, 2024
1 parent 0553062 commit ba362a6
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions backend/danswer/llm/chat_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,20 @@ def _completion(
self._record_call(processed_prompt)

try:
# Get vertex_credentials directly from custom_config if model provider is Vertex AI
vertex_credentials = None
if self._model_provider == "vertex_ai":
if not self._custom_config:
raise ValueError(
"Vertex AI requires 'vertex_credentials' to be provided in custom_config"
)
vertex_credentials = self._custom_config.get("vertex_credentials")
if not vertex_credentials:
raise ValueError(
"Vertex AI credentials not found."
"Please provide 'vertex_credentials' with vertex ai service account details in json format."
)

return litellm.completion(
# model choice
model=f"{self.config.model_provider}/{self.config.deployment_name or self.config.model_name}",
Expand Down Expand Up @@ -390,6 +404,7 @@ def _completion(
if structured_response_format
else {}
),
**({"vertex_credentials": vertex_credentials} if vertex_credentials else {}),
**self._model_kwargs,
)
except Exception as e:
Expand Down

0 comments on commit ba362a6

Please sign in to comment.