Scaffold project
This commit is contained in:
227
internal/providers/google/google.go
Normal file
227
internal/providers/google/google.go
Normal file
@@ -0,0 +1,227 @@
|
||||
package google
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"google.golang.org/genai"
|
||||
|
||||
"github.com/yourusername/go-llm-gateway/internal/api"
|
||||
"github.com/yourusername/go-llm-gateway/internal/config"
|
||||
)
|
||||
|
||||
const Name = "google"
|
||||
|
||||
// Provider implements the Google Generative AI integration.
|
||||
type Provider struct {
|
||||
cfg config.ProviderConfig
|
||||
client *genai.Client
|
||||
}
|
||||
|
||||
// New constructs a Provider using the provided configuration.
|
||||
func New(cfg config.ProviderConfig) *Provider {
|
||||
var client *genai.Client
|
||||
if cfg.APIKey != "" {
|
||||
var err error
|
||||
client, err = genai.NewClient(context.Background(), &genai.ClientConfig{
|
||||
APIKey: cfg.APIKey,
|
||||
})
|
||||
if err != nil {
|
||||
// Log error but don't fail construction - will fail on Generate
|
||||
fmt.Printf("warning: failed to create google client: %v\n", err)
|
||||
}
|
||||
}
|
||||
return &Provider{
|
||||
cfg: cfg,
|
||||
client: client,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Provider) Name() string { return Name }
|
||||
|
||||
// Generate routes the Open Responses request to Gemini.
|
||||
func (p *Provider) Generate(ctx context.Context, req *api.ResponseRequest) (*api.Response, error) {
|
||||
if p.cfg.APIKey == "" {
|
||||
return nil, fmt.Errorf("google api key missing")
|
||||
}
|
||||
if p.client == nil {
|
||||
return nil, fmt.Errorf("google client not initialized")
|
||||
}
|
||||
|
||||
model := chooseModel(req.Model, p.cfg.Model)
|
||||
|
||||
// Convert Open Responses messages to Gemini format
|
||||
var contents []*genai.Content
|
||||
|
||||
for _, msg := range req.Input {
|
||||
var parts []*genai.Part
|
||||
for _, block := range msg.Content {
|
||||
if block.Type == "input_text" {
|
||||
parts = append(parts, genai.NewPartFromText(block.Text))
|
||||
}
|
||||
}
|
||||
|
||||
role := "user"
|
||||
if msg.Role == "assistant" || msg.Role == "model" {
|
||||
role = "model"
|
||||
}
|
||||
|
||||
contents = append(contents, &genai.Content{
|
||||
Role: role,
|
||||
Parts: parts,
|
||||
})
|
||||
}
|
||||
|
||||
// Generate content
|
||||
resp, err := p.client.Models.GenerateContent(ctx, model, contents, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("google api error: %w", err)
|
||||
}
|
||||
|
||||
// Convert Gemini response to Open Responses format
|
||||
output := make([]api.Message, 0, 1)
|
||||
var text string
|
||||
|
||||
if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
|
||||
for _, part := range resp.Candidates[0].Content.Parts {
|
||||
if part != nil {
|
||||
text += part.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
output = append(output, api.Message{
|
||||
Role: "assistant",
|
||||
Content: []api.ContentBlock{
|
||||
{Type: "output_text", Text: text},
|
||||
},
|
||||
})
|
||||
|
||||
// Extract usage info if available
|
||||
var inputTokens, outputTokens int
|
||||
if resp.UsageMetadata != nil {
|
||||
inputTokens = int(resp.UsageMetadata.PromptTokenCount)
|
||||
outputTokens = int(resp.UsageMetadata.CandidatesTokenCount)
|
||||
}
|
||||
|
||||
return &api.Response{
|
||||
ID: uuid.NewString(),
|
||||
Object: "response",
|
||||
Created: time.Now().Unix(),
|
||||
Model: model,
|
||||
Provider: Name,
|
||||
Output: output,
|
||||
Usage: api.Usage{
|
||||
InputTokens: inputTokens,
|
||||
OutputTokens: outputTokens,
|
||||
TotalTokens: inputTokens + outputTokens,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GenerateStream handles streaming requests to Google.
|
||||
func (p *Provider) GenerateStream(ctx context.Context, req *api.ResponseRequest) (<-chan *api.StreamChunk, <-chan error) {
|
||||
chunkChan := make(chan *api.StreamChunk)
|
||||
errChan := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
defer close(chunkChan)
|
||||
defer close(errChan)
|
||||
|
||||
if p.cfg.APIKey == "" {
|
||||
errChan <- fmt.Errorf("google api key missing")
|
||||
return
|
||||
}
|
||||
if p.client == nil {
|
||||
errChan <- fmt.Errorf("google client not initialized")
|
||||
return
|
||||
}
|
||||
|
||||
model := chooseModel(req.Model, p.cfg.Model)
|
||||
|
||||
// Convert messages
|
||||
var contents []*genai.Content
|
||||
|
||||
for _, msg := range req.Input {
|
||||
var parts []*genai.Part
|
||||
for _, block := range msg.Content {
|
||||
if block.Type == "input_text" {
|
||||
parts = append(parts, genai.NewPartFromText(block.Text))
|
||||
}
|
||||
}
|
||||
|
||||
role := "user"
|
||||
if msg.Role == "assistant" || msg.Role == "model" {
|
||||
role = "model"
|
||||
}
|
||||
|
||||
contents = append(contents, &genai.Content{
|
||||
Role: role,
|
||||
Parts: parts,
|
||||
})
|
||||
}
|
||||
|
||||
// Create stream
|
||||
stream := p.client.Models.GenerateContentStream(ctx, model, contents, nil)
|
||||
|
||||
// Process stream
|
||||
for resp, err := range stream {
|
||||
if err != nil {
|
||||
errChan <- fmt.Errorf("google stream error: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
var text string
|
||||
if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
|
||||
for _, part := range resp.Candidates[0].Content.Parts {
|
||||
if part != nil {
|
||||
text += part.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
delta := &api.StreamDelta{}
|
||||
if text != "" {
|
||||
delta.Content = []api.ContentBlock{
|
||||
{Type: "output_text", Text: text},
|
||||
}
|
||||
}
|
||||
|
||||
streamChunk := &api.StreamChunk{
|
||||
Object: "response.chunk",
|
||||
Created: time.Now().Unix(),
|
||||
Model: model,
|
||||
Provider: Name,
|
||||
Delta: delta,
|
||||
}
|
||||
|
||||
select {
|
||||
case chunkChan <- streamChunk:
|
||||
case <-ctx.Done():
|
||||
errChan <- ctx.Err()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Send final chunk
|
||||
select {
|
||||
case chunkChan <- &api.StreamChunk{Object: "response.chunk", Done: true}:
|
||||
case <-ctx.Done():
|
||||
errChan <- ctx.Err()
|
||||
}
|
||||
}()
|
||||
|
||||
return chunkChan, errChan
|
||||
}
|
||||
|
||||
func chooseModel(requested, defaultModel string) string {
|
||||
if requested != "" {
|
||||
return requested
|
||||
}
|
||||
if defaultModel != "" {
|
||||
return defaultModel
|
||||
}
|
||||
return "gemini-2.0-flash-exp"
|
||||
}
|
||||
Reference in New Issue
Block a user