拆分 main 功能模块

Co-Authored-By: Claude <noreply@anthropic.com>
2026-06-17 11:29:51 +08:00
parent 132ab2a1cb
commit ccc1260fe0
21 changed files with 2318 additions and 2074 deletions
@@ -0,0 +1,298 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"aichat/conversation"
+	"aichat/llm"
+	"aichat/message"
+	"aichat/stream"
+	"aichat/toolrouter"
+	"aichat/utils"
+
+	"github.com/gin-gonic/gin"
+	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
+)
+
+type activeProfileRequest struct {
+	Name string `json:"name"`
+}
+
+func (s *Server) indexHandler(c *gin.Context) {
+	profile := s.aiState.ActiveProfile()
+	c.HTML(http.StatusOK, "chat.html", gin.H{
+		"Title":      "AI 对话",
+		"Model":      profile.Config.Model,
+		"OpenAIName": profile.Config.Name,
+	})
+}
+
+func (s *Server) listOpenAIHandler(c *gin.Context) {
+	c.JSON(http.StatusOK, s.aiState.ListProfiles())
+}
+
+func (s *Server) switchOpenAIHandler(c *gin.Context) {
+	var req activeProfileRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
+		return
+	}
+	profile, err := s.aiState.SwitchActive(req.Name)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"active":  profile.Config.Name,
+		"profile": llm.PublicConfig(profile, true),
+	})
+}
+
+func (s *Server) listSearchHandler(c *gin.Context) {
+	c.JSON(http.StatusOK, s.searchState.ListProfiles())
+}
+
+func (s *Server) switchSearchHandler(c *gin.Context) {
+	var req activeProfileRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
+		return
+	}
+	profile, err := s.searchState.SwitchActive(req.Name)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+	profile.APIKey = ""
+	profile.Active = true
+	c.JSON(http.StatusOK, gin.H{
+		"active":  profile.Name,
+		"profile": profile,
+	})
+}
+
+func (s *Server) listConversationsHandler(c *gin.Context) {
+	convs, err := s.store.List()
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+	c.JSON(http.StatusOK, convs)
+}
+
+func (s *Server) createConversationHandler(c *gin.Context) {
+	conv, err := s.store.Create()
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "创建对话失败: " + err.Error()})
+		return
+	}
+	c.JSON(http.StatusOK, conv)
+}
+
+func (s *Server) getConversationHandler(c *gin.Context) {
+	conv, err := s.store.Get(c.Param("id"))
+	if err != nil {
+		status := http.StatusInternalServerError
+		if err.Error() == "对话不存在" {
+			status = http.StatusNotFound
+		}
+		c.JSON(status, gin.H{"error": err.Error()})
+		return
+	}
+	c.JSON(http.StatusOK, conv)
+}
+
+func (s *Server) deleteConversationHandler(c *gin.Context) {
+	if err := s.store.Delete(c.Param("id")); err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+	c.Status(http.StatusNoContent)
+}
+
+// chatHandler 流式 SSE 对话接口
+func (s *Server) chatHandler(c *gin.Context) {
+	var req message.ChatRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
+		return
+	}
+	if len(req.Messages) == 0 {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "消息不能为空"})
+		return
+	}
+	profile, err := s.aiState.GetProfile(req.OpenAIName)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	// SSE 头先写出，后续插件/模型过程都通过 trace 事件实时展示。
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+	c.Writer.WriteHeader(http.StatusOK)
+	flusher, ok := c.Writer.(http.Flusher)
+	if !ok {
+		return
+	}
+	emit := func(frame stream.Frame) {
+		stream.WriteSSEJSON(c.Writer, frame)
+		flusher.Flush()
+	}
+	emitTrace := func(tool, stage, status, message string, data map[string]any) {
+		emit(stream.Frame{Type: "trace", Tool: tool, Stage: stage, Status: status, Message: message, Data: data})
+	}
+	emitError := func(err error) {
+		emit(stream.Frame{Type: "error", Error: err.Error()})
+	}
+
+	// 超时 context
+	timeout := time.Duration(profile.Config.Timeout) * time.Second
+	ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
+	defer cancel()
+	usage := stream.NewTracker()
+	ctx = stream.ContextWithTracker(ctx, usage)
+
+	// 用 Function Calling 工具循环替代旧的路由+隐藏上下文机制
+	messages, err := toolrouter.RunAgentToolLoop(ctx, s.toolRouterState, profile, req.Messages, s.searchState, s.sqlState, emit)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "Agent 工具循环失败:", err)
+		messages, err = message.BuildArkMessages(req.Messages)
+		if err != nil {
+			emitError(err)
+			return
+		}
+	}
+	promptTokens := stream.EstimateChatMessagesTokens(req.Messages)
+
+	if llm.IsOllamaProfile(profile) && message.HasImageMessage(req.Messages) {
+		emitTrace("model", "request", "running", "正在通过 Ollama 原生接口调用视觉模型", nil)
+		err = stream.StreamOllamaChat(ctx, profile, messages, promptTokens, usage, emit, func(content string) {
+			if req.ConversationID != "" {
+				if err := conversation.SaveMessages(s.store, req.ConversationID, req.Messages, content); err != nil {
+					fmt.Fprintln(os.Stderr, "保存对话失败:", err)
+				}
+			}
+		})
+		if err != nil {
+			emitError(err)
+		}
+		return
+	}
+
+	emitTrace("model", "request", "running", "正在调用模型生成回答", nil)
+	modelStream, err := profile.Client.CreateChatCompletionStream(ctx, model.CreateChatCompletionRequest{
+		Model:     profile.Config.Model,
+		Messages:  messages,
+		MaxTokens: utils.IntPtr(4096),
+	}.WithStream(true))
+	if err != nil {
+		emitError(err)
+		return
+	}
+	defer modelStream.Close()
+	emitTrace("model", "stream", "running", "模型已开始输出", nil)
+
+	var full strings.Builder
+	completionTokens := 0
+	streamStarted := time.Now()
+	windowStarted := streamStarted
+	windowTokens := 0
+	peakTokensPerSecond := 0.0
+	parseThinkTags := llm.ShouldParseThinkTags(profile)
+	thinkParser := &stream.Parser{}
+	emitDelta := func(delta string) {
+		if delta == "" {
+			return
+		}
+		now := time.Now()
+		deltaTokens := stream.EstimateTokenCount(delta)
+		windowTokens += deltaTokens
+		windowElapsed := now.Sub(windowStarted).Seconds()
+		if windowElapsed >= 1 {
+			windowSpeed := float64(windowTokens) / windowElapsed
+			if windowSpeed > peakTokensPerSecond {
+				peakTokensPerSecond = windowSpeed
+			}
+			windowStarted = now
+			windowTokens = 0
+		} else if peakTokensPerSecond == 0 && windowElapsed > 0.25 {
+			peakTokensPerSecond = float64(windowTokens) / windowElapsed
+		}
+		full.WriteString(delta)
+		completionTokens += deltaTokens
+		usage.SetModel(promptTokens, completionTokens)
+		stats := usage.Snapshot(stream.TokensPerSecond(completionTokens, streamStarted), peakTokensPerSecond)
+		emit(stream.Frame{Type: "delta", Text: delta, Stats: &stats})
+	}
+	emitModelContent := func(delta string) {
+		if delta == "" {
+			return
+		}
+		if !parseThinkTags {
+			emitDelta(delta)
+			return
+		}
+		visible, reasoning := thinkParser.Accept(delta)
+		if reasoning != "" {
+			emit(stream.Frame{Type: "reasoning", Text: reasoning})
+		}
+		emitDelta(visible)
+	}
+	for {
+		resp, err := modelStream.Recv()
+		if errors.Is(err, io.EOF) {
+			if parseThinkTags {
+				visible, reasoning := thinkParser.Flush()
+				if reasoning != "" {
+					emit(stream.Frame{Type: "reasoning", Text: reasoning})
+				}
+				emitDelta(visible)
+			}
+			usage.SetModel(promptTokens, completionTokens)
+			if windowTokens > 0 {
+				windowElapsed := time.Since(windowStarted).Seconds()
+				if windowElapsed > 0.25 {
+					windowSpeed := float64(windowTokens) / windowElapsed
+					if windowSpeed > peakTokensPerSecond {
+						peakTokensPerSecond = windowSpeed
+					}
+				}
+			}
+			if peakTokensPerSecond == 0 {
+				peakTokensPerSecond = stream.TokensPerSecond(completionTokens, streamStarted)
+			}
+			if req.ConversationID != "" {
+				if err := conversation.SaveMessages(s.store, req.ConversationID, req.Messages, full.String()); err != nil {
+					fmt.Fprintln(os.Stderr, "保存对话失败:", err)
+				}
+			}
+			finalStats := usage.Snapshot(stream.TokensPerSecond(completionTokens, streamStarted), peakTokensPerSecond)
+			emit(stream.Frame{Type: "stats", Stats: &finalStats})
+			emitTrace("model", "stream", "success", "回答生成完成", nil)
+			fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
+			flusher.Flush()
+			return
+		}
+		if err != nil {
+			emitError(err)
+			return
+		}
+		if len(resp.Choices) > 0 {
+			emitModelContent(resp.Choices[0].Delta.Content)
+			// 思考过程 reasoning_content 单独事件推送
+			if resp.Choices[0].Delta.ReasoningContent != nil && *resp.Choices[0].Delta.ReasoningContent != "" {
+				emit(stream.Frame{Type: "reasoning", Text: *resp.Choices[0].Delta.ReasoningContent})
+			}
+		}
+	}
+}