hunjing/shenshi/internal/transform/gemini.go

// Gemini 出站适配器：IR → Google Generative Language streamGenerateContent（alt=sse）。
package transform

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"hunjing/shenshi/internal/ir"
)

// GeminiAdaptor Gemini generateContent 端点适配器。
type GeminiAdaptor struct {
	EndpointName  string
	BaseURL       string
	APIKey        string
	UpstreamModel string
	Client        *http.Client
}

// Name 实现 Adaptor。
func (g *GeminiAdaptor) Name() string { return g.EndpointName }

type gemPart struct {
	Text string `json:"text"`
}

type gemContent struct {
	Role  string    `json:"role,omitempty"` // user / model（Gemini 无 system 角色）
	Parts []gemPart `json:"parts"`
}

type gemGeneration struct {
	Temperature *float64 `json:"temperature,omitempty"`
	MaxTokens   *int     `json:"maxOutputTokens,omitempty"`
}

type gemRequest struct {
	Contents          []gemContent   `json:"contents"`
	Generation        *gemGeneration `json:"generationConfig,omitempty"`
	SystemInstruction *gemContent    `json:"systemInstruction,omitempty"`
}

type gemChunk struct {
	Candidates []struct {
		Content struct {
			Parts []gemPart `json:"parts"`
		} `json:"content"`
		FinishReason *string `json:"finishReason"`
	} `json:"candidates"`
	UsageMetadata *struct {
		PromptTokenCount     int `json:"promptTokenCount"`
		CandidatesTokenCount int `json:"candidatesTokenCount"`
	} `json:"usageMetadata"`
}

// Stream 实现 Adaptor：Gemini SSE 数据帧转 IR Delta。
func (g *GeminiAdaptor) Stream(ctx context.Context, req *ir.ChatRequest, out chan<- ir.Delta) error {
	body := gemRequest{}
	for _, m := range req.Messages {
		switch m.Role {
		case "system":
			if body.SystemInstruction == nil {
				body.SystemInstruction = &gemContent{}
			}
			body.SystemInstruction.Parts = append(body.SystemInstruction.Parts, gemPart{Text: m.Content})
		case "assistant":
			body.Contents = append(body.Contents, gemContent{Role: "model", Parts: []gemPart{{Text: m.Content}}})
		default: // user / tool 结果按 user 侧回传（工具循环随百工 M5 完整化）
			body.Contents = append(body.Contents, gemContent{Role: "user", Parts: []gemPart{{Text: m.Content}}})
		}
	}
	if req.Temperature != nil || req.MaxTokens != nil {
		body.Generation = &gemGeneration{Temperature: req.Temperature, MaxTokens: req.MaxTokens}
	}

	payload, err := json.Marshal(body)
	if err != nil {
		return err
	}
	url := strings.TrimRight(g.BaseURL, "/") +
		"/v1beta/models/" + g.UpstreamModel + ":streamGenerateContent?alt=sse"
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
	if err != nil {
		return err
	}
	httpReq.Header.Set("content-type", "application/json")
	httpReq.Header.Set("x-goog-api-key", g.APIKey)

	client := g.Client
	if client == nil {
		client = &http.Client{Timeout: 120 * time.Second}
	}
	resp, err := client.Do(httpReq)
	if err != nil {
		return &UpstreamError{StatusCode: 0, Message: err.Error()}
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode >= 400 {
		buf := new(bytes.Buffer)
		_, _ = buf.ReadFrom(resp.Body)
		return &UpstreamError{StatusCode: resp.StatusCode, Message: buf.String()}
	}

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if data == "" {
			continue
		}
		var chunk gemChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}
		if len(chunk.Candidates) > 0 {
			c := chunk.Candidates[0]
			for _, part := range c.Content.Parts {
				if part.Text != "" {
					out <- ir.Delta{Delta: part.Text}
				}
			}
			if c.FinishReason != nil {
				finish := "stop"
				if strings.EqualFold(*c.FinishReason, "SAFETY") || strings.EqualFold(*c.FinishReason, "RECITATION") {
					finish = "content_filter"
				}
				var up, uc int
				if chunk.UsageMetadata != nil {
					up, uc = chunk.UsageMetadata.PromptTokenCount, chunk.UsageMetadata.CandidatesTokenCount
				}
				out <- ir.Delta{FinishReason: finish, UsagePromptTokens: up, UsageCompletionTokens: uc}
			}
		}
	}
	return scanner.Err()
}