CamTalk/backend/internal/ai/llm/openai.go

package llm

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

	"go.uber.org/zap"
)

// OpenAIService 基于 OpenAI Chat Completions API 的 LLM 实现。
type OpenAIService struct {
	apiKey   string
	model    string
	endpoint string
	timeout  time.Duration
	logger   *zap.SugaredLogger
	client   *http.Client
}

// NewOpenAIService 创建 OpenAI LLM 服务。
// model、endpoint 由 config 层保证非空。
func NewOpenAIService(apiKey, model, endpoint string, timeoutSec, httpClientTimeoutSec int, logger *zap.SugaredLogger) *OpenAIService {
	timeout := time.Duration(timeoutSec) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	httpClientTimeout := time.Duration(httpClientTimeoutSec) * time.Second
	if httpClientTimeout <= 0 {
		httpClientTimeout = 60 * time.Second
	}
	return &OpenAIService{
		apiKey:   apiKey,
		model:    model,
		endpoint: endpoint,
		timeout:  timeout,
		logger:   logger,
		client:   &http.Client{Timeout: httpClientTimeout},
	}
}

// --- OpenAI API 请求/响应结构 ---

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
	Stream   bool          `json:"stream"`
}

type chatMessage struct {
	Role    string        `json:"role"`
	Content []contentPart `json:"content"`
}

type contentPart struct {
	Type     string    `json:"type"`
	Text     string    `json:"text"`
	ImageURL *imageURL `json:"image_url,omitempty"`
}

type imageURL struct {
	URL string `json:"url"`
}

// streamDelta SSE 流式响应的单个 delta。
type streamDelta struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Usage *struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
	Model string `json:"model"`
}

// ChatStream 实现 llm.Service。调用 OpenAI Chat Completions API 流式推理。
func (o *OpenAIService) ChatStream(ctx context.Context, req Request) (<-chan Chunk, error) {
	// 构建请求
	messages := o.buildMessages(req)

	body := chatRequest{
		Model:    o.model,
		Messages: messages,
		Stream:   true,
	}

	payload, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("llm: marshal request: %w", err)
	}
	if err != nil {
		return nil, fmt.Errorf("llm: marshal request: %w", err)
	}

	// 创建带超时的 context
	ctx, cancel := context.WithTimeout(ctx, o.timeout)

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint+"/chat/completions", bytes.NewReader(payload))
	if err != nil {
		cancel()
		return nil, fmt.Errorf("llm: create request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+o.apiKey)

	resp, err := o.client.Do(httpReq)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("llm: send request: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		cancel()
		bodyBytes, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("llm: api error (status %d): %s", resp.StatusCode, string(bodyBytes))
	}

	// 启动 goroutine 解析 SSE 流
	ch := make(chan Chunk, 64)
	go func() {
		defer close(ch)
		defer cancel()
		defer resp.Body.Close()

		o.parseSSEStream(resp.Body, ch)
	}()

	return ch, nil
}

// parseSSEStream 解析 SSE 流，将 delta 发送到 channel。
func (o *OpenAIService) parseSSEStream(body io.Reader, ch chan<- Chunk) {
	scanner := bufio.NewScanner(body)
	scanner.Buffer(make([]byte, 0, 64*1024), 256*1024)

	var fullText strings.Builder
	var lastModel string

	for scanner.Scan() {
		line := scanner.Text()

		// SSE 格式：data: {...}
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			// 流结束，发送最终 chunk
			ch <- Chunk{Delta: "", Done: true, Model: lastModel}
			return
		}

		var delta streamDelta
		if err := json.Unmarshal([]byte(data), &delta); err != nil {
			o.logger.Warnw("llm: unmarshal delta failed", "error", err, "data", data)
			continue
		}

		if delta.Model != "" {
			lastModel = delta.Model
		}

		// 提取增量文本
		if len(delta.Choices) > 0 {
			content := delta.Choices[0].Delta.Content
			if content != "" {
				fullText.WriteString(content)
				ch <- Chunk{Delta: content, Done: false, Model: lastModel}
			}

			// 某些模型在最后一个 choice 中携带 usage
			if delta.Choices[0].FinishReason != nil && delta.Usage != nil {
				ch <- Chunk{
					Delta: "",
					Done:  true,
					Model: lastModel,
					TokensUsed: &TokenUsage{
						Prompt:     delta.Usage.PromptTokens,
						Completion: delta.Usage.CompletionTokens,
						Total:      delta.Usage.TotalTokens,
					},
				}
				return
			}
		}
	}

	// scanner 结束但没收到 [DONE]
	if err := scanner.Err(); err != nil {
		o.logger.Warnw("llm: scan error", "error", err)
	}
	ch <- Chunk{Delta: "", Done: true, Model: lastModel}
}

// buildMessages 构建 OpenAI Chat API 的 messages 数组。
func (o *OpenAIService) buildMessages(req Request) []chatMessage {
	var messages []chatMessage

	// System prompt
	messages = append(messages, chatMessage{
		Role:    "system",
		Content: []contentPart{{Type: "text", Text: BuildSystemPrompt(req.Language, "")}},
	})

	// 历史消息
	for _, msg := range req.History {
		messages = append(messages, chatMessage{
			Role:    msg.Role,
			Content: []contentPart{{Type: "text", Text: msg.Content}},
		})
	}

	// 当前用户消息（图像 + 文本）
	var parts []contentPart
	if len(req.Image) > 0 {
		b64 := base64.StdEncoding.EncodeToString(req.Image)
		parts = append(parts, contentPart{
			Type:     "image_url",
			ImageURL: &imageURL{URL: "data:image/jpeg;base64," + b64},
		})
	}
	parts = append(parts, contentPart{Type: "text", Text: req.Text})
	messages = append(messages, chatMessage{Role: "user", Content: parts})

	return messages
}