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"
	"eai_agentplatform/backend/internal/model"
	"eai_agentplatform/backend/internal/store"
)

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

// 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(route *config.RouteConfig) *Client {
	return &Client{
		baseURL:     strings.TrimRight(route.BaseURL, "/"),
		apiKey:      route.APIKey,
		model:       route.Model,
		maxTokens:   route.MaxTokens,
		temperature: route.Temperature,
		hc:          &http.Client{Timeout: 120 * time.Second},
	}
}

// NewClientLegacy 兼容旧接口（从 LLMConfig 创建）
func NewClientLegacy(cfg LLMConfig) *Client {
	return &Client{
		baseURL:     strings.TrimRight(cfg.BaseURL, "/"),
		apiKey:      cfg.APIKey,
		model:       cfg.Model,
		maxTokens:   cfg.MaxTokens,
		temperature: cfg.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)
	}
	// #region debug-point C:llm-post
	if payload, err := json.Marshal(map[string]any{
		"sessionId":    "knowledge-chat-401",
		"runId":        "pre-fix",
		"hypothesisId": "C",
		"location":     "backend-go/internal/ai/llm.go:Client.post:request",
		"msg":          "[DEBUG] llm client request",
		"data": map[string]any{
			"url":        c.url(path),
			"model":      c.model,
			"hasAPIKey":  c.apiKey != "",
			"authHeader": req.Header.Get("Authorization") != "",
		},
		"ts": time.Now().UnixMilli(),
	}); err == nil {
		go http.Post("http://127.0.0.1:7777/event", "application/json", strings.NewReader(string(payload)))
	}
	// #endregion
	resp, err := c.hc.Do(req)
	if err != nil {
		return nil, err
	}
	// #region debug-point C:llm-response
	if payload, err := json.Marshal(map[string]any{
		"sessionId":    "knowledge-chat-401",
		"runId":        "pre-fix",
		"hypothesisId": "C",
		"location":     "backend-go/internal/ai/llm.go:Client.post:response",
		"msg":          "[DEBUG] llm client response",
		"data": map[string]any{
			"url":    c.url(path),
			"status": resp.StatusCode,
		},
		"ts": time.Now().UnixMilli(),
	}); err == nil {
		go http.Post("http://127.0.0.1:7777/event", "application/json", strings.NewReader(string(payload)))
	}
	// #endregion
	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)
		// #region debug-point C:llm-non200
		if payload, err := json.Marshal(map[string]any{
			"sessionId":    "knowledge-chat-401",
			"runId":        "pre-fix",
			"hypothesisId": "C",
			"location":     "backend-go/internal/ai/llm.go:GenerateStream:non200",
			"msg":          "[DEBUG] llm stream non-200",
			"data": map[string]any{
				"status":      resp.StatusCode,
				"bodyPreview": truncate(string(data), 240),
				"model":       c.model,
				"baseURL":     c.baseURL,
			},
			"ts": time.Now().UnixMilli(),
		}); err == nil {
			go http.Post("http://127.0.0.1:7777/event", "application/json", strings.NewReader(string(payload)))
		}
		// #endregion
		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 链调用（参考 pj034 chat_completion_with_fallback）
// ──────────────────────────────────────────────

// GenerateWithFallback 依次尝试主路由 + 回退链，首个成功返回
func GenerateWithFallback(primary *config.RouteConfig, messages []Message) (string, error) {
	// 构建路由链
	chain := []*config.RouteConfig{primary}
	fallbacks, err := config.GetFallbackRoutes(primary.RouteID)
	if err == nil && len(fallbacks) > 0 {
		chain = append(chain, fallbacks...)
	}

	var lastErr error
	for _, route := range chain {
		client := NewClient(route)
		content, err := client.Generate(messages)
		if err == nil {
			return content, nil
		}
		lastErr = fmt.Errorf("[%s] %w", route.RouteID, err)
	}
	return "", fmt.Errorf("所有路由均失败: %w", lastErr)
}

// GenerateFullWithFallback 非流式 + 回退链：返回结果与「实际命中」的路由（用于审计）
func GenerateFullWithFallback(primary *config.RouteConfig, messages []Message) (*ChatResult, *config.RouteConfig, error) {
	chain := []*config.RouteConfig{primary}
	if fallbacks, err := config.GetFallbackRoutes(primary.RouteID); err == nil && len(fallbacks) > 0 {
		chain = append(chain, fallbacks...)
	}

	var lastErr error
	for _, route := range chain {
		client := NewClient(route)
		res, err := client.GenerateFull(messages)
		if err == nil {
			res.Provider = route.Provider
			return res, route, nil
		}
		lastErr = fmt.Errorf("[%s] %w", route.RouteID, err)
	}
	return nil, nil, fmt.Errorf("所有路由均失败: %w", lastErr)
}

// GenerateStreamWithFallback 流式 + 回退链：返回实际命中的路由
func GenerateStreamWithFallback(primary *config.RouteConfig, messages []Message, onChunk func(string)) (*config.RouteConfig, error) {
	chain := []*config.RouteConfig{primary}
	if fallbacks, err := config.GetFallbackRoutes(primary.RouteID); err == nil && len(fallbacks) > 0 {
		chain = append(chain, fallbacks...)
	}

	var lastErr error
	for _, route := range chain {
		if requiresAPIKey(route) && strings.TrimSpace(route.APIKey) == "" {
			lastErr = fmt.Errorf("[%s] 未配置 API Key", route.RouteID)
			continue
		}
		client := NewClient(route)
		if err := client.GenerateStream(messages, onChunk); err == nil {
			return route, nil
		} else {
			lastErr = fmt.Errorf("[%s] %w", route.RouteID, err)
		}
	}
	return nil, fmt.Errorf("所有路由均失败: %w", lastErr)
}

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

// ──────────────────────────────────────────────
// 配置解析（兼容旧接口）
// ──────────────────────────────────────────────

// LLMConfig 简单的 LLM 连接配置（旧版，逐步淘汰）
type LLMConfig struct {
	BaseURL     string
	APIKey      string
	Model       string
	EmbedModel  string
	MaxTokens   int
	Temperature float64
}

// ResolveLLM 旧版：从 DB/环境 解析 LLM 配置
func ResolveLLM(cfg *config.Config) (LLMConfig, bool) {
	get := func(key, def string) string {
		var sc model.SystemConfig
		if err := store.DB.Where("config_key = ?", key).First(&sc).Error; err == nil && strings.TrimSpace(sc.ConfigValue) != "" {
			return sc.ConfigValue
		}
		return def
	}

	baseURL := get("llm_base_url", cfg.LLMBaseURL)
	apiKey := get("llm_api_key", cfg.LLMAPIKey)
	modelName := get("llm_model", cfg.LLMModel)
	embedModel := get("embed_model", cfg.EmbedModel)

	// DB 空时降级到 JSON secrets（flat 格式，pj034 兼容）
	if baseURL == "" || apiKey == "" {
		if baseURL == "" {
			baseURL = config.GetProviderBaseURL("ollama")
		}
		if apiKey == "" {
			apiKey = config.GetProviderAPIKey("ollama")
		}
	}

	c := LLMConfig{
		BaseURL:     baseURL,
		APIKey:      apiKey,
		Model:       modelName,
		EmbedModel:  embedModel,
		MaxTokens:   2048,
		Temperature: 0.7,
	}
	if c.BaseURL == "" || c.Model == "" {
		return c, false
	}
	return c, true
}

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