831 lines
24 KiB
Go
831 lines
24 KiB
Go
package routers
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"ops/agents"
|
|
"ops/models"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// SSE frame types sent to frontend
|
|
type sseEvent struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
Tool string `json:"tool,omitempty"`
|
|
Stage string `json:"stage,omitempty"`
|
|
Status string `json:"status,omitempty"`
|
|
Message string `json:"message,omitempty"`
|
|
Data map[string]interface{} `json:"data,omitempty"`
|
|
Stats *tokenUsageStats `json:"stats,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
type tokenUsageStats struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
ToolPromptTokens int `json:"tool_prompt_tokens"`
|
|
ToolCompletionTokens int `json:"tool_completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
CompletionTokensPerSec float64 `json:"completion_tokens_per_sec"`
|
|
PeakCompletionTokensPerSec float64 `json:"peak_completion_tokens_per_sec"`
|
|
Estimated bool `json:"estimated"`
|
|
}
|
|
|
|
const maxImageDataSize = 4 * 1024 * 1024
|
|
|
|
var allowedImageTypes = map[string]bool{
|
|
"image/jpeg": true,
|
|
"image/png": true,
|
|
"image/webp": true,
|
|
"image/gif": true,
|
|
}
|
|
|
|
// chatRequestFromFrontend is the expected POST body
|
|
type chatRequest struct {
|
|
Messages []chatMessage `json:"messages"`
|
|
OpenAIName string `json:"openaiName,omitempty"`
|
|
}
|
|
|
|
type chatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
ImageURL string `json:"image_url,omitempty"`
|
|
ImageURLAlias string `json:"imageURL,omitempty"`
|
|
}
|
|
|
|
// openaiChatRequest is the request sent to the upstream OpenAI-compatible API
|
|
type openaiChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []openaiMessage `json:"messages"`
|
|
Stream bool `json:"stream"`
|
|
MaxTokens int `json:"max_tokens,omitempty"`
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
}
|
|
|
|
type openaiMessage struct {
|
|
Role string `json:"role"`
|
|
Content any `json:"content"`
|
|
}
|
|
|
|
type openaiContentPart struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *openaiImageURL `json:"image_url,omitempty"`
|
|
}
|
|
|
|
type openaiImageURL struct {
|
|
URL string `json:"url"`
|
|
Detail string `json:"detail,omitempty"`
|
|
}
|
|
|
|
type openaiResponseMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// openaiStreamChunk is one SSE data line from the upstream
|
|
type openaiStreamChunk struct {
|
|
ID string `json:"id,omitempty"`
|
|
Object string `json:"object,omitempty"`
|
|
Created int64 `json:"created,omitempty"`
|
|
Model string `json:"model,omitempty"`
|
|
Choices []openaiChoice `json:"choices"`
|
|
Usage *openaiUsage `json:"usage,omitempty"`
|
|
}
|
|
|
|
type openaiChatResponse struct {
|
|
Choices []openaiResponseChoice `json:"choices"`
|
|
Usage *openaiUsage `json:"usage,omitempty"`
|
|
}
|
|
|
|
type openaiResponseChoice struct {
|
|
Message openaiResponseMessage `json:"message"`
|
|
}
|
|
|
|
type toolSelection struct {
|
|
Name string `json:"name"`
|
|
Reason string `json:"reason"`
|
|
}
|
|
|
|
type toolRoutingDecision struct {
|
|
Tools []toolSelection `json:"tools"`
|
|
Reason string `json:"reason"`
|
|
}
|
|
|
|
type toolRoutingResult struct {
|
|
Decision toolRoutingDecision
|
|
Selected []string
|
|
Messages []openaiMessage
|
|
Response string
|
|
Usage *openaiUsage
|
|
}
|
|
|
|
type openaiChoice struct {
|
|
Index int `json:"index"`
|
|
Delta openaiDelta `json:"delta"`
|
|
Finish *string `json:"finish_reason,omitempty"`
|
|
}
|
|
|
|
type openaiDelta struct {
|
|
Role string `json:"role,omitempty"`
|
|
Content string `json:"content,omitempty"`
|
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
|
Reasoning string `json:"reasoning,omitempty"`
|
|
Thinking string `json:"thinking,omitempty"`
|
|
}
|
|
|
|
type openaiUsage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
}
|
|
|
|
func ApiAIChat(r *gin.RouterGroup) {
|
|
r.GET("/openai", handleOpenAIProfiles)
|
|
r.POST("/chat", handleChat)
|
|
|
|
admin := r.Group("/admin")
|
|
admin.POST("/config", handleAIChatAdminGetConfig)
|
|
admin.POST("/config/update", handleAIChatAdminUpdateConfig)
|
|
admin.POST("/refresh", handleAIChatAdminRefreshCache)
|
|
}
|
|
|
|
func handleOpenAIProfiles(ctx *gin.Context) {
|
|
cfg := getAIChatConfig()
|
|
active := ""
|
|
profiles := make([]map[string]interface{}, 0, len(cfg.OpenAI))
|
|
for _, profile := range cfg.OpenAI {
|
|
if profile.Active {
|
|
active = profile.Name
|
|
}
|
|
profiles = append(profiles, map[string]interface{}{
|
|
"name": profile.Name,
|
|
"active": profile.Active,
|
|
"baseUrl": profile.BaseUrl,
|
|
"model": profile.Model,
|
|
"timeout": profile.Timeout,
|
|
"maxTokens": profile.MaxTokens,
|
|
})
|
|
}
|
|
ReturnJson(ctx, "apiOK", gin.H{
|
|
"enabled": cfg.Enabled,
|
|
"active": active,
|
|
"profiles": profiles,
|
|
"toolRouter": gin.H{
|
|
"enabled": cfg.ToolRouter.Enabled,
|
|
"openaiName": cfg.ToolRouter.OpenAIName,
|
|
"timeout": cfg.ToolRouter.Timeout,
|
|
"maxTokens": cfg.ToolRouter.MaxTokens,
|
|
},
|
|
})
|
|
}
|
|
|
|
func handleChat(ctx *gin.Context) {
|
|
data, _ := SeparateData(ctx)
|
|
|
|
if data == nil {
|
|
sendSSEError(ctx, "请求数据为空")
|
|
return
|
|
}
|
|
|
|
var req chatRequest
|
|
if err := decodeJSON(data, &req); err != nil {
|
|
sendSSEError(ctx, "解析消息失败: "+err.Error())
|
|
return
|
|
}
|
|
|
|
if len(req.Messages) == 0 {
|
|
sendSSEError(ctx, "消息不能为空")
|
|
return
|
|
}
|
|
|
|
// Check ai config
|
|
cfg := getAIChatConfig()
|
|
profile, ok := selectOpenAIProfile(cfg, req.OpenAIName)
|
|
if !cfg.Enabled || !ok || profile.Model == "" || profile.ApiKey == "" {
|
|
sendSSEError(ctx, "AI 聊天未配置，请在后台配置 API Key 和模型")
|
|
return
|
|
}
|
|
toolRouterProfile, hasToolRouterProfile := selectOpenAIProfile(cfg, cfg.ToolRouter.OpenAIName)
|
|
chatMsgs := convertToChatMessages(req.Messages)
|
|
|
|
// Set up SSE headers before routing/tools so progress can stream immediately.
|
|
ctx.Writer.Header().Set("Content-Type", "text/event-stream")
|
|
ctx.Writer.Header().Set("Cache-Control", "no-cache")
|
|
ctx.Writer.Header().Set("Connection", "keep-alive")
|
|
ctx.Writer.Header().Set("X-Accel-Buffering", "no")
|
|
ctx.Writer.WriteHeader(http.StatusOK)
|
|
flusher, _ := ctx.Writer.(http.Flusher)
|
|
tracker := newTokenUsageTracker()
|
|
|
|
emitTrace := func(tool, stage, status, message string, data map[string]interface{}) {
|
|
sendSSE(ctx, flusher, sseEvent{Type: "trace", Tool: tool, Stage: stage, Status: status, Message: message, Data: data})
|
|
}
|
|
emitStats := func(stats tokenUsageStats) {
|
|
sendSSE(ctx, flusher, sseEvent{Type: "stats", Stats: &stats})
|
|
}
|
|
|
|
toolConfigs := []agents.ToolConfig{}
|
|
if cfg.ToolRouter.Enabled {
|
|
toolConfigs = buildToolConfigs(cfg.ToolRouter.Tools)
|
|
if hasToolRouterProfile && toolRouterProfile.Model != "" && toolRouterProfile.ApiKey != "" {
|
|
emitTrace("tool_router", "route", "running", "正在进行工具路由", nil)
|
|
routeResult, routeErr := routeTools(ctx.Request.Context(), toolRouterProfile, cfg.ToolRouter, chatMsgs)
|
|
if routeErr != nil {
|
|
emitTrace("tool_router", "route", "error", "工具路由失败，将继续普通回答", map[string]interface{}{"error": routeErr.Error()})
|
|
toolConfigs = []agents.ToolConfig{}
|
|
} else if routeResult != nil {
|
|
tracker.addToolUsage(routeResult.Usage, estimateOpenAIMessagesTokens(routeResult.Messages), estimateTokenCount(routeResult.Response))
|
|
data := map[string]interface{}{
|
|
"tools": routeResult.Selected,
|
|
"selections": routeResult.Decision.Tools,
|
|
"reason": routeResult.Decision.Reason,
|
|
}
|
|
message := "工具路由结果：无需调用工具"
|
|
if len(routeResult.Selected) > 0 {
|
|
message = "工具路由结果：将调用 " + strings.Join(routeResult.Selected, ", ")
|
|
}
|
|
emitTrace("tool_router", "route", "success", message, data)
|
|
toolConfigs = filterToolConfigs(toolConfigs, routeResult.Selected)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Enrich messages with tools (pre-process)
|
|
chatMsgs = agents.EnrichMessages(ctx.Request.Context(), chatMsgs, toolConfigs, emitTrace)
|
|
|
|
// Build OpenAI-compatible request
|
|
openaiMsgs, err := convertToOpenAIMessages(chatMsgs)
|
|
if err != nil {
|
|
sendSSE(ctx, flusher, sseEvent{Type: "error", Error: err.Error()})
|
|
sendSSEDone(ctx, flusher)
|
|
return
|
|
}
|
|
apiReq := openaiChatRequest{
|
|
Model: profile.Model,
|
|
Messages: openaiMsgs,
|
|
Stream: true,
|
|
MaxTokens: profile.MaxTokens,
|
|
Temperature: 0.7,
|
|
}
|
|
|
|
// Add system prompt if configured
|
|
if profile.SystemPrompt != "" {
|
|
apiReq.Messages = append([]openaiMessage{{Role: "system", Content: profile.SystemPrompt}}, apiReq.Messages...)
|
|
}
|
|
|
|
modelPromptTokens := estimateOpenAIMessagesTokens(apiReq.Messages)
|
|
completionTokens := 0
|
|
modelUsageReceived := false
|
|
streamStarted := time.Now()
|
|
windowStarted := streamStarted
|
|
windowTokens := 0
|
|
peakTokensPerSecond := 0.0
|
|
emitTrace("model", "stream", "running", "正在请求模型回复", nil)
|
|
|
|
err = streamOpenAI(ctx.Request.Context(), profile, apiReq, func(chunk openaiStreamChunk) {
|
|
for _, choice := range chunk.Choices {
|
|
reasoningText := choice.Delta.ReasoningContent
|
|
if reasoningText == "" {
|
|
reasoningText = choice.Delta.Reasoning
|
|
}
|
|
if reasoningText == "" {
|
|
reasoningText = choice.Delta.Thinking
|
|
}
|
|
if reasoningText != "" {
|
|
sendSSE(ctx, flusher, sseEvent{Type: "reasoning", Text: reasoningText})
|
|
}
|
|
|
|
if choice.Delta.Content != "" {
|
|
deltaTokens := estimateTokenCount(choice.Delta.Content)
|
|
completionTokens += deltaTokens
|
|
windowTokens += deltaTokens
|
|
elapsedWindow := time.Since(windowStarted).Seconds()
|
|
if elapsedWindow >= 1 {
|
|
peakTokensPerSecond = maxFloat(peakTokensPerSecond, float64(windowTokens)/elapsedWindow)
|
|
windowStarted = time.Now()
|
|
windowTokens = 0
|
|
} else if peakTokensPerSecond == 0 && elapsedWindow > 0.25 {
|
|
peakTokensPerSecond = maxFloat(peakTokensPerSecond, float64(windowTokens)/elapsedWindow)
|
|
}
|
|
stats := tracker.setModelEstimate(modelPromptTokens, completionTokens).snapshot(tokensPerSecond(completionTokens, streamStarted), peakTokensPerSecond)
|
|
sendSSE(ctx, flusher, sseEvent{Type: "delta", Text: choice.Delta.Content, Stats: &stats})
|
|
}
|
|
}
|
|
if chunk.Usage != nil {
|
|
modelUsageReceived = true
|
|
stats := tracker.setModelUsage(chunk.Usage).snapshot(tokensPerSecond(tracker.completionTokens, streamStarted), peakTokensPerSecond)
|
|
emitStats(stats)
|
|
}
|
|
})
|
|
if err != nil {
|
|
sendSSE(ctx, flusher, sseEvent{Type: "error", Error: "请求失败: " + err.Error()})
|
|
sendSSEDone(ctx, flusher)
|
|
return
|
|
}
|
|
|
|
if windowTokens > 0 {
|
|
elapsedWindow := time.Since(windowStarted).Seconds()
|
|
if elapsedWindow > 0 {
|
|
peakTokensPerSecond = maxFloat(peakTokensPerSecond, float64(windowTokens)/elapsedWindow)
|
|
}
|
|
}
|
|
emitTrace("model", "stream", "success", "模型回复完成", nil)
|
|
if modelUsageReceived {
|
|
emitStats(tracker.snapshot(tokensPerSecond(tracker.completionTokens, streamStarted), peakTokensPerSecond))
|
|
} else {
|
|
emitStats(tracker.setModelEstimate(modelPromptTokens, completionTokens).snapshot(tokensPerSecond(completionTokens, streamStarted), peakTokensPerSecond))
|
|
}
|
|
sendSSEDone(ctx, flusher)
|
|
flusher.Flush()
|
|
}
|
|
|
|
func streamOpenAI(ctx context.Context, cfg models.ConfigsAIChatOpenAI_, req openaiChatRequest, onData func(openaiStreamChunk)) error {
|
|
bodyBytes, err := json.Marshal(req)
|
|
if err != nil {
|
|
return fmt.Errorf("序列化请求失败: %w", err)
|
|
}
|
|
|
|
url := strings.TrimRight(cfg.BaseUrl, "/") + "/chat/completions"
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyBytes))
|
|
if err != nil {
|
|
return fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Authorization", "Bearer "+cfg.ApiKey)
|
|
httpReq.Header.Set("Accept", "text/event-stream")
|
|
|
|
client := &http.Client{Timeout: time.Duration(cfg.Timeout) * time.Second}
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return fmt.Errorf("连接上游服务失败: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return fmt.Errorf("上游返回 %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
scanner.Buffer(make([]byte, 0, 64*1024), 256*1024)
|
|
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
|
|
if strings.TrimSpace(line) == "" {
|
|
continue
|
|
}
|
|
if !strings.HasPrefix(line, "data: ") {
|
|
continue
|
|
}
|
|
|
|
payload := strings.TrimPrefix(line, "data: ")
|
|
payload = strings.TrimSpace(payload)
|
|
|
|
if payload == "[DONE]" {
|
|
continue
|
|
}
|
|
|
|
var chunk openaiStreamChunk
|
|
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
|
continue
|
|
}
|
|
|
|
onData(chunk)
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
return fmt.Errorf("读取流失败: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Initialize with system prompt if present
|
|
func sendSSE(ctx *gin.Context, flusher http.Flusher, event sseEvent) {
|
|
data, err := json.Marshal(event)
|
|
if err != nil {
|
|
return
|
|
}
|
|
_, _ = fmt.Fprintf(ctx.Writer, "data: %s\n\n", string(data))
|
|
flusher.Flush()
|
|
}
|
|
|
|
func sendSSEDone(ctx *gin.Context, flusher http.Flusher) {
|
|
_, _ = fmt.Fprint(ctx.Writer, "data: [DONE]\n\n")
|
|
flusher.Flush()
|
|
}
|
|
|
|
func sendSSEError(ctx *gin.Context, message string) {
|
|
ctx.Writer.Header().Set("Content-Type", "text/event-stream")
|
|
ctx.Writer.Header().Set("Cache-Control", "no-cache")
|
|
ctx.Writer.Header().Set("Connection", "keep-alive")
|
|
ctx.Writer.WriteHeader(http.StatusOK)
|
|
flusher, _ := ctx.Writer.(http.Flusher)
|
|
|
|
sendSSE(ctx, flusher, sseEvent{
|
|
Type: "error",
|
|
Error: message,
|
|
})
|
|
|
|
sendSSEDone(ctx, flusher)
|
|
flusher.Flush()
|
|
}
|
|
|
|
func convertToChatMessages(msgs []chatMessage) []agents.ChatMessage {
|
|
result := make([]agents.ChatMessage, 0, len(msgs))
|
|
for _, m := range msgs {
|
|
imageURL := m.ImageURL
|
|
if imageURL == "" {
|
|
imageURL = m.ImageURLAlias
|
|
}
|
|
result = append(result, agents.ChatMessage{Role: m.Role, Content: m.Content, ImageURL: imageURL})
|
|
}
|
|
return result
|
|
}
|
|
|
|
func convertToOpenAIMessages(msgs []agents.ChatMessage) ([]openaiMessage, error) {
|
|
result := make([]openaiMessage, 0, len(msgs))
|
|
for _, m := range msgs {
|
|
content, err := buildOpenAIContent(m)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
result = append(result, openaiMessage{Role: m.Role, Content: content})
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func buildOpenAIContent(m agents.ChatMessage) (any, error) {
|
|
if strings.TrimSpace(m.ImageURL) == "" {
|
|
return m.Content, nil
|
|
}
|
|
|
|
imageURL, err := normalizeImageURL(m.ImageURL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
parts := []openaiContentPart{
|
|
{
|
|
Type: "image_url",
|
|
ImageURL: &openaiImageURL{
|
|
URL: imageURL,
|
|
Detail: "auto",
|
|
},
|
|
},
|
|
}
|
|
if m.Content != "" {
|
|
parts = append(parts, openaiContentPart{Type: "text", Text: m.Content})
|
|
}
|
|
return parts, nil
|
|
}
|
|
|
|
func normalizeImageURL(raw string) (string, error) {
|
|
value := strings.TrimSpace(raw)
|
|
if value == "" {
|
|
return "", errors.New("图片地址为空")
|
|
}
|
|
|
|
if strings.HasPrefix(strings.ToLower(value), "data:") {
|
|
return normalizeImageDataURI(value)
|
|
}
|
|
|
|
parsed, err := url.Parse(value)
|
|
if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
|
return "", errors.New("图片地址无效，仅支持 http/https URL 或 base64 data URI")
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func normalizeImageDataURI(raw string) (string, error) {
|
|
commaIndex := strings.Index(raw, ",")
|
|
if commaIndex == -1 {
|
|
return "", errors.New("图片 data URI 格式无效")
|
|
}
|
|
|
|
metadata := strings.TrimSpace(raw[len("data:"):commaIndex])
|
|
payload := strings.TrimSpace(raw[commaIndex+1:])
|
|
if metadata == "" || payload == "" {
|
|
return "", errors.New("图片 data URI 格式无效")
|
|
}
|
|
|
|
metadataParts := strings.Split(metadata, ";")
|
|
mimeType := strings.ToLower(strings.TrimSpace(metadataParts[0]))
|
|
if !allowedImageTypes[mimeType] {
|
|
return "", errors.New("图片格式不支持，仅支持 jpeg/png/webp/gif")
|
|
}
|
|
|
|
hasBase64 := false
|
|
for _, part := range metadataParts[1:] {
|
|
if strings.EqualFold(strings.TrimSpace(part), "base64") {
|
|
hasBase64 = true
|
|
break
|
|
}
|
|
}
|
|
if !hasBase64 {
|
|
return "", errors.New("图片 data URI 必须使用 base64 编码")
|
|
}
|
|
|
|
if len(payload) > maxImageDataSize*4/3+16 {
|
|
return "", errors.New("图片过大，请选择小于 4MB 的图片")
|
|
}
|
|
decoded, err := base64.StdEncoding.DecodeString(payload)
|
|
if err != nil {
|
|
return "", errors.New("图片 base64 数据无效")
|
|
}
|
|
if len(decoded) > maxImageDataSize {
|
|
return "", errors.New("图片过大，请选择小于 4MB 的图片")
|
|
}
|
|
|
|
return "data:" + mimeType + ";base64," + payload, nil
|
|
}
|
|
|
|
type tokenUsageTracker struct {
|
|
promptTokens int
|
|
completionTokens int
|
|
toolPromptTokens int
|
|
toolCompletionTokens int
|
|
estimated bool
|
|
}
|
|
|
|
func newTokenUsageTracker() *tokenUsageTracker {
|
|
return &tokenUsageTracker{estimated: true}
|
|
}
|
|
|
|
func (t *tokenUsageTracker) addToolUsage(usage *openaiUsage, estimatedPromptTokens, estimatedCompletionTokens int) {
|
|
if usage != nil {
|
|
t.toolPromptTokens += usage.PromptTokens
|
|
t.toolCompletionTokens += usage.CompletionTokens
|
|
return
|
|
}
|
|
t.toolPromptTokens += estimatedPromptTokens
|
|
t.toolCompletionTokens += estimatedCompletionTokens
|
|
t.estimated = true
|
|
}
|
|
|
|
func (t *tokenUsageTracker) setModelEstimate(promptTokens, completionTokens int) *tokenUsageTracker {
|
|
t.promptTokens = promptTokens
|
|
t.completionTokens = completionTokens
|
|
t.estimated = true
|
|
return t
|
|
}
|
|
|
|
func (t *tokenUsageTracker) setModelUsage(usage *openaiUsage) *tokenUsageTracker {
|
|
if usage == nil {
|
|
return t
|
|
}
|
|
t.promptTokens = usage.PromptTokens
|
|
t.completionTokens = usage.CompletionTokens
|
|
return t
|
|
}
|
|
|
|
func (t *tokenUsageTracker) snapshot(completionTokensPerSec, peakCompletionTokensPerSec float64) tokenUsageStats {
|
|
totalTokens := t.promptTokens + t.completionTokens + t.toolPromptTokens + t.toolCompletionTokens
|
|
return tokenUsageStats{
|
|
PromptTokens: t.promptTokens,
|
|
CompletionTokens: t.completionTokens,
|
|
ToolPromptTokens: t.toolPromptTokens,
|
|
ToolCompletionTokens: t.toolCompletionTokens,
|
|
TotalTokens: totalTokens,
|
|
CompletionTokensPerSec: completionTokensPerSec,
|
|
PeakCompletionTokensPerSec: peakCompletionTokensPerSec,
|
|
Estimated: t.estimated,
|
|
}
|
|
}
|
|
|
|
func estimateOpenAIMessagesTokens(messages []openaiMessage) int {
|
|
total := 0
|
|
for _, message := range messages {
|
|
total += estimateTokenCount(message.Role) + 4
|
|
total += estimateOpenAIContentTokens(message.Content)
|
|
}
|
|
return total
|
|
}
|
|
|
|
func estimateOpenAIContentTokens(content any) int {
|
|
switch value := content.(type) {
|
|
case string:
|
|
return estimateTokenCount(value)
|
|
case []openaiContentPart:
|
|
total := 0
|
|
for _, part := range value {
|
|
switch part.Type {
|
|
case "text":
|
|
total += estimateTokenCount(part.Text)
|
|
case "image_url":
|
|
total += 85
|
|
}
|
|
}
|
|
return total
|
|
case []interface{}:
|
|
data, err := json.Marshal(value)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return estimateTokenCount(string(data))
|
|
default:
|
|
data, err := json.Marshal(value)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return estimateTokenCount(string(data))
|
|
}
|
|
}
|
|
|
|
func estimateTokenCount(text string) int {
|
|
text = strings.TrimSpace(text)
|
|
if text == "" {
|
|
return 0
|
|
}
|
|
|
|
tokens := 0
|
|
asciiRunes := 0
|
|
flushASCII := func() {
|
|
if asciiRunes > 0 {
|
|
tokens += (asciiRunes + 3) / 4
|
|
asciiRunes = 0
|
|
}
|
|
}
|
|
|
|
for _, r := range text {
|
|
if r <= 127 {
|
|
if r == ' ' || r == '\n' || r == '\t' || r == '\r' {
|
|
flushASCII()
|
|
continue
|
|
}
|
|
asciiRunes++
|
|
continue
|
|
}
|
|
flushASCII()
|
|
tokens++
|
|
}
|
|
flushASCII()
|
|
if tokens == 0 {
|
|
return 1
|
|
}
|
|
return tokens
|
|
}
|
|
|
|
func tokensPerSecond(tokens int, start time.Time) float64 {
|
|
elapsed := time.Since(start).Seconds()
|
|
if tokens <= 0 || elapsed <= 0 {
|
|
return 0
|
|
}
|
|
return float64(tokens) / elapsed
|
|
}
|
|
|
|
func maxFloat(a, b float64) float64 {
|
|
if b > a {
|
|
return b
|
|
}
|
|
return a
|
|
}
|
|
|
|
func buildToolConfigs(configs []models.ConfigsAIChatTool_) []agents.ToolConfig {
|
|
result := make([]agents.ToolConfig, 0, len(configs))
|
|
for _, c := range configs {
|
|
result = append(result, agents.ToolConfig{
|
|
Name: c.Name,
|
|
Enabled: c.Enabled,
|
|
Description: c.Description,
|
|
})
|
|
}
|
|
return result
|
|
}
|
|
|
|
func selectOpenAIProfile(cfg models.ConfigsAIChat_, name string) (models.ConfigsAIChatOpenAI_, bool) {
|
|
if name != "" {
|
|
for _, p := range cfg.OpenAI {
|
|
if p.Name == name {
|
|
return p, true
|
|
}
|
|
}
|
|
return models.ConfigsAIChatOpenAI_{}, false
|
|
}
|
|
for _, p := range cfg.OpenAI {
|
|
if p.Active {
|
|
return p, true
|
|
}
|
|
}
|
|
if len(cfg.OpenAI) > 0 {
|
|
return cfg.OpenAI[0], true
|
|
}
|
|
return models.ConfigsAIChatOpenAI_{}, false
|
|
}
|
|
|
|
func routeTools(ctx context.Context, profile models.ConfigsAIChatOpenAI_, router models.ConfigsAIChatToolRouter_, messages []agents.ChatMessage) (*toolRoutingResult, error) {
|
|
lastUserContent := strings.TrimSpace(agents.LastUserContent(messages))
|
|
if lastUserContent == "" {
|
|
return nil, nil
|
|
}
|
|
|
|
openaiMsgs := []openaiMessage{{Role: "user", Content: lastUserContent}}
|
|
|
|
toolNames := make([]string, 0, len(router.Tools))
|
|
for _, t := range router.Tools {
|
|
if t.Enabled {
|
|
toolNames = append(toolNames, t.Name+" - "+t.Description)
|
|
}
|
|
}
|
|
if len(toolNames) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
sysPrompt := "请根据用户的最新一条消息，判断需要启用哪些工具来完成用户需求。\n可选工具：\n" + strings.Join(toolNames, "\n") + "\n\n回复格式要求：\n```json\n{\"tools\":[{\"name\":\"工具名称\",\"reason\":\"选择原因\"}],\"reason\":\"整体判断理由\"}\n```\n仅输出 JSON 代码块。如果没有需要启用的工具，返回 {\"tools\":[]}。"
|
|
openaiMsgs = append([]openaiMessage{{Role: "system", Content: sysPrompt}}, openaiMsgs...)
|
|
|
|
req := openaiChatRequest{
|
|
Model: profile.Model,
|
|
Messages: openaiMsgs,
|
|
Stream: false,
|
|
MaxTokens: router.MaxTokens,
|
|
Temperature: 0.1,
|
|
}
|
|
|
|
bodyBytes, err := json.Marshal(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
url := strings.TrimRight(profile.BaseUrl, "/") + "/chat/completions"
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyBytes))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Authorization", "Bearer "+profile.ApiKey)
|
|
client := &http.Client{Timeout: time.Duration(router.Timeout) * time.Second}
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("工具路由返回 %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var result openaiChatResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(result.Choices) == 0 {
|
|
return &toolRoutingResult{Messages: openaiMsgs, Usage: result.Usage}, nil
|
|
}
|
|
|
|
response := result.Choices[0].Message.Content
|
|
decision := extractToolRoutingDecision(response)
|
|
selected := make([]string, 0, len(decision.Tools))
|
|
for _, t := range decision.Tools {
|
|
name := strings.TrimSpace(t.Name)
|
|
if name != "" {
|
|
selected = append(selected, name)
|
|
}
|
|
}
|
|
return &toolRoutingResult{Decision: decision, Selected: selected, Messages: openaiMsgs, Response: response, Usage: result.Usage}, nil
|
|
}
|
|
|
|
func extractToolRoutingDecision(response string) toolRoutingDecision {
|
|
start := strings.Index(response, "{")
|
|
end := strings.LastIndex(response, "}")
|
|
if start == -1 || end == -1 || end <= start {
|
|
return toolRoutingDecision{}
|
|
}
|
|
var parsed toolRoutingDecision
|
|
if err := json.Unmarshal([]byte(response[start:end+1]), &parsed); err != nil {
|
|
return toolRoutingDecision{}
|
|
}
|
|
return parsed
|
|
}
|
|
|
|
func filterToolConfigs(configs []agents.ToolConfig, selected []string) []agents.ToolConfig {
|
|
if len(selected) == 0 {
|
|
return []agents.ToolConfig{}
|
|
}
|
|
selectedSet := make(map[string]bool, len(selected))
|
|
for _, s := range selected {
|
|
selectedSet[s] = true
|
|
}
|
|
filtered := make([]agents.ToolConfig, 0, len(configs))
|
|
for _, c := range configs {
|
|
if selectedSet[c.Name] {
|
|
filtered = append(filtered, c)
|
|
}
|
|
}
|
|
return filtered
|
|
}
|