拆分 main 功能模块

Co-Authored-By: Claude <noreply@anthropic.com>

server/handlers.go

@@ -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})
+			}
+		}
+	}
+}

server/routes.go

@@ -0,0 +1,14 @@
+package server
+
+func (s *Server) registerRoutes() {
+	s.router.GET("/", s.indexHandler)
+	s.router.POST("/api/chat", s.chatHandler)
+	s.router.GET("/api/openai", s.listOpenAIHandler)
+	s.router.POST("/api/openai/active", s.switchOpenAIHandler)
+	s.router.GET("/api/search", s.listSearchHandler)
+	s.router.POST("/api/search/active", s.switchSearchHandler)
+	s.router.GET("/api/conversations", s.listConversationsHandler)
+	s.router.POST("/api/conversations", s.createConversationHandler)
+	s.router.GET("/api/conversations/:id", s.getConversationHandler)
+	s.router.DELETE("/api/conversations/:id", s.deleteConversationHandler)
+}

server/server.go

@@ -0,0 +1,65 @@
+package server
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+	"os"
+
+	searchagent "aichat/agents/search"
+	sqlquery "aichat/agents/sql"
+	"aichat/config"
+	"aichat/conversation"
+	"aichat/llm"
+	"aichat/toolrouter"
+
+	"github.com/gin-gonic/gin"
+)
+
+type Server struct {
+	cfg             *config.Config
+	aiState         *llm.State
+	searchState     *searchagent.State
+	sqlState        *sqlquery.State
+	toolRouterState *toolrouter.State
+	store           *conversation.Store
+	router          *gin.Engine
+}
+
+func New(cfg *config.Config, aiState *llm.State, searchState *searchagent.State, sqlState *sqlquery.State, toolRouterState *toolrouter.State, store *conversation.Store) *Server {
+	s := &Server{
+		cfg:             cfg,
+		aiState:         aiState,
+		searchState:     searchState,
+		sqlState:        sqlState,
+		toolRouterState: toolRouterState,
+		store:           store,
+		router:          gin.Default(),
+	}
+	s.router.LoadHTMLGlob("templates/*")
+	s.router.Static("/static", "./static")
+	s.registerRoutes()
+	return s
+}
+
+func (s *Server) Run() error {
+	switch s.cfg.Server.Mode {
+	case "unix":
+		return s.runUnix(s.cfg.Server.Address)
+	default:
+		fmt.Println("服务已启动，监听 TCP:", s.cfg.Server.Address)
+		return s.router.Run(s.cfg.Server.Address)
+	}
+}
+
+func (s *Server) runUnix(socketPath string) error {
+	if _, statErr := os.Stat(socketPath); statErr == nil {
+		os.Remove(socketPath)
+	}
+	ln, err := net.Listen("unix", socketPath)
+	if err != nil {
+		return fmt.Errorf("监听 Unix socket 失败: %w", err)
+	}
+	fmt.Println("服务已启动，监听 Unix socket:", socketPath)
+	return http.Serve(ln, s.router)
+}