Files
2026-08-14 23:41:57 +08:00

471 lines
13 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package agent
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net/http"
"strings"
"time"
)
type LLMClient struct {
cfg Config
apiConfig *APIConfigStore
client *http.Client
}
type StreamChunk struct {
Content string
ToolCalls []ToolCallDelta
FinishReason string
Usage *Usage
Error error
}
type ToolCallDelta struct {
Index int
ID string
Name string
ArgumentsDelta string
}
func NewLLMClient(cfg Config, apiCfg *APIConfigStore) *LLMClient {
return &LLMClient{
cfg: cfg,
apiConfig: apiCfg,
client: &http.Client{
Timeout: cfg.RequestTimeout + 30*time.Second,
},
}
}
// isAnthropic 报告当前是否使用 Anthropic Messages API。
func (c *LLMClient) isAnthropic() bool {
return c.apiConfig.Provider() == ProviderAnthropic
}
func (c *LLMClient) ChatStream(ctx context.Context, messages []Message, tools []ToolDefinition) (<-chan StreamChunk, error) {
if c.isAnthropic() {
body := map[string]any{
"model": c.apiConfig.Model(),
"max_tokens": c.cfg.StreamOutputTokens,
"stream": true,
"temperature": c.cfg.Temperature,
"messages": toAnthropicMessages(messages),
}
if system := extractSystemPrompt(messages); system != "" {
body["system"] = system
}
if len(tools) > 0 {
body["tools"] = toAnthropicTools(tools)
}
return c.stream(ctx, body)
}
body := map[string]any{
"model": c.apiConfig.Model(),
"messages": toAPIMessages(messages),
"stream": true,
"temperature": c.cfg.Temperature,
"max_tokens": c.cfg.StreamOutputTokens,
}
if len(tools) > 0 {
body["tools"] = tools
body["tool_choice"] = "auto"
}
return c.stream(ctx, body)
}
func (c *LLMClient) Compress(ctx context.Context, history []Message) (string, error) {
systemPrompt := `你是一个高信息密度的对话压缩器。把下面的历史对话压缩成一份中文摘要,供后续继续执行同一任务。
必须保留(按重要程度排序):
1. 用户的核心目标与最新要求、尚未完成的任务与下一步计划;
2. 已确认的决策与结论、得分/进度(如 TSec 累计分数与已通关/进行中的题目编号);
3. 关键文件路径、执行过的命令与重要输出要点、代码要点;
4. 错误信息与解决方案、需要继续跟进的问题;
5. 平台提交规则(如 BENCHMARK_TOKEN、unique_code、提交接口与提交方式)。
规则:
- 按时间顺序组织,最新信息优先,可适当合并同类项;
- 若历史包含工具调用(run_bash/read_file/run_coop 等),只保留"做了什么、结果如何"的要点,不要逐字复制命令或输出;
- 不要添加历史中不存在的信息,不要臆测;
- 直接输出摘要正文,不要输出任何解释、标题或多余格式。`
messages := append([]Message{{Role: "system", Content: &systemPrompt}}, history...)
if c.isAnthropic() {
return c.compressAnthropic(ctx, messages)
}
body := map[string]any{
"model": c.apiConfig.Model(),
"messages": toAPIMessages(messages),
"stream": false,
"temperature": 0.2,
"max_tokens": c.cfg.CompactionTokens,
}
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
defer cancel()
payload, err := json.Marshal(body)
if err != nil {
return "", err
}
respBody, err := c.doSyncRequestWithRetry(ctx, false, c.apiConfig.BaseURL(), payload)
if err != nil {
return "", err
}
var result struct {
Choices []struct {
Message struct {
Content *string `json:"content"`
ReasoningContent *string `json:"reasoning_content"`
} `json:"message"`
} `json:"choices"`
}
if err := json.Unmarshal(respBody, &result); err != nil {
return "", err
}
if len(result.Choices) == 0 {
return "", errors.New("模型没有返回压缩摘要")
}
msg := result.Choices[0].Message
if msg.Content != nil && strings.TrimSpace(*msg.Content) != "" {
return strings.TrimSpace(*msg.Content), nil
}
// 推理模型(deepseek-reasoner 系)正文在 reasoning_contentcontent 可能为空
if msg.ReasoningContent != nil && strings.TrimSpace(*msg.ReasoningContent) != "" {
return strings.TrimSpace(*msg.ReasoningContent), nil
}
return "", errors.New("模型没有返回压缩摘要")
}
// stream 发起流式 LLM 请求,对连接失败、限流(429/5xx/551)、首 chunk 前流断开
// 做指数退避重试(3 次:1s→2s→4s)。保留 (<-chan, error) 签名,上层无需改动。
func (c *LLMClient) stream(ctx context.Context, body map[string]any) (<-chan StreamChunk, error) {
anthropic := c.isAnthropic()
url := c.apiConfig.BaseURL()
if anthropic {
url = normalizeAnthropicURL(url)
}
payload, err := json.Marshal(body)
if err != nil {
return nil, err
}
const maxRetries = 3
var lastErr error
for attempt := 0; attempt <= maxRetries; attempt++ {
if attempt > 0 {
backoff := time.Duration(1<<(attempt-1)) * time.Second
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(backoff):
}
}
ch, reqErr := c.doStreamRequest(ctx, anthropic, url, payload)
if reqErr != nil {
lastErr = reqErr
if attempt < maxRetries && isRetriableErr(reqErr) {
log.Printf("[llm] 流式请求失败(第 %d 次),将重试: %v", attempt+1, reqErr)
continue
}
return nil, reqErr
}
// 请求成功建立,但流可能在首个 chunk 前就断开(EOF)。
// peek 第一个 chunk:如果是可重试错误,排空旧 channel 后重试。
first, ok := <-ch
if !ok {
lastErr = errors.New("stream: 空响应")
if attempt < maxRetries {
log.Printf("[llm] 流式响应为空(第 %d 次),重试", attempt+1)
continue
}
return nil, lastErr
}
if first.Error != nil && attempt < maxRetries && isRetriableErr(first.Error) {
for range ch {
}
lastErr = first.Error
log.Printf("[llm] 流式响应首 chunk 前断开(第 %d 次),重试: %v", attempt+1, first.Error)
continue
}
// 正常:转发 first + 后续 chunk
out := make(chan StreamChunk, 64)
go func() {
defer close(out)
out <- first
for chunk := range ch {
out <- chunk
}
}()
return out, nil
}
return nil, lastErr
}
// doStreamRequest 发起单次流式 LLM 请求(不含重试)。
func (c *LLMClient) doStreamRequest(ctx context.Context, anthropic bool, url string, payload []byte) (<-chan StreamChunk, error) {
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if err != nil {
cancel()
return nil, err
}
c.setRequestHeaders(req, anthropic)
resp, err := c.client.Do(req)
if err != nil {
cancel()
return nil, err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
apiErr := c.readAPIError(resp)
resp.Body.Close()
cancel()
return nil, apiErr
}
ch := make(chan StreamChunk, 64)
go func() {
defer cancel()
defer close(ch)
defer resp.Body.Close()
if anthropic {
readAnthropicSSE(ctx, resp, ch)
} else {
c.readSSE(ctx, resp, ch)
}
}()
return ch, nil
}
// doSyncRequestWithRetry 发起同步(非流式)LLM 请求,对连接失败、限流(429/5xx/551)
// 做指数退避重试(3 次:1s→2s→4s)。返回响应体字节,由调用方解析。
func (c *LLMClient) doSyncRequestWithRetry(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) {
const maxRetries = 3
var lastErr error
for attempt := 0; attempt <= maxRetries; attempt++ {
if attempt > 0 {
backoff := time.Duration(1<<(attempt-1)) * time.Second
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(backoff):
}
}
body, err := c.doSyncRequest(ctx, anthropic, url, payload)
if err == nil {
return body, nil
}
lastErr = err
if attempt < maxRetries && isRetriableErr(err) {
log.Printf("[llm] 同步请求失败(第 %d 次),将重试: %v", attempt+1, err)
continue
}
return nil, err
}
return nil, lastErr
}
// doSyncRequest 发起单次同步(非流式)LLM 请求(不含重试)。
func (c *LLMClient) doSyncRequest(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) {
ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload))
if err != nil {
return nil, err
}
c.setRequestHeaders(req, anthropic)
resp, err := c.client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, c.readAPIError(resp)
}
return io.ReadAll(resp.Body)
}
// isRetriableErr 判断错误是否值得重试(网络错误、限流、服务端错误)。
func isRetriableErr(err error) bool {
if err == nil {
return false
}
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
return true
}
msg := err.Error()
if strings.Contains(msg, "EOF") ||
strings.Contains(msg, "connection reset") ||
strings.Contains(msg, "broken pipe") ||
strings.Contains(msg, "timeout") ||
strings.Contains(msg, "deadline exceeded") ||
strings.Contains(msg, "connection refused") ||
strings.Contains(msg, "no such host") {
return true
}
// HTTP 状态码错误:5xx 服务端错误、429 限流、551 网关熔断
if strings.HasPrefix(msg, "HTTP 5") ||
strings.HasPrefix(msg, "HTTP 429") ||
strings.HasPrefix(msg, "HTTP 551") {
return true
}
return false
}
func (c *LLMClient) readSSE(ctx context.Context, resp *http.Response, ch chan<- StreamChunk) {
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
}
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "[DONE]" {
break
}
var chunk struct {
Choices []struct {
Delta struct {
Content *string `json:"content"`
ToolCalls []struct {
Index *int `json:"index"`
ID string `json:"id"`
Function struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
} `json:"function"`
} `json:"tool_calls"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
Usage *Usage `json:"usage"`
}
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue
}
event := StreamChunk{}
for _, choice := range chunk.Choices {
if choice.Delta.Content != nil {
event.Content += *choice.Delta.Content
}
for _, tc := range choice.Delta.ToolCalls {
index := 0
if tc.Index != nil {
index = *tc.Index
}
event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
Index: index,
ID: tc.ID,
Name: tc.Function.Name,
ArgumentsDelta: tc.Function.Arguments,
})
}
if choice.FinishReason != nil {
event.FinishReason = *choice.FinishReason
}
}
if chunk.Usage != nil {
event.Usage = chunk.Usage
}
if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil {
ch <- event
}
}
if err := scanner.Err(); err != nil && ctx.Err() == nil {
ch <- StreamChunk{Error: err}
}
}
func (c *LLMClient) setRequestHeaders(req *http.Request, anthropic bool) {
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
if anthropic {
req.Header.Set("x-api-key", c.apiConfig.APIKey())
req.Header.Set("anthropic-version", "2023-06-01")
} else {
req.Header.Set("Authorization", "Bearer "+c.apiConfig.APIKey())
}
}
func (c *LLMClient) readAPIError(resp *http.Response) error {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024))
var apiErr struct {
Error struct {
Message string `json:"message"`
} `json:"error"`
}
_ = json.Unmarshal(body, &apiErr)
message := strings.TrimSpace(apiErr.Error.Message)
if message == "" {
message = strings.TrimSpace(string(body))
}
if message == "" {
message = fmt.Sprintf("HTTP %d", resp.StatusCode)
}
return errors.New(message)
}
func toAPIMessages(messages []Message) []map[string]any {
result := make([]map[string]any, 0, len(messages))
for _, message := range messages {
item := map[string]any{
"role": message.Role,
"content": contentValue(message.Content),
}
if len(message.ToolCalls) > 0 {
item["tool_calls"] = toAPIToolCalls(message.ToolCalls)
}
if message.ToolCallID != "" {
item["tool_call_id"] = message.ToolCallID
}
result = append(result, item)
}
return result
}
func toAPIToolCalls(calls []ToolCall) []map[string]any {
result := make([]map[string]any, 0, len(calls))
for _, call := range calls {
result = append(result, map[string]any{
"id": call.ID,
"type": "function",
"function": map[string]any{
"name": call.Function.Name,
"arguments": call.Function.Arguments,
},
})
}
return result
}
func contentValue(content *string) any {
if content == nil {
return nil
}
return *content
}