init

internal/chatter/gemini.go

@@ -0,0 +1,158 @@
+package chatter
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"reflect"
+
+	"github.com/google/generative-ai-go/genai"
+	"google.golang.org/api/option"
+)
+
+type GeminiAdapter struct {
+	client *genai.Client
+	model  string
+}
+
+func NewGeminiAdapter(ctx context.Context, apiKey string, modelName string) (*GeminiAdapter, error) {
+	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create gemini client: %w", err)
+	}
+
+	return &GeminiAdapter{
+		client: client,
+		model:  modelName,
+	}, nil
+}
+
+func (a *GeminiAdapter) Generate(ctx context.Context, messages []Message) (string, error) {
+	model := a.client.GenerativeModel(a.model)
+	cs := model.StartChat()
+
+	// Convert history (excluding the last message which is the current prompt)
+	var history []*genai.Content
+	for i := 0; i < len(messages)-1; i++ {
+		history = append(history, &genai.Content{
+			Role:  mapRoleToGemini(messages[i].Role),
+			Parts: []genai.Part{genai.Text(messages[i].Content)},
+		})
+	}
+	cs.History = history
+
+	// Send the last message
+	lastMsg := messages[len(messages)-1].Content
+	resp, err := cs.SendMessage(ctx, genai.Text(lastMsg))
+	if err != nil {
+		return "", fmt.Errorf("gemini generation failed: %w", err)
+	}
+
+	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
+		return "", fmt.Errorf("empty response from gemini")
+	}
+
+	// Extract text from the first part of the first candidate
+	if part, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
+		return string(part), nil
+	}
+
+	return "", fmt.Errorf("unexpected response format")
+}
+
+// GenerateStructured reflects the 'target' to create a schema, then unmarshals the result.
+func (a *GeminiAdapter) GenerateStructured(ctx context.Context, messages []Message, target any) error {
+	t := reflect.TypeOf(target)
+
+	model := a.client.GenerativeModel(a.model)
+
+	// 1. Automatically generate the JSON Schema from the Go struct
+	// 1. Recursively map the Go struct to Gemini's Schema format
+	model.ResponseMIMEType = "application/json"
+	model.ResponseSchema = schemaFromType(t.Elem()) // 2. Convert to the internal genai.Schema format
+
+	var prompt []genai.Part
+	for _, m := range messages {
+		prompt = append(prompt, genai.Text(m.Content))
+	}
+
+	resp, err := model.GenerateContent(ctx, prompt...)
+	if err != nil {
+		return err
+	}
+
+	part := resp.Candidates[0].Content.Parts[0].(genai.Text)
+	err = json.Unmarshal([]byte(part), target)
+	return err
+}
+
+// schemaFromType recursively builds a *genai.Schema tree from a reflect.Type
+func schemaFromType(t reflect.Type) *genai.Schema {
+	// Follow pointers to the base type
+	for t.Kind() == reflect.Ptr {
+		t = t.Elem()
+	}
+
+	switch t.Kind() {
+	case reflect.Struct:
+		props := make(map[string]*genai.Schema)
+		var required []string
+
+		for i := 0; i < t.NumField(); i++ {
+			field := t.Field(i)
+			if !field.IsExported() {
+				continue
+			}
+
+			// Get JSON tag or use field name
+			name := field.Tag.Get("json")
+			if name == "" || name == "-" {
+				name = field.Name
+			}
+
+			props[name] = schemaFromType(field.Type)
+			required = append(required, name)
+		}
+
+		return &genai.Schema{
+			Type:       genai.TypeObject,
+			Properties: props,
+			Required:   required,
+		}
+
+	case reflect.Slice, reflect.Array:
+		return &genai.Schema{
+			Type:  genai.TypeArray,
+			Items: schemaFromType(t.Elem()),
+		}
+
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
+		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		return &genai.Schema{Type: genai.TypeInteger}
+
+	case reflect.Float32, reflect.Float64:
+		return &genai.Schema{Type: genai.TypeNumber}
+
+	case reflect.Bool:
+		return &genai.Schema{Type: genai.TypeBoolean}
+
+	default:
+		return &genai.Schema{Type: genai.TypeString}
+	}
+}
+
+func (a *GeminiAdapter) GetProviderName() string {
+	return "Google Gemini (" + a.model + ")"
+}
+
+// Helper to map your roles to Gemini's expected roles
+func mapRoleToGemini(role Role) string {
+	switch role {
+	case RoleUser:
+		return "user"
+	case RoleAssistant:
+		return "model"
+	default:
+		return "user" // Gemini doesn't have a specific 'system' role in chat history
+	}
+}