feat(vertex): support Vertex AI for generation

Author: badmonster0

Cargo.lock

@@ -91,12 +91,7 @@ tree-sitter-typescript = "0.23.2"
 tree-sitter-xml = "0.7.0"
 tree-sitter-yaml = "0.7.0"
 
-globset = "0.4.16"
-unicase = "2.8.1"
-google-drive3 = "6.0.0"
-hyper-util = "0.1.11"
-hyper-rustls = { version = "0.27.5" }
-yup-oauth2 = "12.1.0"
+yup-oauth2 = { version = "12.1.0" }
 rustls = { version = "0.23.25" }
 http-body-util = "0.1.3"
 yaml-rust2 = "0.10.1"
@@ -117,3 +112,9 @@ aws-sdk-sqs = "1.67.0"
 numpy = "0.25.0"
 infer = "0.19.0"
 serde_with = { version = "3.13.0", features = ["base64"] }
+globset = "0.4.16"
+unicase = "2.8.1"
+google-drive3 = "6.0.0"
+hyper-util = "0.1.11"
+hyper-rustls = { version = "0.27.5" }
+google-cloud-aiplatform-v1 = "0.4.0"

Cargo.toml

@@ -33,6 +33,7 @@
     # Submodules
     "_engine",
     "functions",
+    "llm",
     "sources",
     "targets",
     "storages",

python/cocoindex/__init__.py

@@ -45,6 +45,7 @@ class EmbedText(op.FunctionSpec):
     address: str | None = None
     output_dimension: int | None = None
     task_type: str | None = None
+    api_config: llm.VertexAiConfig | None = None
 
 
 class ExtractByLlm(op.FunctionSpec):

python/cocoindex/functions.py

@@ -8,17 +8,29 @@ class LlmApiType(Enum):
     OPENAI = "OpenAi"
     OLLAMA = "Ollama"
     GEMINI = "Gemini"
+    VERTEX_AI = "VertexAi"
     ANTHROPIC = "Anthropic"
     LITE_LLM = "LiteLlm"
     OPEN_ROUTER = "OpenRouter"
     VOYAGE = "Voyage"
     VLLM = "Vllm"
 
 
+@dataclass
+class VertexAiConfig:
+    """A specification for a Vertex AI LLM."""
+
+    kind = "VertexAi"
+
+    project: str
+    location: str | None = None
+
+
 @dataclass
 class LlmSpec:
     """A specification for a LLM."""
 
     api_type: LlmApiType
     model: str
     address: str | None = None
+    api_config: VertexAiConfig | None = None