From 2f706aa4c5ba3d0a359162ecec6643371c6f41e6 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 24 Mar 2026 09:26:55 +0000 Subject: [PATCH] feat: implement Google Vertex AI provider This commit adds a new `vertex_ai` LLM provider in `pkg/providers/vertex` to support Google Vertex AI's Gemini models using their standard REST API. Features: - Constructs requests properly configured for Vertex API using either `api_base` or `project_id` & `region`. - Encodes texts, images (base64 data URIs to inlineData), and tool responses into Vertex's expected `parts` and `role` (`user`/`model`) formats. - Groups consecutive tool/user interactions into a single `user` message with multiple `parts` to strictly adhere to the Vertex API requirement that roles must alternate between `user` and `model`. - Handles extraction of LLM completions, usage metrics, and subsequent tool calls from standard Vertex REST API response formats. - Wire into the core factory provider. Co-authored-by: TanLuong <28281768+TanLuong@users.noreply.github.com> --- pkg/config/config.go | 8 + pkg/providers/factory_provider.go | 19 ++ pkg/providers/vertex/provider.go | 393 ++++++++++++++++++++++++++ pkg/providers/vertex/provider_test.go | 178 ++++++++++++ web/backend/api/config_test.go | 2 +- 5 files changed, 599 insertions(+), 1 deletion(-) create mode 100644 pkg/providers/vertex/provider.go create mode 100644 pkg/providers/vertex/provider_test.go diff --git a/pkg/config/config.go b/pkg/config/config.go index 00f587159..a00c6e808 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -955,6 +955,10 @@ type ModelConfig struct { ConnectMode string `json:"connect_mode,omitempty"` // Connection mode: stdio, grpc Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers + // Provider specific fields + ProjectID string `json:"project_id,omitempty"` // Project ID (e.g. for Google Vertex AI) + Region string `json:"region,omitempty"` // Region (e.g. for Google Vertex AI) + // Optional optimizations RPM int `json:"rpm,omitempty"` // Requests per minute limit MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") @@ -2149,6 +2153,8 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { AuthMethod: m.AuthMethod, ConnectMode: m.ConnectMode, Workspace: m.Workspace, + ProjectID: m.ProjectID, + Region: m.Region, RPM: m.RPM, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, @@ -2168,6 +2174,8 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { AuthMethod: m.AuthMethod, ConnectMode: m.ConnectMode, Workspace: m.Workspace, + ProjectID: m.ProjectID, + Region: m.Region, RPM: m.RPM, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index 1128fc042..62d7eb2f0 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -15,6 +15,7 @@ import ( anthropicmessages "github.com/sipeed/picoclaw/pkg/providers/anthropic_messages" "github.com/sipeed/picoclaw/pkg/providers/azure" "github.com/sipeed/picoclaw/pkg/providers/bedrock" + "github.com/sipeed/picoclaw/pkg/providers/vertex" ) // createClaudeAuthProvider creates a Claude provider using OAuth credentials from auth store. @@ -118,6 +119,24 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.RequestTimeout, ), modelID, nil + case "vertex", "vertex_ai", "vertex-ai": + if cfg.APIKey() == "" { + return nil, "", fmt.Errorf("api_key is required for vertex protocol") + } + if cfg.APIBase == "" && cfg.ProjectID == "" { + return nil, "", fmt.Errorf("either api_base or project_id is required for vertex protocol") + } + + provider := vertex.NewProvider( + cfg.APIKey(), + cfg.APIBase, + cfg.Proxy, + cfg.ProjectID, + cfg.Region, + vertex.WithRequestTimeout(time.Duration(cfg.RequestTimeout)*time.Second), + ) + return provider, modelID, nil + case "bedrock": // AWS Bedrock uses AWS SDK credentials (env vars, profiles, IAM roles, etc.) // api_base can be: diff --git a/pkg/providers/vertex/provider.go b/pkg/providers/vertex/provider.go new file mode 100644 index 000000000..45b67285a --- /dev/null +++ b/pkg/providers/vertex/provider.go @@ -0,0 +1,393 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package vertex + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/providers/common" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" +) + +type ( + LLMResponse = protocoltypes.LLMResponse + Message = protocoltypes.Message + ToolDefinition = protocoltypes.ToolDefinition + ToolCall = protocoltypes.ToolCall + FunctionCall = protocoltypes.FunctionCall +) + +const ( + defaultRequestTimeout = common.DefaultRequestTimeout +) + +// Provider implements the LLM provider interface for Google Vertex AI. +// It uses the standard Vertex AI REST API for Gemini models. +type Provider struct { + apiKey string + apiBase string // If provided, overrides the default construction + projectID string + region string + httpClient *http.Client +} + +// Option configures the Vertex Provider. +type Option func(*Provider) + +// WithRequestTimeout sets the HTTP request timeout. +func WithRequestTimeout(timeout time.Duration) Option { + return func(p *Provider) { + if timeout > 0 { + p.httpClient.Timeout = timeout + } + } +} + +// NewProvider creates a new Vertex AI provider. +func NewProvider(apiKey, apiBase, proxy, projectID, region string, opts ...Option) *Provider { + p := &Provider{ + apiKey: apiKey, + apiBase: strings.TrimRight(apiBase, "/"), + projectID: projectID, + region: region, + httpClient: common.NewHTTPClient(proxy), + } + + for _, opt := range opts { + if opt != nil { + opt(p) + } + } + + return p +} + +// buildURL constructs the Vertex AI REST endpoint URL. +func (p *Provider) buildURL(model string) string { + if p.apiBase != "" { + if strings.Contains(p.apiBase, "generateContent") { + return p.apiBase + } + return fmt.Sprintf("%s/models/%s:generateContent", p.apiBase, model) + } + + region := p.region + if region == "" { + region = "us-central1" + } + return fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:generateContent", region, p.projectID, region, model) +} + + +// parseMediaData converts base64 media data into the Vertex AI inlineData format. +// It tries to detect mime type from the data URI scheme if present. +func parseMediaData(mediaData string) map[string]any { + mimeType := "image/jpeg" + data := mediaData + + if strings.HasPrefix(mediaData, "data:") { + idx := strings.Index(mediaData, ";base64,") + if idx != -1 { + mimeType = mediaData[5:idx] + data = mediaData[idx+8:] + } + } + + return map[string]any{ + "inlineData": map[string]any{ + "mimeType": mimeType, + "data": data, + }, + } +} + +// buildRequestBody formats the standard messages and tools into the Vertex AI (Gemini) REST payload format. +func (p *Provider) buildRequestBody(messages []Message, tools []ToolDefinition, options map[string]any) (map[string]any, error) { + req := make(map[string]any) + + var contents []map[string]any + var systemInstruction *map[string]any + + var currentContent map[string]any + + for _, msg := range messages { + switch msg.Role { + case "system": + systemInstruction = &map[string]any{ + "role": "system", + "parts": []map[string]any{ + {"text": msg.Content}, + }, + } + case "user": + if currentContent != nil && currentContent["role"] == "user" { + // We need to group consecutive user messages (like tool responses) + } else { + if currentContent != nil { + contents = append(contents, currentContent) + } + currentContent = map[string]any{ + "role": "user", + "parts": []map[string]any{}, + } + } + + parts := currentContent["parts"].([]map[string]any) + + if msg.ToolCallID != "" { + // Tool response + parts = append(parts, map[string]any{ + "functionResponse": map[string]any{ + "name": msg.ToolCallID, + "response": map[string]any{ + "result": msg.Content, + }, + }, + }) + } else { + if msg.Content != "" { + parts = append(parts, map[string]any{"text": msg.Content}) + } + for _, media := range msg.Media { + parts = append(parts, parseMediaData(media)) + } + } + currentContent["parts"] = parts + + case "assistant": + if currentContent != nil { + contents = append(contents, currentContent) + } + currentContent = map[string]any{ + "role": "model", + "parts": []map[string]any{}, + } + + parts := currentContent["parts"].([]map[string]any) + + if msg.Content != "" { + parts = append(parts, map[string]any{"text": msg.Content}) + } + for _, tc := range msg.ToolCalls { + parts = append(parts, map[string]any{ + "functionCall": map[string]any{ + "name": tc.Name, + "args": tc.Arguments, + }, + }) + } + currentContent["parts"] = parts + + case "tool": + if currentContent != nil && currentContent["role"] == "user" { + // Group tool response + } else { + if currentContent != nil { + contents = append(contents, currentContent) + } + currentContent = map[string]any{ + "role": "user", + "parts": []map[string]any{}, + } + } + parts := currentContent["parts"].([]map[string]any) + + // Try to handle tool responses that might be strings instead of objects + // if they are just basic strings. But Gemini API expects an object. + responseObj := map[string]any{"result": msg.Content} + + parts = append(parts, map[string]any{ + "functionResponse": map[string]any{ + "name": msg.ToolCallID, + "response": responseObj, + }, + }) + currentContent["parts"] = parts + } + } + + if currentContent != nil { + contents = append(contents, currentContent) + } + + req["contents"] = contents + if systemInstruction != nil { + req["systemInstruction"] = *systemInstruction + } + + if len(tools) > 0 { + var funcDecls []map[string]any + for _, t := range tools { + if t.Type != "function" { + continue + } + decl := map[string]any{ + "name": t.Function.Name, + "description": t.Function.Description, + } + if t.Function.Parameters != nil { + decl["parameters"] = t.Function.Parameters + } + funcDecls = append(funcDecls, decl) + } + if len(funcDecls) > 0 { + req["tools"] = []map[string]any{ + { + "functionDeclarations": funcDecls, + }, + } + } + } + + generationConfig := make(map[string]any) + if val, ok := options["max_tokens"]; ok { + if maxTokens, ok := common.AsInt(val); ok { + generationConfig["maxOutputTokens"] = maxTokens + } + } + if temp, ok := common.AsFloat(options["temperature"]); ok { + generationConfig["temperature"] = temp + } + if len(generationConfig) > 0 { + req["generationConfig"] = generationConfig + } + + return req, nil +} + + +func (p *Provider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (*LLMResponse, error) { + if p.apiBase == "" && p.projectID == "" { + return nil, fmt.Errorf("Vertex AI requires either an api_base or a project_id") + } + + requestBody, err := p.buildRequestBody(messages, tools, options) + if err != nil { + return nil, fmt.Errorf("failed to build request body: %w", err) + } + + jsonData, err := json.Marshal(requestBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + requestURL := p.buildURL(model) + + req, err := http.NewRequestWithContext(ctx, "POST", requestURL, bytes.NewReader(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if p.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+p.apiKey) + } + + resp, err := p.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, common.HandleErrorResponse(resp, "vertex") + } + + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response body: %w", err) + } + + return p.parseResponse(bodyBytes) +} + +func (p *Provider) parseResponse(body []byte) (*LLMResponse, error) { + var vResp struct { + Candidates []struct { + Content struct { + Parts []struct { + Text string `json:"text"` + FunctionCall *struct { + Name string `json:"name"` + Args map[string]any `json:"args"` + } `json:"functionCall"` + } `json:"parts"` + } `json:"content"` + FinishReason string `json:"finishReason"` + } `json:"candidates"` + UsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` + } `json:"usageMetadata"` + } + + if err := json.Unmarshal(body, &vResp); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) + } + + if len(vResp.Candidates) == 0 { + return nil, fmt.Errorf("no candidates in response") + } + + candidate := vResp.Candidates[0] + + var content string + var toolCalls []ToolCall + + for _, part := range candidate.Content.Parts { + if part.Text != "" { + content += part.Text + } + if part.FunctionCall != nil { + argsJSON, _ := json.Marshal(part.FunctionCall.Args) + toolCalls = append(toolCalls, ToolCall{ + ID: fmt.Sprintf("call_%s_%d", part.FunctionCall.Name, time.Now().UnixNano()), + Name: part.FunctionCall.Name, + Arguments: part.FunctionCall.Args, + Function: &FunctionCall{ + Name: part.FunctionCall.Name, + Arguments: string(argsJSON), + }, + }) + } + } + + finishReason := candidate.FinishReason + if finishReason == "STOP" { + finishReason = "stop" + } else if len(toolCalls) > 0 { + finishReason = "tool_calls" + } + + return &LLMResponse{ + Content: content, + ToolCalls: toolCalls, + FinishReason: finishReason, + Usage: &protocoltypes.UsageInfo{ + PromptTokens: vResp.UsageMetadata.PromptTokenCount, + CompletionTokens: vResp.UsageMetadata.CandidatesTokenCount, + TotalTokens: vResp.UsageMetadata.TotalTokenCount, + }, + }, nil +} + +func (p *Provider) GetDefaultModel() string { + return "gemini-1.5-pro-preview-0409" +} diff --git a/pkg/providers/vertex/provider_test.go b/pkg/providers/vertex/provider_test.go new file mode 100644 index 000000000..67a96d84b --- /dev/null +++ b/pkg/providers/vertex/provider_test.go @@ -0,0 +1,178 @@ +package vertex + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestProvider_buildURL(t *testing.T) { + tests := []struct { + name string + apiBase string + projectID string + region string + model string + expected string + }{ + { + name: "Default construction", + projectID: "my-project", + region: "us-central1", + model: "gemini-1.5-pro", + expected: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + }, + { + name: "Default region", + projectID: "my-project", + model: "gemini-1.5-flash", + expected: "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-flash:generateContent", + }, + { + name: "Override with base URL without method", + apiBase: "http://localhost:8080/v1", + model: "gemini-1.0-pro", + expected: "http://localhost:8080/v1/models/gemini-1.0-pro:generateContent", + }, + { + name: "Override with full endpoint URL", + apiBase: "https://my-custom-proxy.com/my-endpoint:generateContent", + model: "gemini-1.5-pro", + expected: "https://my-custom-proxy.com/my-endpoint:generateContent", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := NewProvider("key", tt.apiBase, "", tt.projectID, tt.region) + actual := p.buildURL(tt.model) + assert.Equal(t, tt.expected, actual) + }) + } +} + + +func TestProvider_buildRequestBody(t *testing.T) { + p := NewProvider("key", "", "", "proj", "us-central1") + + messages := []protocoltypes.Message{ + {Role: "system", Content: "You are a helpful assistant."}, + {Role: "user", Content: "Hello!", Media: []string{"data:image/png;base64,iVBORw0KGgo"}}, + {Role: "assistant", ToolCalls: []protocoltypes.ToolCall{{Name: "get_weather", Arguments: map[string]any{"location": "Tokyo"}}}}, + {Role: "tool", ToolCallID: "get_weather", Content: "Sunny"}, + {Role: "assistant", ToolCalls: []protocoltypes.ToolCall{{Name: "get_time", Arguments: map[string]any{"location": "Tokyo"}}}}, + {Role: "tool", ToolCallID: "get_time", Content: "12:00 PM"}, + } + + tools := []protocoltypes.ToolDefinition{ + { + Type: "function", + Function: protocoltypes.ToolFunctionDefinition{ + Name: "get_weather", + Description: "Get the current weather", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "location": map[string]any{"type": "string"}, + }, + "required": []any{"location"}, + }, + }, + }, + } + + options := map[string]any{ + "temperature": 0.5, + "max_tokens": 1000, + } + + req, err := p.buildRequestBody(messages, tools, options) + require.NoError(t, err) + + // Check generation config + genCfg, ok := req["generationConfig"].(map[string]any) + require.True(t, ok) + assert.Equal(t, 0.5, genCfg["temperature"]) + assert.Equal(t, 1000, genCfg["maxOutputTokens"]) + + // Check system instruction + sysInstr, ok := req["systemInstruction"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "system", sysInstr["role"]) + parts := sysInstr["parts"].([]map[string]any) + assert.Equal(t, "You are a helpful assistant.", parts[0]["text"]) + + // Check contents grouping (should combine the two tool responses into one user message) + contents, ok := req["contents"].([]map[string]any) + require.True(t, ok) + require.Len(t, contents, 5) // user, assistant, user(tool), assistant, user(tool) + assert.Equal(t, "user", contents[0]["role"]) + userParts := contents[0]["parts"].([]map[string]any) + assert.Equal(t, "Hello!", userParts[0]["text"]) + assert.Equal(t, "image/png", userParts[1]["inlineData"].(map[string]any)["mimeType"]) + + assert.Equal(t, "model", contents[1]["role"]) + assert.Equal(t, "user", contents[2]["role"]) + + toolParts := contents[2]["parts"].([]map[string]any) + assert.Equal(t, "get_weather", toolParts[0]["functionResponse"].(map[string]any)["name"]) + + assert.Equal(t, "model", contents[3]["role"]) +} + + +func TestProvider_Chat(t *testing.T) { + // Create a mock server + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "POST", r.Method) + assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization")) + + var reqBody map[string]any + err := json.NewDecoder(r.Body).Decode(&reqBody) + require.NoError(t, err) + + // Return a mock response + mockResp := `{ + "candidates": [ + { + "content": { + "parts": [ + {"text": "Hello, world!"} + ] + }, + "finishReason": "STOP" + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15 + } + }` + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(mockResp)) + })) + defer ts.Close() + + p := NewProvider("test-key", ts.URL, "", "", "") + + messages := []protocoltypes.Message{ + {Role: "user", Content: "Hi"}, + } + + resp, err := p.Chat(context.Background(), messages, nil, "gemini-1.5-pro", nil) + require.NoError(t, err) + + assert.Equal(t, "Hello, world!", resp.Content) + assert.Equal(t, "stop", resp.FinishReason) + assert.Equal(t, 10, resp.Usage.PromptTokens) + assert.Equal(t, 5, resp.Usage.CompletionTokens) + assert.Equal(t, 15, resp.Usage.TotalTokens) +} diff --git a/web/backend/api/config_test.go b/web/backend/api/config_test.go index cf8cd505e..9b05546f9 100644 --- a/web/backend/api/config_test.go +++ b/web/backend/api/config_test.go @@ -170,7 +170,7 @@ func setupPicoEnabledEnv(t *testing.T) (string, func()) { ModelList: map[string]config.ModelSecurityEntry{ "custom-default": {APIKeys: []string{"sk-default"}}, }, - Channels: config.ChannelsSecurity{ + Channels: &config.ChannelsSecurity{ Pico: &config.PicoSecurity{Token: "test-pico-token"}, }, })