google.go

package goai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

func init() {
	register("google", newGoogleProvider)
}

// googleProvider implements Provider for Google Gemini models
type googleProvider struct {
	apiKey      string
	modelID     string
	apiEndpoint string
}

// newGoogleProvider creates a new Google provider
func newGoogleProvider(model string, opts Options) (Provider, error) {
	if opts.APIKey == "" {
		return nil, fmt.Errorf("%w: google", ErrNoAPIKey)
	}
	
	endpoint := opts.Endpoint
	if endpoint == "" {
		// Google uses the API key in the URL
		endpoint = fmt.Sprintf("https://generativelanguage.googleapis.com/v1beta/models/%s:generateContent?key=%s", model, opts.APIKey)
	}
	
	return &googleProvider{
		apiKey:      opts.APIKey,
		modelID:     model,
		apiEndpoint: endpoint,
	}, nil
}

// chat sends a prompt to Google Gemini and returns the response
func (p *googleProvider) chat(ctx context.Context, prompt string, opts Options) (*Response, error) {
	start := time.Now()
	
	requestBody := map[string]interface{}{
		"contents": []map[string]interface{}{
			{
				"parts": []map[string]string{
					{"text": prompt},
				},
			},
		},
		"generationConfig": map[string]interface{}{
			"temperature":    opts.Temperature,
			"maxOutputTokens": opts.MaxTokens,
		},
	}
	
	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		return nil, &Error{
			Provider: "google",
			Message:  fmt.Sprintf("failed to marshal request: %v", err),
			Err:      err,
		}
	}
	
	req, err := http.NewRequestWithContext(ctx, "POST", p.apiEndpoint, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, &Error{
			Provider: "google",
			Message:  fmt.Sprintf("failed to create request: %v", err),
			Err:      err,
		}
	}
	
	req.Header.Set("Content-Type", "application/json")
	
	client := &http.Client{Timeout: opts.Timeout}
	resp, err := client.Do(req)
	if err != nil {
		return nil, &Error{
			Provider: "google",
			Message:  fmt.Sprintf("request failed: %v", err),
			Err:      err,
		}
	}
	defer resp.Body.Close()
	
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, &Error{
			Provider: "google",
			Message:  fmt.Sprintf("failed to read response: %v", err),
			Err:      err,
		}
	}
	
	if resp.StatusCode != http.StatusOK {
		return nil, &Error{
			Provider:   "google",
			StatusCode: resp.StatusCode,
			Message:    string(body),
		}
	}
	
	var apiResponse struct {
		Candidates []struct {
			Content struct {
				Parts []struct {
					Text string `json:"text"`
				} `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
		UsageMetadata struct {
			TotalTokenCount int `json:"totalTokenCount"`
		} `json:"usageMetadata"`
	}
	
	if err := json.Unmarshal(body, &apiResponse); err != nil {
		return nil, &Error{
			Provider: "google",
			Message:  fmt.Sprintf("failed to parse response: %v", err),
			Err:      err,
		}
	}
	
	if len(apiResponse.Candidates) == 0 || len(apiResponse.Candidates[0].Content.Parts) == 0 {
		return nil, &Error{
			Provider: "google",
			Message:  "no response from Google",
			Err:      ErrNoResponse,
		}
	}
	
	return &Response{
		Content:    apiResponse.Candidates[0].Content.Parts[0].Text,
		TokensUsed: apiResponse.UsageMetadata.TotalTokenCount,
		LatencyMs:  time.Since(start).Milliseconds(),
		Model:      p.modelID,
		Provider:   "google",
		Metadata:   map[string]interface{}{"model": p.modelID},
	}, nil
}

// name returns the provider name
func (p *googleProvider) name() string {
	return "google"
}

// model returns the model identifier
func (p *googleProvider) model() string {
	return p.modelID
}