aichat/server/handlers.go

package server

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	searchagent "aichat/agents/search"
	"aichat/contextwindow"
	"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"`
}

type createConversationRequest struct {
	Preset       string `json:"preset"`
	PresetPrompt string `json:"preset_prompt"`
}

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) searchState() (*searchagent.State, bool) {
	if s == nil || s.toolManager == nil {
		return nil, false
	}
	raw, ok := s.toolManager.RawState(searchagent.ToolName)
	if !ok {
		return nil, false
	}
	state, ok := raw.(*searchagent.State)
	return state, ok && state != nil
}

func (s *Server) listSearchHandler(c *gin.Context) {
	state, ok := s.searchState()
	if !ok {
		c.JSON(http.StatusOK, searchagent.ListResponse{})
		return
	}
	c.JSON(http.StatusOK, state.ListProfiles())
}

func (s *Server) switchSearchHandler(c *gin.Context) {
	state, ok := s.searchState()
	if !ok {
		c.JSON(http.StatusBadRequest, gin.H{"error": "搜索工具未加载"})
		return
	}
	var req activeProfileRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
		return
	}
	profile, err := state.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) {
	var req createConversationRequest
	if c.Request.Body != nil {
		_ = c.ShouldBindJSON(&req)
	}
	preset := req.PresetPrompt
	if strings.TrimSpace(preset) == "" {
		preset = req.Preset
	}
	conv, err := s.store.CreateWithPreset(preset)
	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)
}

func conversationPresetPrompt(store *conversation.Store, id string) string {
	id = strings.TrimSpace(id)
	if id == "" || store == nil {
		return ""
	}
	conv, err := store.Get(id)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(conv.PresetPrompt)
}

func conversationContextMessages(messages []message.ChatMessage, preset string) []message.ChatMessage {
	preset = strings.TrimSpace(preset)
	if preset == "" {
		return append([]message.ChatMessage(nil), messages...)
	}
	cleaned := filterRequestPresetMessages(messages)
	contextMessages := make([]message.ChatMessage, 0, len(cleaned)+1)
	contextMessages = append(contextMessages, message.ChatMessage{Role: "system", Content: preset, Hidden: true})
	contextMessages = append(contextMessages, cleaned...)
	return contextMessages
}

func filterRequestPresetMessages(messages []message.ChatMessage) []message.ChatMessage {
	filtered := make([]message.ChatMessage, 0, len(messages))
	for _, msg := range messages {
		if msg.Hidden && strings.EqualFold(msg.Role, "system") {
			continue
		}
		filtered = append(filtered, msg)
	}
	return filtered
}

// 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)

	contextMessages := conversationContextMessages(req.Messages, conversationPresetPrompt(s.store, req.ConversationID))
	chatWindow := contextwindow.ApplyChatWindow(contextMessages, profile.Config.ContextWindowTokens)
	if chatWindow.Removed > 0 || chatWindow.Overflow {
		emitTrace("context_window", "chat", "success", "已清理对话历史上下文", map[string]any{"max_tokens": chatWindow.MaxTokens, "before_tokens": chatWindow.BeforeTokens, "after_tokens": chatWindow.AfterTokens, "removed_messages": chatWindow.Removed, "overflow": chatWindow.Overflow, "base_overflow": chatWindow.BaseOverflow})
	}
	contextMessages = chatWindow.ChatMessages

	// 用 Function Calling 工具循环替代旧的路由+隐藏上下文机制
	messages, err := toolrouter.RunAgentToolLoop(ctx, s.toolRouterState, profile, contextMessages, s.toolManager, emit)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Agent 工具循环失败:", err)
		messages, err = message.BuildArkMessages(contextMessages)
		if err != nil {
			emitError(err)
			return
		}
	}
	arkWindow := contextwindow.ApplyArkWindow(messages, profile.Config.ContextWindowTokens)
	if arkWindow.Removed > 0 || arkWindow.Overflow {
		emitTrace("context_window", "model", "success", "已清理最终模型上下文", map[string]any{"max_tokens": arkWindow.MaxTokens, "before_tokens": arkWindow.BeforeTokens, "after_tokens": arkWindow.AfterTokens, "removed_messages": arkWindow.Removed, "overflow": arkWindow.Overflow, "base_overflow": arkWindow.BaseOverflow})
	}
	messages = arkWindow.ArkMessages
	promptTokens := arkWindow.AfterTokens

	if llm.IsOllamaProfile(profile) && message.HasImageMessage(contextMessages) {
		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})
			}
		}
	}
}