Changeset 0.7.2 (#40)

server/handlers/completion.go

@@ -4,14 +4,15 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log"
 	"net/http"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 
 	"git.gobha.me/xcaliber/chat-switchboard/database"
 	"git.gobha.me/xcaliber/chat-switchboard/providers"
+	"git.gobha.me/xcaliber/chat-switchboard/tools"
 )
 
 // ── Request Types ───────────────────────────
@@ -45,9 +46,11 @@ func NewCompletionHandler() *CompletionHandler {
 //   2. Resolve api_config (from request, chat, or user default)
 //   3. Load conversation history from messages
 //   4. Persist user message
-//   5. Call provider (stream or non-stream)
-//   6. Stream SSE to client / return JSON
-//   7. Persist assistant message with token counts
+//   5. Attach tool definitions if model supports tools
+//   6. Call provider (stream or non-stream)
+//   7. Tool loop: if LLM requests tool calls, execute them and re-call
+//   8. Stream SSE to client / return JSON
+//   9. Persist assistant message with token counts
 
 func (h *CompletionHandler) Complete(c *gin.Context) {
 	var req completionRequest
@@ -153,6 +156,11 @@ func (h *CompletionHandler) Complete(c *gin.Context) {
 		provReq.TopP = req.TopP
 	}
 
+	// Attach tool definitions if model supports tool calling and tools are available
+	if caps.ToolCalling && tools.HasTools() {
+		provReq.Tools = h.buildToolDefs()
+	}
+
 	// Determine streaming
 	stream := true // default
 	if req.Stream != nil {
@@ -166,7 +174,26 @@ func (h *CompletionHandler) Complete(c *gin.Context) {
 	}
 }
 
-// ── Streaming Completion (SSE) ──────────────
+// buildToolDefs converts registered tools to provider-format tool definitions.
+func (h *CompletionHandler) buildToolDefs() []providers.ToolDef {
+	allTools := tools.AllDefinitions()
+	defs := make([]providers.ToolDef, len(allTools))
+	for i, t := range allTools {
+		defs[i] = providers.ToolDef{
+			Type: "function",
+			Function: providers.FunctionDef{
+				Name:        t.Name,
+				Description: t.Description,
+				Parameters:  t.Parameters,
+			},
+		}
+	}
+	return defs
+}
+
+// ── Streaming Completion (SSE) with Tool Loop ──
+
+const maxToolIterations = 10 // safety limit on tool call rounds
 
 func (h *CompletionHandler) streamCompletion(
 	c *gin.Context,
@@ -175,12 +202,6 @@ func (h *CompletionHandler) streamCompletion(
 	req providers.CompletionRequest,
 	channelID, userID, model string,
 ) {
-	ch, err := provider.StreamCompletion(c.Request.Context(), cfg, req)
-	if err != nil {
-		c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
-		return
-	}
-
 	// Set SSE headers
 	c.Header("Content-Type", "text/event-stream")
 	c.Header("Cache-Control", "no-cache")
@@ -188,45 +209,150 @@ func (h *CompletionHandler) streamCompletion(
 	c.Header("X-Accel-Buffering", "no") // Disable nginx buffering
 
 	c.Status(http.StatusOK)
-
-	var fullContent string
 	flusher, _ := c.Writer.(http.Flusher)
 
-	for event := range ch {
-		if event.Error != nil {
-			fmt.Fprintf(c.Writer, "data: {\"error\":\"%s\"}\n\n", event.Error.Error())
-			if flusher != nil {
-				flusher.Flush()
-			}
-			break
-		}
-
-		if event.Delta != "" {
-			fullContent += event.Delta
-			// OpenAI-compatible SSE format
-			fmt.Fprintf(c.Writer, "data: {\"choices\":[{\"delta\":{\"content\":%q},\"finish_reason\":null}],\"model\":%q}\n\n",
-				event.Delta, model)
-			if flusher != nil {
-				flusher.Flush()
-			}
-		}
-
-		if event.Done {
-			finishReason := event.FinishReason
-			if finishReason == "" {
-				finishReason = "stop"
-			}
-			fmt.Fprintf(c.Writer, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":%q}],\"model\":%q}\n\n",
-				finishReason, model)
-			io.WriteString(c.Writer, "data: [DONE]\n\n")
-			if flusher != nil {
-				flusher.Flush()
-			}
-			break
+	flush := func() {
+		if flusher != nil {
+			flusher.Flush()
 		}
 	}
 
-	// Persist assistant response
+	// sendSSE sends a standard OpenAI-compatible SSE data line.
+	sendSSE := func(data string) {
+		fmt.Fprintf(c.Writer, "data: %s\n\n", data)
+		flush()
+	}
+
+	// sendEvent sends a named SSE event (for tool status — non-standard).
+	sendEvent := func(event, data string) {
+		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, data)
+		flush()
+	}
+
+	var fullContent string
+
+	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
+		}
+
+		var iterContent string
+		var toolCalls []providers.ToolCall
+
+		for event := range ch {
+			if event.Error != nil {
+				sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(event.Error.Error())))
+				return
+			}
+
+			// 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"
+					}
+					fullContent += iterContent
+					sendSSE(fmt.Sprintf(
+						`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
+						finishReason, model))
+					sendSSE("[DONE]")
+
+					// Persist assistant response
+					if fullContent != "" {
+						if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
+							log.Printf("Failed to persist assistant message: %v", err)
+						}
+					}
+					return
+				}
+				break
+			}
+		}
+
+		// ── Tool execution round ────────────────
+		if len(toolCalls) == 0 {
+			// Stream ended without tool calls or finish — shouldn't happen
+			sendSSE("[DONE]")
+			if iterContent != "" {
+				fullContent += iterContent
+				if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
+					log.Printf("Failed to persist assistant message: %v", err)
+				}
+			}
+			return
+		}
+
+		// 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)
+			result := tools.ExecuteCall(c.Request.Context(), execCtx, call)
+
+			// Notify client about tool result
+			resultJSON, _ := json.Marshal(map[string]interface{}{
+				"tool_call_id": result.ToolCallID,
+				"name":         result.Name,
+				"content":      result.Content,
+				"is_error":     result.IsError,
+			})
+			sendEvent("tool_result", string(resultJSON))
+
+			// Append tool result to conversation for next iteration
+			req.Messages = append(req.Messages, providers.Message{
+				Role:       "tool",
+				Content:    result.Content,
+				ToolCallID: result.ToolCallID,
+				Name:       result.Name,
+			})
+		}
+
+		fullContent += iterContent
+		// Loop: send updated messages back to provider
+	}
+
+	// If we hit max iterations, send what we have
+	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]")
+
 	if fullContent != "" {
 		if _, err := h.persistMessage(channelID, userID, "assistant", fullContent, model, 0, 0, nil); err != nil {
 			log.Printf("Failed to persist assistant message: %v", err)
@@ -234,7 +360,7 @@ func (h *CompletionHandler) streamCompletion(
 	}
 }
 
-// ── Non-Streaming Completion ────────────────
+// ── Non-Streaming Completion with Tool Loop ──
 
 func (h *CompletionHandler) syncCompletion(
 	c *gin.Context,
@@ -243,39 +369,108 @@ func (h *CompletionHandler) syncCompletion(
 	req providers.CompletionRequest,
 	channelID, userID, model string,
 ) {
-	resp, err := provider.ChatCompletion(c.Request.Context(), cfg, req)
-	if err != nil {
-		c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
-		return
+	var totalInput, totalOutput int
+	var finalContent string
+
+	for iteration := 0; iteration < maxToolIterations; iteration++ {
+		resp, err := provider.ChatCompletion(c.Request.Context(), cfg, req)
+		if err != nil {
+			c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
+			return
+		}
+
+		totalInput += resp.InputTokens
+		totalOutput += resp.OutputTokens
+
+		// No tool calls — normal response
+		if len(resp.ToolCalls) == 0 || resp.FinishReason != "tool_calls" {
+			finalContent = resp.Content
+
+			// Persist assistant response
+			if _, err := h.persistMessage(channelID, userID, "assistant", finalContent, model, totalInput, totalOutput, nil); err != nil {
+				log.Printf("Failed to persist assistant message: %v", err)
+			}
+
+			// Return OpenAI-compatible response
+			c.JSON(http.StatusOK, gin.H{
+				"model": resp.Model,
+				"choices": []gin.H{
+					{
+						"message": gin.H{
+							"role":    "assistant",
+							"content": finalContent,
+						},
+						"finish_reason": resp.FinishReason,
+					},
+				},
+				"usage": gin.H{
+					"prompt_tokens":     totalInput,
+					"completion_tokens": totalOutput,
+					"total_tokens":      totalInput + totalOutput,
+				},
+			})
+			return
+		}
+
+		// ── Tool execution round ────────────────
+		assistantMsg := providers.Message{
+			Role:    "assistant",
+			Content: resp.Content,
+		}
+		for _, tc := range resp.ToolCalls {
+			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
+		}
+		req.Messages = append(req.Messages, assistantMsg)
+
+		execCtx := tools.ExecutionContext{
+			UserID:    userID,
+			ChannelID: channelID,
+		}
+		for _, tc := range resp.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)
+			result := tools.ExecuteCall(c.Request.Context(), execCtx, call)
+
+			req.Messages = append(req.Messages, providers.Message{
+				Role:       "tool",
+				Content:    result.Content,
+				ToolCallID: result.ToolCallID,
+				Name:       result.Name,
+			})
+		}
 	}
 
-	// Persist assistant response
-	if _, err := h.persistMessage(channelID, userID, "assistant", resp.Content, model, resp.InputTokens, resp.OutputTokens, nil); err != nil {
-		log.Printf("Failed to persist assistant message: %v", err)
-	}
-
-	// Return OpenAI-compatible response
+	// Max iterations reached
 	c.JSON(http.StatusOK, gin.H{
-		"model": resp.Model,
+		"model": model,
 		"choices": []gin.H{
 			{
 				"message": gin.H{
 					"role":    "assistant",
-					"content": resp.Content,
+					"content": "[Tool execution limit reached]",
 				},
-				"finish_reason": resp.FinishReason,
+				"finish_reason": "stop",
 			},
 		},
-		"usage": gin.H{
-			"prompt_tokens":     resp.InputTokens,
-			"completion_tokens": resp.OutputTokens,
-			"total_tokens":      resp.InputTokens + resp.OutputTokens,
-		},
 	})
 }
 
 // ── Model Capabilities ──────────────────────
 
+// escapeJSON escapes a string for safe embedding in a JSON string value.
+func escapeJSON(s string) string {
+	s = strings.ReplaceAll(s, `\`, `\\`)
+	s = strings.ReplaceAll(s, `"`, `\"`)
+	s = strings.ReplaceAll(s, "\n", `\n`)
+	s = strings.ReplaceAll(s, "\r", `\r`)
+	s = strings.ReplaceAll(s, "\t", `\t`)
+	return s
+}
+
 // getModelCapabilities looks up capabilities from model_configs DB,
 // then overlays with known model defaults and heuristic detection.
 func (h *CompletionHandler) getModelCapabilities(model, apiConfigID string) providers.ModelCapabilities {