MuyueWorkspace/internal/api/handlers_chat.go

package api

import (
	"context"
	"encoding/json"
	"net/http"
	"regexp"
	"strings"

	"github.com/muyue/muyue/internal/agent"
	"github.com/muyue/muyue/internal/orchestrator"
)

var thinkingTagRegex = regexp.MustCompile(`(?s)<[Tt]hink[^>]*>.*?</[Tt]hink>`)

const maxToolIterations = 15

func (s *Server) handleChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}
	var body struct {
		Message string `json:"message"`
		Stream  bool   `json:"stream"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}
	if body.Message == "" {
		writeError(w, "no message", http.StatusMethodNotAllowed)
		return
	}

	s.convStore.Add("user", body.Message)

	if s.convStore.NeedsSummarization() {
		s.autoSummarize()
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		writeError(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	orb.SetSystemPrompt(agent.StudioSystemPrompt())
	orb.SetTools(s.agentToolsJSON)

	if body.Stream {
		s.handleStreamChat(w, orb, body.Message)
	} else {
		s.handleNonStreamChat(w, orb, body.Message)
	}
}

func (s *Server) handleStreamChat(w http.ResponseWriter, orb *orchestrator.Orchestrator, userMessage string) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(http.StatusOK)
	flusher, canFlush := w.(http.Flusher)

	writeSSE := func(data map[string]interface{}) {
		b, _ := json.Marshal(data)
		w.Write([]byte("data: " + string(b) + "\n\n"))
		if canFlush {
			flusher.Flush()
		}
	}

	ctx := context.Background()
	messages := s.buildContextMessages(userMessage)

	var finalContent string
	var allToolCalls []map[string]interface{}
	var allToolResults []map[string]interface{}

	for i := 0; i < maxToolIterations; i++ {
		resp, err := orb.SendWithTools(messages)
		if err != nil {
			writeSSE(map[string]interface{}{"error": err.Error()})
			return
		}

		choice := resp.Choices[0]
		content := cleanThinkingTags(choice.Message.Content)

		if content != "" {
			words := strings.Fields(content)
			for i, w := range words {
				chunk := w
				if i < len(words)-1 {
					chunk += " "
				}
				writeSSE(map[string]interface{}{"content": chunk})
			}
			finalContent = content
		}

		if len(choice.Message.ToolCalls) == 0 {
			break
		}

		assistantMsg := orchestrator.Message{
			Role:      "assistant",
			Content:   content,
			ToolCalls: choice.Message.ToolCalls,
		}
		messages = append(messages, assistantMsg)

		for _, tc := range choice.Message.ToolCalls {
			toolCallData := map[string]interface{}{
				"tool_call_id": tc.ID,
				"name":         tc.Function.Name,
				"args":         tc.Function.Arguments,
			}
			allToolCalls = append(allToolCalls, toolCallData)
			writeSSE(map[string]interface{}{"tool_call": toolCallData})

			call := agent.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: json.RawMessage(tc.Function.Arguments),
			}

			result, execErr := s.agentRegistry.Execute(ctx, call)
			if execErr != nil {
				result = agent.ToolResponse{
					Content: execErr.Error(),
					IsError: true,
				}
			}

			resultData := map[string]interface{}{
				"tool_call_id": tc.ID,
				"content":      result.Content,
				"is_error":     result.IsError,
			}
			writeSSE(map[string]interface{}{"tool_result": resultData})

			allToolResults = append(allToolResults, map[string]interface{}{
				"tool_call_id": tc.ID,
				"name":         tc.Function.Name,
				"args":         tc.Function.Arguments,
				"result":       result.Content,
				"is_error":     result.IsError,
			})

			messages = append(messages, orchestrator.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: tc.ID,
				Name:       tc.Function.Name,
			})
		}

		finalContent = ""
	}

	storeContent := finalContent
	if len(allToolCalls) > 0 {
		storeObj := map[string]interface{}{
			"content":      storeContent,
			"tool_calls":   allToolCalls,
			"tool_results": allToolResults,
		}
		storeJSON, _ := json.Marshal(storeObj)
		storeContent = string(storeJSON)
	}
	s.convStore.Add("assistant", storeContent)

	writeSSE(map[string]interface{}{"done": "true"})
}

func (s *Server) handleNonStreamChat(w http.ResponseWriter, orb *orchestrator.Orchestrator, userMessage string) {
	ctx := context.Background()
	messages := s.buildContextMessages(userMessage)

	var finalContent string

	for i := 0; i < maxToolIterations; i++ {
		resp, err := orb.SendWithTools(messages)
		if err != nil {
			writeError(w, err.Error(), http.StatusInternalServerError)
			return
		}

		choice := resp.Choices[0]
		content := cleanThinkingTags(choice.Message.Content)

		if content != "" {
			finalContent = content
		}

		if len(choice.Message.ToolCalls) == 0 {
			break
		}

		assistantMsg := orchestrator.Message{
			Role:      "assistant",
			Content:   content,
			ToolCalls: choice.Message.ToolCalls,
		}
		messages = append(messages, assistantMsg)

		for _, tc := range choice.Message.ToolCalls {
			call := agent.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: json.RawMessage(tc.Function.Arguments),
			}

			result, execErr := s.agentRegistry.Execute(ctx, call)
			if execErr != nil {
				result = agent.ToolResponse{
					Content: execErr.Error(),
					IsError: true,
				}
			}

			messages = append(messages, orchestrator.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: tc.ID,
				Name:       tc.Function.Name,
			})
		}

		finalContent = ""
	}

	if finalContent == "" {
		finalContent = "(tool calls completed, no text response)"
	}

	s.convStore.Add("assistant", finalContent)
	writeJSON(w, map[string]string{"content": finalContent})
}

func cleanThinkingTags(content string) string {
	return strings.TrimSpace(thinkingTagRegex.ReplaceAllString(content, ""))
}

const contextWindowMessages = 20

func (s *Server) buildContextMessages(userMessage string) []orchestrator.Message {
	history := s.convStore.Get()
	start := 0
	if len(history) > contextWindowMessages {
		start = len(history) - contextWindowMessages
	}

	messages := make([]orchestrator.Message, 0, len(history[start:])+1)

	summary := s.convStore.GetSummary()
	if summary != "" {
		messages = append(messages, orchestrator.Message{
			Role:    "system",
			Content: "Résumé de la conversation précédente:\n" + summary,
		})
	}

	for _, m := range history[start:] {
		content := m.Content
		if m.Role == "assistant" {
			var parsed struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					ToolCallID string `json:"tool_call_id"`
					Name       string `json:"name"`
					Args       string `json:"args"`
				} `json:"tool_calls"`
			}
			if err := json.Unmarshal([]byte(content), &parsed); err == nil && parsed.Content != "" {
				content = parsed.Content
			}
		}
		role := m.Role
		if role == "system" {
			continue
		}
		messages = append(messages, orchestrator.Message{
			Role:    role,
			Content: content,
		})
	}

	messages = append(messages, orchestrator.Message{
		Role:    "user",
		Content: userMessage,
	})

	return messages
}

func (s *Server) autoSummarize() {
	messages := s.convStore.Get()
	if len(messages) < 10 {
		return
	}

	half := len(messages) / 2
	var oldText string
	for _, m := range messages[:half] {
		oldText += m.Role + ": " + m.Content + "\n\n"
	}

	summary := s.convStore.GetSummary()
	if summary != "" {
		oldText = "Résumé précédent:\n" + summary + "\n\nNouveaux échanges:\n" + oldText
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		return
	}
	orb.SetSystemPrompt(summarizePrompt)

	result, err := orb.Send(oldText)
	if err != nil {
		return
	}

	s.convStore.SetSummary(result)
	s.convStore.TrimOld(len(messages) - half)
	s.convStore.Add("system", "[Conversation résumée automatiquement]")
}

func (s *Server) handleChatHistory(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeError(w, "GET only", http.StatusMethodNotAllowed)
		return
	}
	messages := s.convStore.Get()
	writeJSON(w, map[string]interface{}{
		"messages": messages,
		"tokens":   s.convStore.ApproxTokenCount(),
	})
}

func (s *Server) handleChatClear(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}
	s.convStore.Clear()
	writeJSON(w, map[string]string{"status": "ok"})
}