// Package llm OpenAI 兼容 Chat 客户端(移植自 opsx-jenkins 生产实现)。 package llm import ( "bufio" "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strings" "time" ) type Client struct { Endpoint string APIKey string Model string MaxTokens int } func NewClient(endpoint, apiKey, model string, maxTokens int) *Client { if maxTokens <= 0 { maxTokens = 1200 } return &Client{ Endpoint: strings.TrimRight(endpoint, "/"), APIKey: apiKey, Model: model, MaxTokens: maxTokens, } } type chatMessage struct { Role string `json:"role"` Content string `json:"content"` } type chatRequest struct { Model string `json:"model"` Messages []chatMessage `json:"messages"` Temperature float64 `json:"temperature"` MaxTokens int `json:"max_tokens"` Stream bool `json:"stream,omitempty"` } type chatResponse struct { Choices []struct { Message struct { Content string `json:"content"` } `json:"message"` } `json:"choices"` Error *struct { Message string `json:"message"` } `json:"error"` } // Chat 单轮对话;endpoint/apiKey 为空时返回错误由调用方降级 func (c *Client) Chat(ctx context.Context, system, user string) (string, error) { if c.Endpoint == "" || c.APIKey == "" || c.Model == "" { return "", fmt.Errorf("LLM 未配置") } payload := chatRequest{ Model: c.Model, Messages: []chatMessage{ {Role: "system", Content: system}, {Role: "user", Content: user}, }, Temperature: 0.2, MaxTokens: c.MaxTokens, } body, err := json.Marshal(payload) if err != nil { return "", err } req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.Endpoint+"/chat/completions", bytes.NewReader(body)) if err != nil { return "", err } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.APIKey) resp, err := (&http.Client{Timeout: 90 * time.Second}).Do(req) if err != nil { return "", err } defer resp.Body.Close() data, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) if err != nil { return "", err } var out chatResponse if err := json.Unmarshal(data, &out); err != nil { return "", fmt.Errorf("响应解析失败: %s", truncate(string(data), 200)) } if out.Error != nil && out.Error.Message != "" { return "", fmt.Errorf("LLM 错误: %s", out.Error.Message) } if resp.StatusCode < 200 || resp.StatusCode >= 300 { return "", fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(data), 200)) } if len(out.Choices) == 0 { return "", fmt.Errorf("空响应") } return strings.TrimSpace(out.Choices[0].Message.Content), nil } func truncate(s string, n int) string { if len(s) <= n { return s } return s[:n] + "..." } // StreamDelta 单个流式增量 type StreamDelta struct { Content string // 正文增量 Reasoning string // 推理过程增量(DeepSeek R1 类模型) Usage *Usage // 最后一个 chunk 可能带用量 } // Usage token 用量 type Usage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` TotalTokens int `json:"total_tokens"` } type streamChunk struct { Choices []struct { Delta struct { Content string `json:"content"` Reasoning string `json:"reasoning_content"` FinishReason string `json:"finish_reason"` } `json:"delta"` } `json:"choices"` Usage *Usage `json:"usage"` } // ChatStream 流式对话:onDelta 逐段回调(正文/推理),返回完整正文。 // 流式请求不设整体超时,由 ctx 控制取消。 func (c *Client) ChatStream(ctx context.Context, system, user string, onDelta func(StreamDelta)) (string, error) { if c.Endpoint == "" || c.APIKey == "" || c.Model == "" { return "", fmt.Errorf("LLM 未配置") } payload := chatRequest{ Model: c.Model, Messages: []chatMessage{ {Role: "system", Content: system}, {Role: "user", Content: user}, }, Temperature: 0.2, MaxTokens: c.MaxTokens, Stream: true, } body, err := json.Marshal(payload) if err != nil { return "", err } req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.Endpoint+"/chat/completions", bytes.NewReader(body)) if err != nil { return "", err } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+c.APIKey) // 流式专用 client:Timeout 0,取消由 ctx 控制 resp, err := (&http.Client{}).Do(req) if err != nil { return "", err } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { data, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) return "", fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(data), 200)) } var full strings.Builder scanner := bufio.NewScanner(resp.Body) scanner.Buffer(make([]byte, 64*1024), 1024*1024) for scanner.Scan() { line := strings.TrimSpace(scanner.Text()) if !strings.HasPrefix(line, "data:") { continue } data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) if data == "" || data == "[DONE]" { continue } var chunk streamChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { continue } if len(chunk.Choices) == 0 { if chunk.Usage != nil && onDelta != nil { onDelta(StreamDelta{Usage: chunk.Usage}) } continue } d := chunk.Choices[0].Delta if d.Content != "" || d.Reasoning != "" { full.WriteString(d.Content) if onDelta != nil { onDelta(StreamDelta{Content: d.Content, Reasoning: d.Reasoning}) } } } if err := scanner.Err(); err != nil { if ctx.Err() != nil { return full.String(), nil // 客户端主动取消:返回已收到的内容 } return full.String(), err } if full.Len() == 0 { return "", fmt.Errorf("流式响应为空") } return full.String(), nil }