Add Vertex AI support

This commit is contained in:
2026-03-02 16:20:57 +00:00
parent 2188e3cba8
commit 38d44f104a
5 changed files with 58 additions and 11 deletions

View File

@@ -11,6 +11,7 @@ Simplify LLM integration by exposing a single, consistent API that routes reques
- **Azure OpenAI** (Azure-deployed models) - **Azure OpenAI** (Azure-deployed models)
- **Anthropic** (Claude) - **Anthropic** (Claude)
- **Google Generative AI** (Gemini) - **Google Generative AI** (Gemini)
- **Vertex AI** (Google Cloud-hosted Gemini models)
Instead of managing multiple SDK integrations in your application, call one endpoint and let the gateway handle provider-specific implementations. Instead of managing multiple SDK integrations in your application, call one endpoint and let the gateway handle provider-specific implementations.
@@ -24,7 +25,8 @@ latticelm (unified API)
├─→ OpenAI SDK ├─→ OpenAI SDK
├─→ Azure OpenAI (OpenAI SDK + Azure auth) ├─→ Azure OpenAI (OpenAI SDK + Azure auth)
├─→ Anthropic SDK ├─→ Anthropic SDK
─→ Google Gen AI SDK ─→ Google Gen AI SDK
└─→ Vertex AI (Google Gen AI SDK + GCP auth)
``` ```
## Key Features ## Key Features
@@ -45,11 +47,12 @@ latticelm (unified API)
## 🎉 Status: **WORKING!** ## 🎉 Status: **WORKING!**
**All four providers integrated with official Go SDKs:** **All providers integrated with official Go SDKs:**
- OpenAI → `github.com/openai/openai-go/v3` - OpenAI → `github.com/openai/openai-go/v3`
- Azure OpenAI → `github.com/openai/openai-go/v3` (with Azure auth) - Azure OpenAI → `github.com/openai/openai-go/v3` (with Azure auth)
- Anthropic → `github.com/anthropics/anthropic-sdk-go` - Anthropic → `github.com/anthropics/anthropic-sdk-go`
- Google → `google.golang.org/genai` - Google → `google.golang.org/genai`
- Vertex AI → `google.golang.org/genai` (with GCP auth)
**Compiles successfully** (36MB binary) **Compiles successfully** (36MB binary)
**Provider auto-selection** (gpt→Azure/OpenAI, claude→Anthropic, gemini→Google) **Provider auto-selection** (gpt→Azure/OpenAI, claude→Anthropic, gemini→Google)

View File

@@ -14,6 +14,12 @@ providers:
type: "openai" type: "openai"
api_key: "YOUR_OPENAI_API_KEY" api_key: "YOUR_OPENAI_API_KEY"
endpoint: "https://api.openai.com" endpoint: "https://api.openai.com"
# Vertex AI (Google Cloud) - optional
# Uses Application Default Credentials (ADC) or service account
# vertexai:
# type: "vertexai"
# project: "your-gcp-project-id"
# location: "us-central1" # or other GCP region
# Azure OpenAI - optional # Azure OpenAI - optional
# azureopenai: # azureopenai:
# type: "azureopenai" # type: "azureopenai"
@@ -48,6 +54,8 @@ models:
provider: "anthropic" provider: "anthropic"
- name: "gpt-4o-mini" - name: "gpt-4o-mini"
provider: "openai" provider: "openai"
# - name: "gemini-2.0-flash-exp"
# provider: "vertexai" # Use Vertex AI instead of Google AI API
# - name: "gpt-4o" # - name: "gpt-4o"
# provider: "azureopenai" # provider: "azureopenai"
# provider_model_id: "my-gpt4o-deployment" # optional: defaults to name # provider_model_id: "my-gpt4o-deployment" # optional: defaults to name

View File

@@ -48,6 +48,8 @@ type ProviderEntry struct {
APIKey string `yaml:"api_key"` APIKey string `yaml:"api_key"`
Endpoint string `yaml:"endpoint"` Endpoint string `yaml:"endpoint"`
APIVersion string `yaml:"api_version"` APIVersion string `yaml:"api_version"`
Project string `yaml:"project"` // For Vertex AI
Location string `yaml:"location"` // For Vertex AI
} }
// ModelEntry maps a model name to a provider entry. // ModelEntry maps a model name to a provider entry.
@@ -78,6 +80,12 @@ type AzureAnthropicConfig struct {
Model string `yaml:"model"` Model string `yaml:"model"`
} }
// VertexAIConfig contains Vertex AI-specific settings used internally by the Google provider.
type VertexAIConfig struct {
Project string `yaml:"project"`
Location string `yaml:"location"`
}
// Load reads and parses a YAML configuration file, expanding ${VAR} env references. // Load reads and parses a YAML configuration file, expanding ${VAR} env references.
func Load(path string) (*Config, error) { func Load(path string) (*Config, error) {
data, err := os.ReadFile(path) data, err := os.ReadFile(path)

View File

@@ -19,7 +19,7 @@ type Provider struct {
client *genai.Client client *genai.Client
} }
// New constructs a Provider using the provided configuration. // New constructs a Provider using the Google AI API with API key authentication.
func New(cfg config.ProviderConfig) *Provider { func New(cfg config.ProviderConfig) *Provider {
var client *genai.Client var client *genai.Client
if cfg.APIKey != "" { if cfg.APIKey != "" {
@@ -38,13 +38,36 @@ func New(cfg config.ProviderConfig) *Provider {
} }
} }
// NewVertexAI constructs a Provider targeting Vertex AI.
// Vertex AI uses the same genai SDK but with GCP project/location configuration
// and Application Default Credentials (ADC) or service account authentication.
func NewVertexAI(vertexCfg config.VertexAIConfig) *Provider {
var client *genai.Client
if vertexCfg.Project != "" && vertexCfg.Location != "" {
var err error
client, err = genai.NewClient(context.Background(), &genai.ClientConfig{
Project: vertexCfg.Project,
Location: vertexCfg.Location,
Backend: genai.BackendVertexAI,
})
if err != nil {
// Log error but don't fail construction - will fail on Generate
fmt.Printf("warning: failed to create vertex ai client: %v\n", err)
}
}
return &Provider{
cfg: config.ProviderConfig{
// Vertex AI doesn't use API key, but set empty for consistency
APIKey: "",
},
client: client,
}
}
func (p *Provider) Name() string { return Name } func (p *Provider) Name() string { return Name }
// Generate routes the request to Gemini and returns a ProviderResult. // Generate routes the request to Gemini and returns a ProviderResult.
func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) { func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
if p.cfg.APIKey == "" {
return nil, fmt.Errorf("google api key missing")
}
if p.client == nil { if p.client == nil {
return nil, fmt.Errorf("google client not initialized") return nil, fmt.Errorf("google client not initialized")
} }
@@ -96,10 +119,6 @@ func (p *Provider) GenerateStream(ctx context.Context, messages []api.Message, r
defer close(deltaChan) defer close(deltaChan)
defer close(errChan) defer close(errChan)
if p.cfg.APIKey == "" {
errChan <- fmt.Errorf("google api key missing")
return
}
if p.client == nil { if p.client == nil {
errChan <- fmt.Errorf("google client not initialized") errChan <- fmt.Errorf("google client not initialized")
return return

View File

@@ -60,7 +60,8 @@ func NewRegistry(entries map[string]config.ProviderEntry, models []config.ModelE
} }
func buildProvider(entry config.ProviderEntry) (Provider, error) { func buildProvider(entry config.ProviderEntry) (Provider, error) {
if entry.APIKey == "" { // Vertex AI doesn't require APIKey, so check for it separately
if entry.Type != "vertexai" && entry.APIKey == "" {
return nil, nil return nil, nil
} }
@@ -97,6 +98,14 @@ func buildProvider(entry config.ProviderEntry) (Provider, error) {
APIKey: entry.APIKey, APIKey: entry.APIKey,
Endpoint: entry.Endpoint, Endpoint: entry.Endpoint,
}), nil }), nil
case "vertexai":
if entry.Project == "" || entry.Location == "" {
return nil, fmt.Errorf("project and location are required for vertexai")
}
return googleprovider.NewVertexAI(config.VertexAIConfig{
Project: entry.Project,
Location: entry.Location,
}), nil
default: default:
return nil, fmt.Errorf("unknown provider type %q", entry.Type) return nil, fmt.Errorf("unknown provider type %q", entry.Type)
} }