gooseek/backend/internal/llm/gemini.go

package llm

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

type GeminiClient struct {
	baseClient
	apiKey  string
	baseURL string
	client  *http.Client
}

func NewGeminiClient(cfg ProviderConfig) (*GeminiClient, error) {
	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "https://generativelanguage.googleapis.com/v1beta"
	}

	return &GeminiClient{
		baseClient: baseClient{
			providerID: cfg.ProviderID,
			modelKey:   cfg.ModelKey,
		},
		apiKey:  cfg.APIKey,
		baseURL: strings.TrimSuffix(baseURL, "/"),
		client:  &http.Client{},
	}, nil
}

type geminiRequest struct {
	Contents         []geminiContent        `json:"contents"`
	SystemInstruction *geminiContent        `json:"systemInstruction,omitempty"`
	GenerationConfig  geminiGenerationConfig `json:"generationConfig,omitempty"`
	Tools            []geminiTool           `json:"tools,omitempty"`
}

type geminiContent struct {
	Role  string       `json:"role,omitempty"`
	Parts []geminiPart `json:"parts"`
}

type geminiPart struct {
	Text string `json:"text,omitempty"`
}

type geminiGenerationConfig struct {
	MaxOutputTokens int     `json:"maxOutputTokens,omitempty"`
	Temperature     float64 `json:"temperature,omitempty"`
	TopP            float64 `json:"topP,omitempty"`
}

type geminiTool struct {
	FunctionDeclarations []geminiFunctionDecl `json:"functionDeclarations,omitempty"`
}

type geminiFunctionDecl struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	Parameters  interface{} `json:"parameters"`
}

type geminiStreamResponse struct {
	Candidates []struct {
		Content struct {
			Parts []struct {
				Text string `json:"text"`
			} `json:"parts"`
		} `json:"content"`
		FinishReason string `json:"finishReason,omitempty"`
	} `json:"candidates"`
}

func (c *GeminiClient) StreamText(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
	contents := make([]geminiContent, 0)
	var systemInstruction *geminiContent

	for _, m := range req.Messages {
		if m.Role == RoleSystem {
			systemInstruction = &geminiContent{
				Parts: []geminiPart{{Text: m.Content}},
			}
			continue
		}

		role := "user"
		if m.Role == RoleAssistant {
			role = "model"
		}

		contents = append(contents, geminiContent{
			Role:  role,
			Parts: []geminiPart{{Text: m.Content}},
		})
	}

	geminiReq := geminiRequest{
		Contents:          contents,
		SystemInstruction: systemInstruction,
		GenerationConfig: geminiGenerationConfig{
			MaxOutputTokens: req.Options.MaxTokens,
			Temperature:     req.Options.Temperature,
			TopP:            req.Options.TopP,
		},
	}

	if len(req.Tools) > 0 {
		decls := make([]geminiFunctionDecl, len(req.Tools))
		for i, t := range req.Tools {
			decls[i] = geminiFunctionDecl{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  t.Schema,
			}
		}
		geminiReq.Tools = []geminiTool{{FunctionDeclarations: decls}}
	}

	body, err := json.Marshal(geminiReq)
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("%s/models/%s:streamGenerateContent?key=%s&alt=sse",
		c.baseURL, c.modelKey, c.apiKey)

	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("gemini API error: %d - %s", resp.StatusCode, string(body))
	}

	ch := make(chan StreamChunk, 100)
	go func() {
		defer close(ch)
		defer resp.Body.Close()

		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			line := scanner.Text()
			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")

			var response geminiStreamResponse
			if err := json.Unmarshal([]byte(data), &response); err != nil {
				continue
			}

			if len(response.Candidates) > 0 {
				candidate := response.Candidates[0]
				for _, part := range candidate.Content.Parts {
					if part.Text != "" {
						ch <- StreamChunk{ContentChunk: part.Text}
					}
				}
				if candidate.FinishReason != "" {
					ch <- StreamChunk{FinishReason: candidate.FinishReason}
				}
			}
		}
	}()

	return ch, nil
}

func (c *GeminiClient) GenerateText(ctx context.Context, req StreamRequest) (string, error) {
	ch, err := c.StreamText(ctx, req)
	if err != nil {
		return "", err
	}
	return readAllChunks(ch), nil
}