471 lines
13 KiB
Go
471 lines
13 KiB
Go
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_content，content 可能为空
|
||
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
|
||
}
|