feat: support google vertex llm provider (#599)

Close #523. Considering that our current `llama_index` can already support `Vertex`, should we merge `ANTHROPIC_VERTEX` into the same type as `VERTEX` in this commit? --------- Co-authored-by: Mini256 <[email protected]>
pingcap · Jan 25, 2025 · 3abfe48 · 3abfe48
1 parent ec67725
commit 3abfe48
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 14 deletions.
diff --git a/backend/app/rag/chat.py b/backend/app/rag/chat.py
@@ -151,9 +151,9 @@ def __init__(
                     user_id=self.user.id if self.user else None,
                     browser_id=self.browser_id,
                     origin=origin,
-                    visibility=ChatVisibility.PUBLIC.value
+                    visibility=ChatVisibility.PUBLIC
                     if not self.user
-                    else ChatVisibility.PRIVATE.value,
+                    else ChatVisibility.PRIVATE,
                 ),
             )
             chat_id = self.db_chat_obj.id

diff --git a/backend/app/rag/llms/provider.py b/backend/app/rag/llms/provider.py
@@ -7,7 +7,9 @@
 class LLMProvider(str, enum.Enum):
     OPENAI = "openai"
     GEMINI = "gemini"
+    # TODO: remove ANTHROPIC_VERTEX as it's deprecated.
     ANTHROPIC_VERTEX = "anthropic_vertex"
+    VERTEX = "vertex"
     OPENAI_LIKE = "openai_like"
     BEDROCK = "bedrock"
     OLLAMA = "ollama"
@@ -117,12 +119,12 @@ class LLMProviderOption(BaseModel):
         default_credentials="****",
     ),
     LLMProviderOption(
-        provider=LLMProvider.ANTHROPIC_VERTEX,
-        provider_display_name="Anthropic Vertex AI",
-        provider_description="Anthropic's Claude models are now generally available through Vertex AI.",
-        provider_url="https://docs.anthropic.com/en/api/claude-on-vertex-ai",
-        default_llm_model="claude-3-5-sonnet@20241022",
-        llm_model_description="",
+        provider=LLMProvider.VERTEX,
+        provider_display_name="Vertex AI",
+        provider_description="Vertex AI is a fully-managed, unified AI development platform for building and using generative AI.",
+        provider_url="https://cloud.google.com/vertex-ai",
+        default_llm_model="gemini-1.5-flash",
+        llm_model_description="Find more in https://cloud.google.com/model-garden",
         credentials_display_name="Google Credentials JSON",
         credentials_description="The JSON Object of Google Credentials, refer to https://cloud.google.com/docs/authentication/provide-credentials-adc#on-prem",
         credentials_type="dict",

diff --git a/backend/app/rag/llms/resolver.py b/backend/app/rag/llms/resolver.py
@@ -61,18 +61,19 @@ def resolve_llm(
             llm.aws_secret_access_key = secret_access_key
             llm.region_name = region_name
             return llm
-        case LLMProvider.ANTHROPIC_VERTEX:
+        case LLMProvider.VERTEX | LLMProvider.ANTHROPIC_VERTEX:
             google_creds: service_account.Credentials = (
                 service_account.Credentials.from_service_account_info(
                     credentials,
                     scopes=["https://www.googleapis.com/auth/cloud-platform"],
                 )
             )
             google_creds.refresh(request=Request())
-            if "max_tokens" not in config:
-                config.update(max_tokens=4096)
+            config.setdefault("max_tokens", 4096)
+            config.setdefault("context_window", 200 * 1000)
             return Vertex(
                 model=model,
+                project=credentials["project_id"],
                 credentials=google_creds,
                 **config,
             )

diff --git a/backend/app/utils/dspy.py b/backend/app/utils/dspy.py
@@ -74,7 +74,10 @@ def get_dspy_lm_by_llama_llm(llama_llm: BaseLLM) -> dspy.LM:
                 "Bedrock model " + llama_llm.model + " is not supported by dspy."
             )
     elif type(llama_llm) is Vertex:
-        raise ValueError("Vertex is not supported by dspy.")
+        return dspy.GoogleVertexAI(
+            model=llama_llm.model,
+            max_output_tokens=llama_llm.max_tokens or 8192,
+        )
     elif type(llama_llm) is Ollama:
         return DspyOllamaLocal(
             model=llama_llm.model,

diff --git a/frontend/app/src/pages/docs/llm.mdx b/frontend/app/src/pages/docs/llm.mdx
@@ -34,9 +34,9 @@ To learn more about OpenAI, please visit [OpenAI](https://platform.openai.com/).
 
 To learn more about Google Gemini, please visit [Google Gemini](https://gemini.google.com/).
 
-### Anthropic Vertex AI
+### Vertex AI
 
-To learn more about Anthropic Vertex AI, please visit [Anthropic Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude)
+To learn more about Vertex AI, please visit [Vertex AI](https://cloud.google.com/vertex-ai).
 
 ### Amazon Bedrock