Changeset 0.9.3 (#53)

server/handlers/stream_loop.go

@@ -0,0 +1,202 @@
+package handlers
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"git.gobha.me/xcaliber/chat-switchboard/providers"
+	"git.gobha.me/xcaliber/chat-switchboard/tools"
+)
+
+// streamResult holds the accumulated output from a streaming completion
+// with tool execution. Callers are responsible for persistence.
+type streamResult struct {
+	Content      string
+	ToolActivity []map[string]interface{}
+}
+
+// streamWithToolLoop is the single, canonical streaming implementation.
+// It handles SSE setup, multi-round tool execution, reasoning_content
+// forwarding, and client notifications. Returns the accumulated content
+// and tool activity for the caller to persist however it needs to.
+//
+// Used by:
+//   - streamCompletion (normal chat — persists via persistMessage)
+//   - Regenerate       (regen/edit — persists as sibling with tree metadata)
+//
+// If you add features here (new event types, new tool capabilities, new
+// persistence fields), both callers get them automatically.
+func streamWithToolLoop(
+	c *gin.Context,
+	provider providers.Provider,
+	cfg providers.ProviderConfig,
+	req *providers.CompletionRequest,
+	model, userID, channelID string,
+) streamResult {
+	// Set SSE headers
+	c.Header("Content-Type", "text/event-stream")
+	c.Header("Cache-Control", "no-cache")
+	c.Header("Connection", "keep-alive")
+	c.Header("X-Accel-Buffering", "no")
+
+	c.Status(http.StatusOK)
+	flusher, _ := c.Writer.(http.Flusher)
+
+	flush := func() {
+		if flusher != nil {
+			flusher.Flush()
+		}
+	}
+
+	sendSSE := func(data string) {
+		fmt.Fprintf(c.Writer, "data: %s\n\n", data)
+		flush()
+	}
+
+	sendEvent := func(event, data string) {
+		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, data)
+		flush()
+	}
+
+	var result streamResult
+
+	for iteration := 0; iteration < maxToolIterations; iteration++ {
+		ch, err := provider.StreamCompletion(c.Request.Context(), cfg, *req)
+		if err != nil {
+			sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
+			return result
+		}
+
+		var iterContent string
+		var iterReasoning string
+		var toolCalls []providers.ToolCall
+
+		for event := range ch {
+			if event.Error != nil {
+				sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(event.Error.Error())))
+				return result
+			}
+
+			// Stream reasoning deltas to client
+			if event.Reasoning != "" {
+				iterReasoning += event.Reasoning
+				sendSSE(fmt.Sprintf(
+					`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q}`,
+					event.Reasoning, model))
+			}
+
+			// Stream text deltas to client
+			if event.Delta != "" {
+				iterContent += event.Delta
+				sendSSE(fmt.Sprintf(
+					`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q}`,
+					event.Delta, model))
+			}
+
+			if event.Done {
+				if event.FinishReason == "tool_calls" && len(event.ToolCalls) > 0 {
+					toolCalls = event.ToolCalls
+				} else {
+					// Normal completion — send finish event
+					finishReason := event.FinishReason
+					if finishReason == "" {
+						finishReason = "stop"
+					}
+					if iterReasoning != "" {
+						result.Content += "<think>" + iterReasoning + "</think>"
+					}
+					result.Content += iterContent
+					sendSSE(fmt.Sprintf(
+						`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
+						finishReason, model))
+					sendSSE("[DONE]")
+					return result
+				}
+				break
+			}
+		}
+
+		// ── Tool execution round ────────────────
+		if len(toolCalls) == 0 {
+			// Stream ended without tool calls or finish — shouldn't happen
+			sendSSE("[DONE]")
+			if iterReasoning != "" {
+				result.Content += "<think>" + iterReasoning + "</think>"
+			}
+			result.Content += iterContent
+			return result
+		}
+
+		// Notify client about tool calls
+		toolCallsJSON, _ := json.Marshal(toolCalls)
+		sendEvent("tool_use", string(toolCallsJSON))
+
+		// Append assistant message (with tool_calls, possibly with text) to conversation
+		assistantMsg := providers.Message{
+			Role:    "assistant",
+			Content: iterContent,
+		}
+		for _, tc := range toolCalls {
+			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
+		}
+		req.Messages = append(req.Messages, assistantMsg)
+
+		// Execute all tool calls
+		execCtx := tools.ExecutionContext{
+			UserID:    userID,
+			ChannelID: channelID,
+		}
+		for _, tc := range toolCalls {
+			call := tools.ToolCall{
+				ID:        tc.ID,
+				Name:      tc.Function.Name,
+				Arguments: tc.Function.Arguments,
+			}
+
+			log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID)
+			toolResult := tools.ExecuteCall(c.Request.Context(), execCtx, call)
+
+			// Notify client about tool result
+			resultJSON, _ := json.Marshal(map[string]interface{}{
+				"tool_call_id": toolResult.ToolCallID,
+				"name":         toolResult.Name,
+				"content":      toolResult.Content,
+				"is_error":     toolResult.IsError,
+			})
+			sendEvent("tool_result", string(resultJSON))
+
+			// Collect for persistence
+			result.ToolActivity = append(result.ToolActivity, map[string]interface{}{
+				"id":        tc.ID,
+				"name":      tc.Function.Name,
+				"arguments": tc.Function.Arguments,
+				"result":    toolResult.Content,
+				"is_error":  toolResult.IsError,
+			})
+
+			// Append tool result to conversation for next iteration
+			req.Messages = append(req.Messages, providers.Message{
+				Role:       "tool",
+				Content:    toolResult.Content,
+				ToolCallID: toolResult.ToolCallID,
+				Name:       toolResult.Name,
+			})
+		}
+
+		result.Content += iterContent
+		// Loop: send updated messages back to provider
+	}
+
+	// Hit max iterations
+	log.Printf("⚠️ Tool loop hit max iterations (%d) for channel %s", maxToolIterations, channelID)
+	sendSSE(fmt.Sprintf(
+		`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q}`,
+		escapeJSON("[Tool execution limit reached]"), model))
+	sendSSE("[DONE]")
+
+	return result
+}