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

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) {
	SetupSSEHeaders(w)
	flusher, canFlush := w.(http.Flusher)


	sseWriter := NewSSEWriter(w)


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

	engine := NewChatEngine(orb, s.agentRegistry, s.agentToolsJSON)
	engine.OnChunk(func(data map[string]interface{}) {
		if data == nil {
			return
		}
		sseWriter.Write(data)
		if canFlush {
			flusher.Flush()
		}
	})

	finalContent, allToolCalls, allToolResults, err := engine.RunWithTools(ctx, messages)
	if err != nil {
		sseWriter.Write(map[string]interface{}{"error": err.Error()})
		return
	}

	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)

	sseWriter.Write(map[string]interface{}{"done": "true"})
}

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

	engine := NewChatEngine(orb, s.agentRegistry, s.agentToolsJSON)
	finalContent, err := engine.RunNonStream(ctx, messages)
	if err != nil {
		writeError(w, err.Error(), http.StatusInternalServerError)
		return
	}

	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"})
}