package ai import ( "bufio" "bytes" "encoding/json" "fmt" "io" "net/http" "strings" "time" "eaisalestrain/backend/internal/config" "eaisalestrain/backend/internal/model" "eaisalestrain/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) } return c.hc.Do(req) } // ────────────────────────────────────────────── // 非流式调用 // ────────────────────────────────────────────── // 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 链调用(参考 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) } // ────────────────────────────────────────────── // 配置解析(兼容旧接口) // ────────────────────────────────────────────── // 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]) + "..." }