MuyueWorkspace/internal/api/handlers_chat.go

package api

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

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

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 := []orchestrator.Message{
		{Role: "user", Content: userMessage},
	}

	var finalContent string
	var allToolCalls []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 != "" {
			for _, ch := range strings.Split(content, "") {
				writeSSE(map[string]interface{}{"content": ch})
			}
			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})

			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}
		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 := []orchestrator.Message{
		{Role: "user", Content: 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.ReplaceAll(content, "<think", "")
}

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