Skip to content

Commit 0ebb654

Browse files
authored
feat: Add support for VertexAI compatible service (#375)
feat: consolidate Vertex AI compatibility with API key support in Gemini
1 parent 08a1d2e commit 0ebb654

File tree

9 files changed

+633
-77
lines changed

9 files changed

+633
-77
lines changed

.gitignore

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pgstore/*
1515
gitstore/*
1616
objectstore/*
1717
static/*
18+
refs/*
1819

1920
# Authentication data
2021
auths/*
@@ -30,3 +31,7 @@ GEMINI.md
3031
.vscode/*
3132
.claude/*
3233
.serena/*
34+
35+
# macOS
36+
.DS_Store
37+
._*

internal/api/server.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -934,6 +934,7 @@ func (s *Server) UpdateClients(cfg *config.Config) {
934934
geminiAPIKeyCount := len(cfg.GeminiKey)
935935
claudeAPIKeyCount := len(cfg.ClaudeKey)
936936
codexAPIKeyCount := len(cfg.CodexKey)
937+
vertexAICompatCount := len(cfg.VertexCompatAPIKey)
937938
openAICompatCount := 0
938939
for i := range cfg.OpenAICompatibility {
939940
entry := cfg.OpenAICompatibility[i]
@@ -944,13 +945,14 @@ func (s *Server) UpdateClients(cfg *config.Config) {
944945
openAICompatCount += len(entry.APIKeys)
945946
}
946947

947-
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
948-
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)\n",
948+
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + vertexAICompatCount + openAICompatCount
949+
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d Vertex-compat + %d OpenAI-compat)\n",
949950
total,
950951
authFiles,
951952
geminiAPIKeyCount,
952953
claudeAPIKeyCount,
953954
codexAPIKeyCount,
955+
vertexAICompatCount,
954956
openAICompatCount,
955957
)
956958
}

internal/config/config.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ type Config struct {
7070
// GeminiKey defines Gemini API key configurations with optional routing overrides.
7171
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
7272

73+
// VertexCompatAPIKey defines Vertex AI-compatible API key configurations for third-party providers.
74+
// Used for services that use Vertex AI-style paths but with simple API key authentication.
75+
VertexCompatAPIKey []VertexCompatKey `yaml:"vertex-api-key" json:"vertex-api-key"`
76+
7377
// RequestRetry defines the retry times when the request failed.
7478
RequestRetry int `yaml:"request-retry" json:"request-retry"`
7579
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
@@ -343,6 +347,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
343347
// Sanitize Gemini API key configuration and migrate legacy entries.
344348
cfg.SanitizeGeminiKeys()
345349

350+
// Sanitize Vertex-compatible API keys: drop entries without base-url
351+
cfg.SanitizeVertexCompatKeys()
352+
346353
// Sanitize Codex keys: drop entries without base-url
347354
cfg.SanitizeCodexKeys()
348355

@@ -831,6 +838,7 @@ func shouldSkipEmptyCollectionOnPersist(key string, node *yaml.Node) bool {
831838
switch key {
832839
case "generative-language-api-key",
833840
"gemini-api-key",
841+
"vertex-api-key",
834842
"claude-api-key",
835843
"codex-api-key",
836844
"openai-compatibility":

internal/config/vertex_compat.go

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package config
2+
3+
import "strings"
4+
5+
// VertexCompatKey represents the configuration for Vertex AI-compatible API keys.
6+
// This supports third-party services that use Vertex AI-style endpoint paths
7+
// (/publishers/google/models/{model}:streamGenerateContent) but authenticate
8+
// with simple API keys instead of Google Cloud service account credentials.
9+
//
10+
// Example services: zenmux.ai and similar Vertex-compatible providers.
11+
type VertexCompatKey struct {
12+
// APIKey is the authentication key for accessing the Vertex-compatible API.
13+
// Maps to the x-goog-api-key header.
14+
APIKey string `yaml:"api-key" json:"api-key"`
15+
16+
// BaseURL is the base URL for the Vertex-compatible API endpoint.
17+
// The executor will append "/v1/publishers/google/models/{model}:action" to this.
18+
// Example: "https://zenmux.ai/api" becomes "https://zenmux.ai/api/v1/publishers/google/models/..."
19+
BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"`
20+
21+
// ProxyURL optionally overrides the global proxy for this API key.
22+
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
23+
24+
// Headers optionally adds extra HTTP headers for requests sent with this key.
25+
// Commonly used for cookies, user-agent, and other authentication headers.
26+
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
27+
28+
// Models defines the model configurations including aliases for routing.
29+
Models []VertexCompatModel `yaml:"models,omitempty" json:"models,omitempty"`
30+
}
31+
32+
// VertexCompatModel represents a model configuration for Vertex compatibility,
33+
// including the actual model name and its alias for API routing.
34+
type VertexCompatModel struct {
35+
// Name is the actual model name used by the external provider.
36+
Name string `yaml:"name" json:"name"`
37+
38+
// Alias is the model name alias that clients will use to reference this model.
39+
Alias string `yaml:"alias" json:"alias"`
40+
}
41+
42+
// SanitizeVertexCompatKeys deduplicates and normalizes Vertex-compatible API key credentials.
43+
func (cfg *Config) SanitizeVertexCompatKeys() {
44+
if cfg == nil {
45+
return
46+
}
47+
48+
seen := make(map[string]struct{}, len(cfg.VertexCompatAPIKey))
49+
out := cfg.VertexCompatAPIKey[:0]
50+
for i := range cfg.VertexCompatAPIKey {
51+
entry := cfg.VertexCompatAPIKey[i]
52+
entry.APIKey = strings.TrimSpace(entry.APIKey)
53+
if entry.APIKey == "" {
54+
continue
55+
}
56+
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
57+
if entry.BaseURL == "" {
58+
// BaseURL is required for vertex-compat keys
59+
continue
60+
}
61+
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
62+
entry.Headers = NormalizeHeaders(entry.Headers)
63+
64+
// Sanitize models: remove entries without valid alias
65+
sanitizedModels := make([]VertexCompatModel, 0, len(entry.Models))
66+
for _, model := range entry.Models {
67+
model.Alias = strings.TrimSpace(model.Alias)
68+
model.Name = strings.TrimSpace(model.Name)
69+
if model.Alias != "" && model.Name != "" {
70+
sanitizedModels = append(sanitizedModels, model)
71+
}
72+
}
73+
entry.Models = sanitizedModels
74+
75+
// Use API key + base URL as uniqueness key
76+
uniqueKey := entry.APIKey + "|" + entry.BaseURL
77+
if _, exists := seen[uniqueKey]; exists {
78+
continue
79+
}
80+
seen[uniqueKey] = struct{}{}
81+
out = append(out, entry)
82+
}
83+
cfg.VertexCompatAPIKey = out
84+
}

0 commit comments

Comments
 (0)