314 lines
8.7 KiB
Go
314 lines
8.7 KiB
Go
package agent
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// ModelConfig 模型配置
|
||
type ModelConfig struct {
|
||
ModelName string // 对话模型名
|
||
APIKey string // API密钥
|
||
BaseURL string // API地址
|
||
MaxTokens int // 最大Token数
|
||
Temperature float32 // 温度参数
|
||
Timeout time.Duration // HTTP请求超时(0表示默认)
|
||
MaxRetries int // 最大重试次数(0表示默认3次)
|
||
}
|
||
|
||
// CallChatModel 调用大模型聊天接口（OpenAI 兼容格式）
|
||
func CallChatModel(ctx context.Context, cfg *ModelConfig, req *ChatRequest) (*ChatResponse, error) {
|
||
if cfg == nil {
|
||
return nil, fmt.Errorf("model config cannot be empty")
|
||
}
|
||
if cfg.APIKey == "" {
|
||
return nil, fmt.Errorf("APIKey not configured")
|
||
}
|
||
if cfg.ModelName == "" {
|
||
return nil, fmt.Errorf("model name not configured")
|
||
}
|
||
if cfg.BaseURL == "" {
|
||
return nil, fmt.Errorf("API address not configured")
|
||
}
|
||
|
||
timeout := cfg.Timeout
|
||
if timeout <= 0 {
|
||
timeout = 300 * time.Second
|
||
}
|
||
|
||
body, err := buildReqBody(cfg.ModelName, req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
url := trimSlashes(cfg.BaseURL)
|
||
|
||
var lastErr error
|
||
maxRetries := cfg.MaxRetries
|
||
if maxRetries <= 0 {
|
||
maxRetries = 3
|
||
}
|
||
g.Log().Debugf(ctx, "ChatAPI 开始调用 model=%s timeout=%v max_retries=%d body_size=%d", cfg.ModelName, timeout, maxRetries, len(body))
|
||
|
||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||
if attempt > 0 {
|
||
wait := time.Duration(1<<(attempt-1)) * time.Second
|
||
g.Log().Infof(ctx, "ChatAPI 重试第%d次(等待%v)", attempt, wait)
|
||
select {
|
||
case <-ctx.Done():
|
||
return nil, ctx.Err()
|
||
case <-time.After(wait):
|
||
}
|
||
}
|
||
|
||
result, doErr := doChatRequest(ctx, url, cfg.APIKey, body, timeout)
|
||
if doErr == nil {
|
||
g.Log().Debugf(ctx, "ChatAPI 调用成功 url=%s tool_calls=%d content_len=%d",
|
||
url, len(result.ToolCalls), len(result.Content))
|
||
return result, nil
|
||
}
|
||
|
||
lastErr = doErr
|
||
g.Log().Warningf(ctx, "ChatAPI request failed (attempt=%d/%d): %v", attempt+1, maxRetries+1, doErr)
|
||
// 只有限流或服务端错误才重试
|
||
errStr := lastErr.Error()
|
||
if !strings.Contains(errStr, "limit_requests") &&
|
||
!strings.Contains(errStr, "limit_tokens") &&
|
||
!strings.Contains(errStr, "500") &&
|
||
!strings.Contains(errStr, "502") &&
|
||
!strings.Contains(errStr, "503") {
|
||
break
|
||
}
|
||
}
|
||
|
||
g.Log().Errorf(ctx, "ChatAPI failed after %d retries: %v", maxRetries+1, lastErr)
|
||
return nil, lastErr
|
||
}
|
||
|
||
func doChatRequest(ctx context.Context, url, apiKey string, body []byte, timeout time.Duration) (*ChatResponse, error) {
|
||
httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(body))
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create request failed: %w", err)
|
||
}
|
||
httpReq.Header.Set("Authorization", "Bearer "+apiKey)
|
||
httpReq.Header.Set("Content-Type", "application/json")
|
||
|
||
start := time.Now()
|
||
client := &http.Client{Timeout: timeout}
|
||
resp, err := client.Do(httpReq)
|
||
elapsed := time.Since(start)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("request failed (elapsed %v): %w", elapsed, err)
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
|
||
respBody, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("read response failed (status=%d): %w", resp.StatusCode, err)
|
||
}
|
||
|
||
if resp.StatusCode != 200 {
|
||
return nil, fmt.Errorf("API response error status=%d body=%s", resp.StatusCode, string(respBody))
|
||
}
|
||
|
||
g.Log().Debugf(ctx, "ChatAPI 响应完成 status=%d body_len=%d elapsed=%v",
|
||
resp.StatusCode, len(respBody), elapsed)
|
||
|
||
return parseRespBody(ctx, respBody)
|
||
}
|
||
|
||
// ==================== 内部实现 ====================
|
||
|
||
type apiReqBody struct {
|
||
Model string `json:"model"`
|
||
Messages []apiMessage `json:"messages"`
|
||
MaxTokens int `json:"max_tokens,omitempty"`
|
||
Temperature float32 `json:"temperature,omitempty"`
|
||
Stream bool `json:"stream,omitempty"`
|
||
Tools []apiToolDef `json:"tools,omitempty"`
|
||
}
|
||
|
||
// apiMessage 用于JSON序列化的消息体（适配OpenAI format）
|
||
type apiMessage struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
ToolCalls []apiToolCall `json:"tool_calls,omitempty"`
|
||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||
Name string `json:"name,omitempty"`
|
||
}
|
||
|
||
type apiToolDef struct {
|
||
Type string `json:"type"`
|
||
Function apiToolFunction `json:"function"`
|
||
}
|
||
|
||
type apiToolFunction struct {
|
||
Name string `json:"name"`
|
||
Description string `json:"description"`
|
||
Parameters map[string]any `json:"parameters"`
|
||
}
|
||
|
||
type apiRespBody struct {
|
||
Choices []apiChoice `json:"choices"`
|
||
Error *struct {
|
||
Message string `json:"message"`
|
||
Code string `json:"code"`
|
||
} `json:"error,omitempty"`
|
||
}
|
||
|
||
type apiChoice struct {
|
||
Index int `json:"index"`
|
||
Message apiRespMsg `json:"message"`
|
||
FinishReason string `json:"finish_reason"`
|
||
}
|
||
|
||
// apiRespMsg 响应消息体（arguments 使用 json.RawMessage 兼容对象和字符串）
|
||
type apiRespMsg struct {
|
||
Content string `json:"content"`
|
||
ToolCalls []apiRespToolCall `json:"tool_calls,omitempty"`
|
||
}
|
||
|
||
type apiRespToolCall struct {
|
||
ID string `json:"id"`
|
||
Type string `json:"type"`
|
||
Function apiRespFuncCall `json:"function"`
|
||
}
|
||
|
||
type apiRespFuncCall struct {
|
||
Name string `json:"name"`
|
||
Arguments json.RawMessage `json:"arguments"`
|
||
}
|
||
|
||
type apiToolCall struct {
|
||
ID string `json:"id"`
|
||
Type string `json:"type"`
|
||
Function apiReqFuncCall `json:"function"`
|
||
}
|
||
|
||
// apiReqFuncCall 请求中的 function call（arguments 为 json.RawMessage 避免二次编码）
|
||
type apiReqFuncCall struct {
|
||
Name string `json:"name"`
|
||
Arguments json.RawMessage `json:"arguments"`
|
||
}
|
||
|
||
func buildReqBody(model string, req *ChatRequest) ([]byte, error) {
|
||
body := apiReqBody{
|
||
Model: model,
|
||
Messages: toAPIMessages(req.Messages),
|
||
MaxTokens: req.MaxTokens,
|
||
Temperature: req.Temperature,
|
||
Stream: req.Stream,
|
||
}
|
||
if len(req.Tools) > 0 {
|
||
body.Tools = make([]apiToolDef, 0, len(req.Tools))
|
||
for _, t := range req.Tools {
|
||
body.Tools = append(body.Tools, apiToolDef{
|
||
Type: "function",
|
||
Function: apiToolFunction{
|
||
Name: t.Name,
|
||
Description: t.Description,
|
||
Parameters: t.Parameters,
|
||
},
|
||
})
|
||
}
|
||
}
|
||
return json.Marshal(body)
|
||
}
|
||
|
||
func toAPIMessages(msgs []*ChatMessage) []apiMessage {
|
||
out := make([]apiMessage, 0, len(msgs))
|
||
for _, m := range msgs {
|
||
om := apiMessage{
|
||
Role: m.Role,
|
||
Content: m.Content,
|
||
ToolCallID: m.ToolCallID,
|
||
Name: m.Name,
|
||
}
|
||
if len(m.ToolCalls) > 0 {
|
||
om.ToolCalls = make([]apiToolCall, 0, len(m.ToolCalls))
|
||
for _, tc := range m.ToolCalls {
|
||
args := tc.Arguments
|
||
if args == "" || !json.Valid([]byte(args)) {
|
||
args = "{}"
|
||
}
|
||
om.ToolCalls = append(om.ToolCalls, apiToolCall{
|
||
ID: tc.ID,
|
||
Type: "function",
|
||
Function: apiReqFuncCall{
|
||
Name: tc.Name,
|
||
Arguments: json.RawMessage(args),
|
||
},
|
||
})
|
||
}
|
||
}
|
||
out = append(out, om)
|
||
}
|
||
return out
|
||
}
|
||
|
||
func parseRespBody(ctx context.Context, data []byte) (*ChatResponse, error) {
|
||
var resp apiRespBody
|
||
if err := json.Unmarshal(data, &resp); err != nil {
|
||
return nil, fmt.Errorf("parse response failed: %s", string(data))
|
||
}
|
||
if resp.Error != nil {
|
||
return nil, fmt.Errorf("API error(code=%s): %s", resp.Error.Code, resp.Error.Message)
|
||
}
|
||
if len(resp.Choices) == 0 {
|
||
return nil, fmt.Errorf("empty response")
|
||
}
|
||
|
||
msg := resp.Choices[0].Message
|
||
cr := &ChatResponse{Content: msg.Content}
|
||
|
||
// 检测 finish_reason 是否为 length（被 max_tokens 截断）
|
||
if resp.Choices[0].FinishReason == "length" {
|
||
g.Log().Warningf(ctx, "ChatAPI response truncated (finish_reason=length), content_len=%d, consider increasing max_tokens", len(msg.Content))
|
||
}
|
||
|
||
if len(msg.ToolCalls) > 0 {
|
||
cr.ToolCalls = make([]*ToolCall, 0, len(msg.ToolCalls))
|
||
for _, tc := range msg.ToolCalls {
|
||
args := resolveArguments(tc.Function.Arguments)
|
||
cr.ToolCalls = append(cr.ToolCalls, &ToolCall{
|
||
ID: tc.ID,
|
||
Name: tc.Function.Name,
|
||
Arguments: args,
|
||
})
|
||
}
|
||
}
|
||
return cr, nil
|
||
}
|
||
|
||
// resolveArguments 将 json.RawMessage 的参数转为字符串
|
||
// API 可能返回 "arguments": "{\"key\":\"val\"}"（字符串）或 "arguments": {"key":"val"}（对象）
|
||
func resolveArguments(raw json.RawMessage) string {
|
||
if len(raw) == 0 {
|
||
return ""
|
||
}
|
||
// 如果是 JSON 字符串（以 " 开头），直接提取字符串值
|
||
if raw[0] == '"' {
|
||
var s string
|
||
if json.Unmarshal(raw, &s) == nil {
|
||
return s
|
||
}
|
||
}
|
||
// 否则是 JSON 对象，重新序列化回字符串
|
||
return string(raw)
|
||
}
|
||
|
||
func trimSlashes(s string) string {
|
||
for len(s) > 0 && s[len(s)-1] == '/' {
|
||
s = s[:len(s)-1]
|
||
}
|
||
return s
|
||
}
|