pj0235-eai_agentplatform/eai_agentplatform/backend-go/internal/ai/llm.go

package ai

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

	"eai_agentplatform/backend/internal/config"
)

// ──────────────────────────────────────────────
// 基础类型
// ──────────────────────────────────────────────

// Message OpenAI 兼容消息
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// ChatResult 非流式调用结果
type ChatResult struct {
	Content      string
	Model        string
	Provider     string
	FinishReason string
	Usage        struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	}
}

// ──────────────────────────────────────────────
// OpenAI 兼容 HTTP 客户端
// ──────────────────────────────────────────────

// Client 低层 HTTP 客户端（基于 RouteConfig）
type Client struct {
	baseURL     string
	apiKey      string
	model       string
	maxTokens   int
	temperature float64
	hc          *http.Client
}

// NewClient 从 RouteConfig 创建客户端
func NewClient(aiRoute *config.RouteConfig) *Client {
	return &Client{
		baseURL:     strings.TrimRight(aiRoute.BaseURL, "/"),
		apiKey:      aiRoute.APIKey,
		model:       aiRoute.Model,
		maxTokens:   aiRoute.MaxTokens,
		temperature: aiRoute.Temperature,
		hc:          &http.Client{Timeout: 120 * time.Second},
	}
}

func (c *Client) url(path string) string {
	return c.baseURL + path
}

func (c *Client) headers() map[string]string {
	h := map[string]string{"Content-Type": "application/json"}
	if c.apiKey != "" {
		h["Authorization"] = "Bearer " + c.apiKey
	}
	return h
}

func (c *Client) post(path string, body any) (*http.Response, error) {
	j, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest(http.MethodPost, c.url(path), bytes.NewReader(j))
	if err != nil {
		return nil, err
	}
	for k, v := range c.headers() {
		req.Header.Set(k, v)
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// ──────────────────────────────────────────────
// 非流式调用
// ──────────────────────────────────────────────

// Generate 非流式对话，返回完整正文
func (c *Client) Generate(messages []Message) (string, error) {
	resp, err := c.post("/chat/completions", map[string]any{
		"model":       c.model,
		"messages":    messages,
		"stream":      false,
		"temperature": c.temperature,
		"max_tokens":  c.maxTokens,
	})
	if err != nil {
		return "", fmt.Errorf("LLM 服务不可达: %w", err)
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("读取 LLM 响应失败: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
	}

	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return "", fmt.Errorf("LLM 响应解析失败: %w", err)
	}
	if len(out.Choices) == 0 || strings.TrimSpace(out.Choices[0].Message.Content) == "" {
		return "", fmt.Errorf("LLM 返回空正文")
	}
	return out.Choices[0].Message.Content, nil
}

// GenerateFull 非流式调用，返回完整 ChatResult（含 usage）
func (c *Client) GenerateFull(messages []Message) (*ChatResult, error) {
	resp, err := c.post("/chat/completions", map[string]any{
		"model":       c.model,
		"messages":    messages,
		"stream":      false,
		"temperature": c.temperature,
		"max_tokens":  c.maxTokens,
	})
	if err != nil {
		return nil, fmt.Errorf("LLM 服务不可达: %w", err)
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取 LLM 响应失败: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
	}

	var raw struct {
		Model   string `json:"model"`
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
			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"`
	}
	if err := json.Unmarshal(data, &raw); err != nil {
		return nil, fmt.Errorf("LLM 响应解析失败: %w", err)
	}
	if len(raw.Choices) == 0 || strings.TrimSpace(raw.Choices[0].Message.Content) == "" {
		return nil, fmt.Errorf("LLM 返回空正文")
	}

	result := &ChatResult{
		Content:      raw.Choices[0].Message.Content,
		Model:        raw.Model,
		FinishReason: raw.Choices[0].FinishReason,
	}
	if raw.Usage != nil {
		result.Usage.PromptTokens = raw.Usage.PromptTokens
		result.Usage.CompletionTokens = raw.Usage.CompletionTokens
		result.Usage.TotalTokens = raw.Usage.TotalTokens
	}
	return result, nil
}

// ──────────────────────────────────────────────
// 流式调用
// ──────────────────────────────────────────────

// GenerateStream SSE 流式调用，逐 chunk 回调
func (c *Client) GenerateStream(messages []Message, onChunk func(string)) error {
	resp, err := c.post("/chat/completions", map[string]any{
		"model":          c.model,
		"messages":       messages,
		"stream":         true,
		"temperature":    c.temperature,
		"max_tokens":     c.maxTokens,
		"stream_options": map[string]any{"include_usage": true},
	})
	if err != nil {
		return fmt.Errorf("LLM 服务不可达: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		data, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("LLM 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
	}

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "[DONE]" {
			break
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			continue
		}
		if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
			onChunk(chunk.Choices[0].Delta.Content)
		}
	}
	return scanner.Err()
}

// ──────────────────────────────────────────────
// Embedding
// ──────────────────────────────────────────────

// Embed 批量向量化
func (c *Client) Embed(inputs []string) ([][]float64, error) {
	resp, err := c.post("/embeddings", map[string]any{
		"model": c.model, "input": inputs,
	})
	if err != nil {
		return nil, fmt.Errorf("embedding 服务不可达: %w", err)
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取 embedding 响应失败: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("embedding 返回 %d: %s", resp.StatusCode, truncate(string(data), 200))
	}

	var out struct {
		Data []struct {
			Embedding []float64 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return nil, fmt.Errorf("embedding 响应解析失败: %w", err)
	}
	res := make([][]float64, len(out.Data))
	for i, d := range out.Data {
		res[i] = d.Embedding
	}
	return res, nil
}

// ──────────────────────────────────────────────
// Fallback 链调用
// ──────────────────────────────────────────────

// buildRouteChain 构建「主路由 + 回退链」。
//
// primary 为 nil 时必须显式报错而不是继续：链里存 nil 会在循环里 NewClient(nil)，
// 而取 primary.RouteID 更是直接空指针。历史上 batch_extract / contract_review
// 就是直接传了 nil，导致这两个接口必 500。调用方应先用 config.GetRoute 解析路由。
func buildRouteChain(primary *config.RouteConfig) ([]*config.RouteConfig, error) {
	if primary == nil {
		return nil, fmt.Errorf("主路由为 nil：调用方未解析 AI 路由，请先用 config.GetRoute 取路由")
	}
	aiRouteChain := []*config.RouteConfig{primary}
	if fallbacks, err := config.GetFallbackRoutes(primary.RouteID); err == nil && len(fallbacks) > 0 {
		aiRouteChain = append(aiRouteChain, fallbacks...)
	}
	return aiRouteChain, nil
}

// GenerateWithFallback 依次尝试主路由 + 回退链，首个成功返回
func GenerateWithFallback(primary *config.RouteConfig, messages []Message) (string, error) {
	aiRouteChain, err := buildRouteChain(primary)
	if err != nil {
		return "", err
	}

	var lastErr error
	for _, aiRoute := range aiRouteChain {
		client := NewClient(aiRoute)
		start := time.Now()
		content, err := client.Generate(messages)
		if err == nil {
			config.ReportRouteHealth(config.RouteHealth{
				AIRouteID:     aiRoute.RouteID,
				Category:      aiRoute.Category,
				Healthy:       true,
				Checked:       true,
				LatencyMs:     time.Since(start).Milliseconds(),
				LastCheckedAt: time.Now(),
			})
			return content, nil
		}
		config.ReportRouteHealth(config.RouteHealth{
			AIRouteID:     aiRoute.RouteID,
			Category:      aiRoute.Category,
			Healthy:       false,
			Checked:       true,
			LatencyMs:     time.Since(start).Milliseconds(),
			LastCheckedAt: time.Now(),
			LastError:     err.Error(),
		})
		lastErr = fmt.Errorf("[%s] %w", aiRoute.RouteID, err)
	}
	return "", fmt.Errorf("所有路由均失败: %w", lastErr)
}

// GenerateFullWithFallback 非流式 + 回退链：返回结果与「实际命中」的路由（用于审计）
func GenerateFullWithFallback(primary *config.RouteConfig, messages []Message) (*ChatResult, *config.RouteConfig, error) {
	aiRouteChain, err := buildRouteChain(primary)
	if err != nil {
		return nil, nil, err
	}

	var lastErr error
	for _, aiRoute := range aiRouteChain {
		client := NewClient(aiRoute)
		start := time.Now()
		res, err := client.GenerateFull(messages)
		if err == nil {
			config.ReportRouteHealth(config.RouteHealth{
				AIRouteID:     aiRoute.RouteID,
				Category:      aiRoute.Category,
				Healthy:       true,
				Checked:       true,
				LatencyMs:     time.Since(start).Milliseconds(),
				LastCheckedAt: time.Now(),
			})
			res.Provider = aiRoute.Provider
			return res, aiRoute, nil
		}
		config.ReportRouteHealth(config.RouteHealth{
			AIRouteID:     aiRoute.RouteID,
			Category:      aiRoute.Category,
			Healthy:       false,
			Checked:       true,
			LatencyMs:     time.Since(start).Milliseconds(),
			LastCheckedAt: time.Now(),
			LastError:     err.Error(),
		})
		lastErr = fmt.Errorf("[%s] %w", aiRoute.RouteID, err)
	}
	return nil, nil, fmt.Errorf("所有路由均失败: %w", lastErr)
}

// GenerateStreamWithFallback 流式 + 回退链：返回实际命中的路由
func GenerateStreamWithFallback(primary *config.RouteConfig, messages []Message, onChunk func(string)) (*config.RouteConfig, error) {
	aiRouteChain, err := buildRouteChain(primary)
	if err != nil {
		return nil, err
	}

	var lastErr error
	for _, aiRoute := range aiRouteChain {
		if requiresAPIKey(aiRoute) && strings.TrimSpace(aiRoute.APIKey) == "" {
			config.ReportRouteHealth(config.RouteHealth{
				AIRouteID:     aiRoute.RouteID,
				Category:      aiRoute.Category,
				Healthy:       false,
				Checked:       true,
				LastCheckedAt: time.Now(),
				LastError:     "未配置 API Key",
			})
			lastErr = fmt.Errorf("[%s] 未配置 API Key", aiRoute.RouteID)
			continue
		}
		client := NewClient(aiRoute)
		start := time.Now()
		if err := client.GenerateStream(messages, onChunk); err == nil {
			config.ReportRouteHealth(config.RouteHealth{
				AIRouteID:     aiRoute.RouteID,
				Category:      aiRoute.Category,
				Healthy:       true,
				Checked:       true,
				LatencyMs:     time.Since(start).Milliseconds(),
				LastCheckedAt: time.Now(),
			})
			return aiRoute, nil
		} else {
			config.ReportRouteHealth(config.RouteHealth{
				AIRouteID:     aiRoute.RouteID,
				Category:      aiRoute.Category,
				Healthy:       false,
				Checked:       true,
				LatencyMs:     time.Since(start).Milliseconds(),
				LastCheckedAt: time.Now(),
				LastError:     err.Error(),
			})
			lastErr = fmt.Errorf("[%s] %w", aiRoute.RouteID, err)
		}
	}
	return nil, fmt.Errorf("所有路由均失败: %w", lastErr)
}

func requiresAPIKey(aiRoute *config.RouteConfig) bool {
	if aiRoute == nil {
		return false
	}
	baseURL := strings.ToLower(strings.TrimSpace(aiRoute.BaseURL))
	if strings.Contains(baseURL, "openrouter.ai") || strings.Contains(baseURL, "openai.com") {
		return true
	}
	provider := strings.ToLower(strings.TrimSpace(aiRoute.Provider))
	return provider == "openrouter" || provider == "openai"
}

func truncate(s string, n int) string {
	r := []rune(s)
	if len(r) <= n {
		return s
	}
	return string(r[:n]) + "..."
}