Co-authored-by: Jeffrey Smith <jasafpro@gmail.com> Co-committed-by: Jeffrey Smith <jasafpro@gmail.com>
751 lines
23 KiB
Go
751 lines
23 KiB
Go
// Package handlers — tool_loop.go
|
|
//
|
|
// v0.27.2: Extracted core tool loop with sink abstraction.
|
|
//
|
|
// coreToolLoop is the single implementation of multi-round LLM completion
|
|
// with tool execution. All callers (SSE streaming, multi-model streaming,
|
|
// sync JSON, headless scheduler) provide a LoopSink for I/O and a LoopConfig
|
|
// for dependencies. The loop handles provider calls, tool dispatch (server +
|
|
// browser bridge), health recording, budget enforcement, and workspace events.
|
|
//
|
|
// Prior to this extraction, the same logic was duplicated in three places:
|
|
// - streamWithToolLoop (SSE streaming)
|
|
// - streamModelResponse (multi-model SSE)
|
|
// - syncCompletion (non-streaming JSON)
|
|
//
|
|
// Those functions are now thin wrappers over coreToolLoop.
|
|
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"go.starlark.net/starlark"
|
|
|
|
"git.gobha.me/xcaliber/chat-switchboard/events"
|
|
"git.gobha.me/xcaliber/chat-switchboard/providers"
|
|
"git.gobha.me/xcaliber/chat-switchboard/sandbox"
|
|
"git.gobha.me/xcaliber/chat-switchboard/store"
|
|
"git.gobha.me/xcaliber/chat-switchboard/tools"
|
|
)
|
|
|
|
// ── Types ──────────────────────────────────────
|
|
|
|
// defaultMaxRounds is the safety limit on tool call iterations.
|
|
// Callers can override via LoopBudget.MaxRounds.
|
|
const defaultMaxRounds = 10
|
|
|
|
// LoopResult holds the accumulated output from a completion with tool
|
|
// execution. Callers are responsible for persistence.
|
|
type LoopResult struct {
|
|
Content string
|
|
ToolActivity []map[string]interface{}
|
|
InputTokens int
|
|
OutputTokens int
|
|
CacheCreationTokens int
|
|
CacheReadTokens int
|
|
ToolCallCount int // total tool calls across all rounds
|
|
BudgetExceeded string // "" | "tokens" | "tool_calls" | "max_rounds"
|
|
Error error // non-nil if the loop terminated due to a provider error
|
|
}
|
|
|
|
// LoopBudget controls resource limits for the tool loop.
|
|
// Zero values mean "use default" for MaxRounds and "unlimited" for the rest.
|
|
type LoopBudget struct {
|
|
MaxRounds int // max tool-call iterations; 0 = defaultMaxRounds
|
|
MaxToolCalls int // total tool calls across all rounds; 0 = unlimited
|
|
MaxTokens int // total input+output tokens; 0 = unlimited
|
|
}
|
|
|
|
func (b LoopBudget) maxRounds() int {
|
|
if b.MaxRounds > 0 {
|
|
return b.MaxRounds
|
|
}
|
|
return defaultMaxRounds
|
|
}
|
|
|
|
// LoopConfig bundles all dependencies for the core tool loop.
|
|
type LoopConfig struct {
|
|
Provider providers.Provider
|
|
Cfg providers.ProviderConfig
|
|
Req *providers.CompletionRequest
|
|
Model string
|
|
ProviderType string
|
|
ExecCtx tools.ExecutionContext
|
|
Hub *events.Hub // nil for headless (scheduler)
|
|
Health HealthRecorder // nil for headless
|
|
ConfigID string
|
|
Budget LoopBudget
|
|
Streaming bool // true = StreamCompletion, false = ChatCompletion
|
|
// v0.29.2: extension tool dispatch (nil = no extension tools)
|
|
Runner *sandbox.Runner
|
|
ExtTools map[string]*store.PackageRegistration // toolName → package
|
|
}
|
|
|
|
// ── Sink Interface ─────────────────────────────
|
|
|
|
// LoopSink receives events from the core tool loop for I/O.
|
|
// Implementations control what happens with each event: write SSE to
|
|
// a gin.Context, accumulate silently, log to stdout, etc.
|
|
type LoopSink interface {
|
|
// OnDelta receives a text content delta from the LLM.
|
|
OnDelta(delta string)
|
|
|
|
// OnReasoning receives a reasoning/thinking delta from the LLM.
|
|
OnReasoning(delta string)
|
|
|
|
// OnError receives a fatal error (provider failure, stream error).
|
|
OnError(err error)
|
|
|
|
// OnToolUse is called when the LLM requests tool calls.
|
|
OnToolUse(calls []providers.ToolCall)
|
|
|
|
// OnToolResult is called after each tool execution completes.
|
|
OnToolResult(result map[string]interface{})
|
|
|
|
// OnFinish is called when the LLM produces a finish_reason.
|
|
OnFinish(reason string)
|
|
|
|
// OnDone signals end of the entire completion (after finish or budget breach).
|
|
OnDone()
|
|
|
|
// OnMaxRounds is called when the loop hits its iteration limit.
|
|
OnMaxRounds()
|
|
}
|
|
|
|
// ── Core Tool Loop ─────────────────────────────
|
|
|
|
// coreToolLoop is the single, canonical tool-loop implementation.
|
|
// It drives multi-round LLM completion with tool execution, budget
|
|
// enforcement, and health recording. I/O is delegated to the sink.
|
|
//
|
|
// The caller is responsible for:
|
|
// - Setting SSE headers (if streaming to a client)
|
|
// - Calling providers.GetHooks().PreRequest() before entry
|
|
// - Persisting the result after return
|
|
// - Sending [DONE] or HTTP response after return (via sink.OnDone)
|
|
func CoreToolLoop(ctx context.Context, lcfg LoopConfig, sink LoopSink) LoopResult {
|
|
var result LoopResult
|
|
maxRounds := lcfg.Budget.maxRounds()
|
|
|
|
for iteration := 0; iteration < maxRounds; iteration++ {
|
|
var iterContent string
|
|
var iterReasoning string
|
|
var toolCalls []providers.ToolCall
|
|
var iterInput, iterOutput, iterCacheCreate, iterCacheRead int
|
|
|
|
if lcfg.Streaming {
|
|
done := runStreamingRound(ctx, lcfg, sink, &result,
|
|
&iterContent, &iterReasoning, &toolCalls,
|
|
&iterInput, &iterOutput, &iterCacheCreate, &iterCacheRead)
|
|
if done {
|
|
return result
|
|
}
|
|
} else {
|
|
done := runSyncRound(ctx, lcfg, sink, &result,
|
|
&iterContent, &iterReasoning, &toolCalls,
|
|
&iterInput, &iterOutput, &iterCacheCreate, &iterCacheRead)
|
|
if done {
|
|
return result
|
|
}
|
|
}
|
|
|
|
// ── Tool execution round ────────────────
|
|
if len(toolCalls) == 0 {
|
|
// Stream/call ended without tool calls or finish — edge case
|
|
if iterReasoning != "" {
|
|
result.Content += "<think>" + iterReasoning + "</think>"
|
|
}
|
|
result.Content += iterContent
|
|
sink.OnDone()
|
|
return result
|
|
}
|
|
|
|
// Budget check: tool calls (pre-execution)
|
|
result.ToolCallCount += len(toolCalls)
|
|
if lcfg.Budget.MaxToolCalls > 0 && result.ToolCallCount > lcfg.Budget.MaxToolCalls {
|
|
result.BudgetExceeded = "tool_calls"
|
|
result.Content += iterContent
|
|
sink.OnFinish("budget_exceeded")
|
|
sink.OnDone()
|
|
return result
|
|
}
|
|
|
|
sink.OnToolUse(toolCalls)
|
|
|
|
// 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)
|
|
}
|
|
lcfg.Req.Messages = append(lcfg.Req.Messages, assistantMsg)
|
|
|
|
// Execute all tool calls
|
|
for _, tc := range toolCalls {
|
|
call := tools.ToolCall{
|
|
ID: tc.ID,
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
}
|
|
|
|
toolResult := executeToolCall(ctx, lcfg, call)
|
|
|
|
// Notify sink
|
|
resultMap := map[string]interface{}{
|
|
"tool_call_id": toolResult.ToolCallID,
|
|
"name": toolResult.Name,
|
|
"content": toolResult.Content,
|
|
"is_error": toolResult.IsError,
|
|
}
|
|
sink.OnToolResult(resultMap)
|
|
|
|
// Emit workspace.file.changed for live editor updates (v0.21.5)
|
|
emitWorkspaceEvent(lcfg, call, toolResult)
|
|
|
|
// 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
|
|
lcfg.Req.Messages = append(lcfg.Req.Messages, providers.Message{
|
|
Role: "tool",
|
|
Content: toolResult.Content,
|
|
ToolCallID: toolResult.ToolCallID,
|
|
Name: toolResult.Name,
|
|
})
|
|
}
|
|
|
|
result.Content += iterContent
|
|
|
|
// Budget check: tokens (post-round)
|
|
if lcfg.Budget.MaxTokens > 0 && (result.InputTokens+result.OutputTokens) > lcfg.Budget.MaxTokens {
|
|
result.BudgetExceeded = "tokens"
|
|
sink.OnFinish("budget_exceeded")
|
|
sink.OnDone()
|
|
return result
|
|
}
|
|
|
|
// Loop: send updated messages back to provider
|
|
}
|
|
|
|
// Hit max iterations
|
|
result.BudgetExceeded = "max_rounds"
|
|
sink.OnMaxRounds()
|
|
sink.OnDone()
|
|
return result
|
|
}
|
|
|
|
// ── Streaming Round ────────────────────────────
|
|
|
|
// runStreamingRound executes one provider StreamCompletion call and
|
|
// processes the event stream. Returns true if the loop should return
|
|
// (normal finish or error), false if tool execution should follow.
|
|
func runStreamingRound(
|
|
ctx context.Context,
|
|
lcfg LoopConfig,
|
|
sink LoopSink,
|
|
result *LoopResult,
|
|
iterContent, iterReasoning *string,
|
|
toolCalls *[]providers.ToolCall,
|
|
iterInput, iterOutput, iterCacheCreate, iterCacheRead *int,
|
|
) bool {
|
|
callStart := time.Now()
|
|
ch, err := lcfg.Provider.StreamCompletion(ctx, lcfg.Cfg, *lcfg.Req)
|
|
if err != nil {
|
|
recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, err)
|
|
result.Error = err
|
|
sink.OnError(err)
|
|
return true
|
|
}
|
|
|
|
for event := range ch {
|
|
// Apply provider-specific post-processing hooks (v0.22.1)
|
|
if hooks := providers.GetHooks(lcfg.ProviderType); hooks != nil {
|
|
hooks.PostStreamEvent(lcfg.Cfg, &event)
|
|
}
|
|
|
|
if event.Error != nil {
|
|
recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, event.Error)
|
|
result.Error = event.Error
|
|
sink.OnError(event.Error)
|
|
return true
|
|
}
|
|
|
|
// Reasoning deltas
|
|
if event.Reasoning != "" {
|
|
*iterReasoning += event.Reasoning
|
|
sink.OnReasoning(event.Reasoning)
|
|
}
|
|
|
|
// Content deltas
|
|
if event.Delta != "" {
|
|
*iterContent += event.Delta
|
|
sink.OnDelta(event.Delta)
|
|
}
|
|
|
|
if event.Done {
|
|
recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, nil)
|
|
|
|
if event.FinishReason == "tool_calls" && len(event.ToolCalls) > 0 {
|
|
*toolCalls = event.ToolCalls
|
|
result.InputTokens += event.InputTokens
|
|
result.OutputTokens += event.OutputTokens
|
|
result.CacheCreationTokens += event.CacheCreationTokens
|
|
result.CacheReadTokens += event.CacheReadTokens
|
|
return false // continue to tool execution
|
|
}
|
|
|
|
// Normal completion
|
|
finishReason := event.FinishReason
|
|
if finishReason == "" {
|
|
finishReason = "stop"
|
|
}
|
|
if *iterReasoning != "" {
|
|
result.Content += "<think>" + *iterReasoning + "</think>"
|
|
}
|
|
result.Content += *iterContent
|
|
result.InputTokens += event.InputTokens
|
|
result.OutputTokens += event.OutputTokens
|
|
result.CacheCreationTokens += event.CacheCreationTokens
|
|
result.CacheReadTokens += event.CacheReadTokens
|
|
|
|
sink.OnFinish(finishReason)
|
|
sink.OnDone()
|
|
return true
|
|
}
|
|
}
|
|
|
|
// Channel closed without Done event — shouldn't happen
|
|
return false
|
|
}
|
|
|
|
// ── Sync Round ─────────────────────────────────
|
|
|
|
// runSyncRound executes one provider ChatCompletion call. Returns true
|
|
// if the loop should return, false if tool execution should follow.
|
|
func runSyncRound(
|
|
ctx context.Context,
|
|
lcfg LoopConfig,
|
|
sink LoopSink,
|
|
result *LoopResult,
|
|
iterContent, iterReasoning *string,
|
|
toolCalls *[]providers.ToolCall,
|
|
iterInput, iterOutput, iterCacheCreate, iterCacheRead *int,
|
|
) bool {
|
|
callStart := time.Now()
|
|
resp, err := lcfg.Provider.ChatCompletion(ctx, lcfg.Cfg, *lcfg.Req)
|
|
recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, err)
|
|
if err != nil {
|
|
result.Error = err
|
|
sink.OnError(err)
|
|
return true
|
|
}
|
|
|
|
*iterContent = resp.Content
|
|
*iterInput = resp.InputTokens
|
|
*iterOutput = resp.OutputTokens
|
|
*iterCacheCreate = resp.CacheCreationTokens
|
|
*iterCacheRead = resp.CacheReadTokens
|
|
|
|
result.InputTokens += resp.InputTokens
|
|
result.OutputTokens += resp.OutputTokens
|
|
result.CacheCreationTokens += resp.CacheCreationTokens
|
|
result.CacheReadTokens += resp.CacheReadTokens
|
|
|
|
if resp.FinishReason == "tool_calls" && len(resp.ToolCalls) > 0 {
|
|
*toolCalls = resp.ToolCalls
|
|
return false // continue to tool execution
|
|
}
|
|
|
|
// Normal completion — deliver full content as single delta
|
|
result.Content += resp.Content
|
|
sink.OnDelta(resp.Content)
|
|
finishReason := resp.FinishReason
|
|
if finishReason == "" {
|
|
finishReason = "stop"
|
|
}
|
|
sink.OnFinish(finishReason)
|
|
sink.OnDone()
|
|
return true
|
|
}
|
|
|
|
// ── Tool Execution ─────────────────────────────
|
|
|
|
// executeToolCall dispatches a single tool call: server-side first,
|
|
// browser bridge fallback, unknown tool error as last resort.
|
|
func executeToolCall(ctx context.Context, lcfg LoopConfig, call tools.ToolCall) tools.ToolResult {
|
|
toolStart := time.Now()
|
|
|
|
// Server-side tool
|
|
if tools.Get(call.Name) != nil {
|
|
log.Printf("🔧 Executing tool (server): %s (call %s)", call.Name, call.ID)
|
|
result := tools.ExecuteCall(ctx, lcfg.ExecCtx, call)
|
|
if lcfg.Health != nil {
|
|
toolLatency := int(time.Since(toolStart).Milliseconds())
|
|
if result.IsError {
|
|
lcfg.Health.RecordToolError(call.Name, toolLatency, result.Content)
|
|
} else {
|
|
lcfg.Health.RecordToolSuccess(call.Name, toolLatency)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// v0.29.2: Starlark extension tool
|
|
if pkg, ok := lcfg.ExtTools[call.Name]; ok && lcfg.Runner != nil {
|
|
log.Printf("🔧 Executing tool (extension): %s (call %s, pkg %s)", call.Name, call.ID, pkg.ID)
|
|
return executeExtensionTool(ctx, lcfg, pkg, call)
|
|
}
|
|
|
|
// Browser bridge (only available when hub is connected)
|
|
if lcfg.Hub != nil && lcfg.Hub.IsConnected(lcfg.ExecCtx.UserID) {
|
|
log.Printf("🔧 Executing tool (browser): %s (call %s)", call.Name, call.ID)
|
|
return executeBrowserTool(lcfg.Hub, lcfg.ExecCtx.UserID, call)
|
|
}
|
|
|
|
// Unknown tool
|
|
log.Printf("⚠️ Unknown tool: %s (call %s)", call.Name, call.ID)
|
|
return tools.ToolResult{
|
|
ToolCallID: call.ID,
|
|
Name: call.Name,
|
|
Content: `{"error":"unknown tool or browser not connected"}`,
|
|
IsError: true,
|
|
}
|
|
}
|
|
|
|
// executeExtensionTool calls a starlark extension's on_tool_call entry point
|
|
// and serializes the return value to a JSON string for the tool result.
|
|
func executeExtensionTool(ctx context.Context, lcfg LoopConfig, pkg *store.PackageRegistration, call tools.ToolCall) tools.ToolResult {
|
|
// Parse JSON arguments into a Starlark dict.
|
|
var rawArgs map[string]interface{}
|
|
if call.Arguments != "" {
|
|
if err := json.Unmarshal([]byte(call.Arguments), &rawArgs); err != nil {
|
|
return tools.ToolResult{
|
|
ToolCallID: call.ID,
|
|
Name: call.Name,
|
|
Content: `{"error":"invalid tool arguments JSON"}`,
|
|
IsError: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build the call dict passed to on_tool_call(call).
|
|
callDict := starlark.NewDict(3)
|
|
_ = callDict.SetKey(starlark.String("tool_name"), starlark.String(call.Name))
|
|
_ = callDict.SetKey(starlark.String("tool_call_id"), starlark.String(call.ID))
|
|
_ = callDict.SetKey(starlark.String("arguments"), jsonToStarlark(rawArgs))
|
|
|
|
rc := &sandbox.RunContext{
|
|
UserID: lcfg.ExecCtx.UserID,
|
|
ChannelID: lcfg.ExecCtx.ChannelID,
|
|
}
|
|
val, output, err := lcfg.Runner.CallEntryPoint(ctx, pkg, "on_tool_call",
|
|
starlark.Tuple{callDict}, nil, rc)
|
|
if output != "" {
|
|
log.Printf(" 🔧 ext tool %s print: %s", pkg.ID, output)
|
|
}
|
|
if err != nil {
|
|
log.Printf("⚠️ ext tool %s on_tool_call error: %v", pkg.ID, err)
|
|
return tools.ToolResult{
|
|
ToolCallID: call.ID,
|
|
Name: call.Name,
|
|
Content: fmt.Sprintf(`{"error":%q}`, err.Error()),
|
|
IsError: true,
|
|
}
|
|
}
|
|
|
|
// Serialize the Starlark return value to JSON.
|
|
content, jsonErr := json.Marshal(starlarkValueToGo(val))
|
|
if jsonErr != nil {
|
|
content = []byte(fmt.Sprintf(`{"error":%q}`, jsonErr.Error()))
|
|
}
|
|
return tools.ToolResult{
|
|
ToolCallID: call.ID,
|
|
Name: call.Name,
|
|
Content: string(content),
|
|
}
|
|
}
|
|
|
|
// jsonToStarlark recursively converts a decoded-JSON value (map/slice/scalar)
|
|
// to its Starlark equivalent.
|
|
func jsonToStarlark(v interface{}) starlark.Value {
|
|
if v == nil {
|
|
return starlark.None
|
|
}
|
|
switch val := v.(type) {
|
|
case map[string]interface{}:
|
|
d := starlark.NewDict(len(val))
|
|
for k, v := range val {
|
|
_ = d.SetKey(starlark.String(k), jsonToStarlark(v))
|
|
}
|
|
return d
|
|
case []interface{}:
|
|
elems := make([]starlark.Value, len(val))
|
|
for i, v := range val {
|
|
elems[i] = jsonToStarlark(v)
|
|
}
|
|
return starlark.NewList(elems)
|
|
case string:
|
|
return starlark.String(val)
|
|
case float64:
|
|
if val == float64(int64(val)) {
|
|
return starlark.MakeInt64(int64(val))
|
|
}
|
|
return starlark.Float(val)
|
|
case bool:
|
|
return starlark.Bool(val)
|
|
default:
|
|
return starlark.String(fmt.Sprintf("%v", val))
|
|
}
|
|
}
|
|
|
|
// starlarkValueToGo recursively converts a Starlark value to a Go value
|
|
// suitable for json.Marshal.
|
|
func starlarkValueToGo(v starlark.Value) interface{} {
|
|
if v == nil || v == starlark.None {
|
|
return nil
|
|
}
|
|
switch val := v.(type) {
|
|
case starlark.String:
|
|
return string(val)
|
|
case starlark.Int:
|
|
i64, ok := val.Int64()
|
|
if ok {
|
|
return i64
|
|
}
|
|
return val.String()
|
|
case starlark.Float:
|
|
return float64(val)
|
|
case starlark.Bool:
|
|
return bool(val)
|
|
case *starlark.Dict:
|
|
m := make(map[string]interface{}, val.Len())
|
|
for _, item := range val.Items() {
|
|
k, ok := starlark.AsString(item[0])
|
|
if !ok {
|
|
k = item[0].String()
|
|
}
|
|
m[k] = starlarkValueToGo(item[1])
|
|
}
|
|
return m
|
|
case *starlark.List:
|
|
list := make([]interface{}, val.Len())
|
|
for i := 0; i < val.Len(); i++ {
|
|
list[i] = starlarkValueToGo(val.Index(i))
|
|
}
|
|
return list
|
|
default:
|
|
return v.String()
|
|
}
|
|
}
|
|
|
|
// emitWorkspaceEvent sends workspace.file.changed when a write tool succeeds.
|
|
func emitWorkspaceEvent(lcfg LoopConfig, call tools.ToolCall, result tools.ToolResult) {
|
|
if lcfg.Hub == nil || result.IsError || lcfg.ExecCtx.WorkspaceID == "" {
|
|
return
|
|
}
|
|
if call.Name != "workspace_write" && call.Name != "workspace_patch" {
|
|
return
|
|
}
|
|
var toolArgs struct {
|
|
Path string `json:"path"`
|
|
}
|
|
if json.Unmarshal([]byte(call.Arguments), &toolArgs) == nil && toolArgs.Path != "" {
|
|
lcfg.Hub.SendToUser(lcfg.ExecCtx.UserID, events.Event{
|
|
Label: "workspace.file.changed",
|
|
Payload: events.MustJSON(map[string]string{
|
|
"workspace_id": lcfg.ExecCtx.WorkspaceID,
|
|
"path": toolArgs.Path,
|
|
}),
|
|
})
|
|
}
|
|
}
|
|
|
|
// ── Sink Implementations ───────────────────────
|
|
|
|
// sseSink writes SSE events to a gin.Context writer for interactive streaming.
|
|
type sseSink struct {
|
|
w http.ResponseWriter
|
|
flush func()
|
|
model string
|
|
}
|
|
|
|
func newSSESink(c *gin.Context, model string) *sseSink {
|
|
flusher, _ := c.Writer.(http.Flusher)
|
|
return &sseSink{
|
|
w: c.Writer,
|
|
model: model,
|
|
flush: func() {
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
func (s *sseSink) sendData(data string) {
|
|
fmt.Fprintf(s.w, "data: %s\n\n", data)
|
|
s.flush()
|
|
}
|
|
|
|
func (s *sseSink) sendEvent(event, data string) {
|
|
fmt.Fprintf(s.w, "event: %s\ndata: %s\n\n", event, data)
|
|
s.flush()
|
|
}
|
|
|
|
func (s *sseSink) OnDelta(delta string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q}`,
|
|
delta, s.model))
|
|
}
|
|
|
|
func (s *sseSink) OnReasoning(delta string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q}`,
|
|
delta, s.model))
|
|
}
|
|
|
|
func (s *sseSink) OnError(err error) {
|
|
s.sendData(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
|
|
}
|
|
|
|
func (s *sseSink) OnToolUse(calls []providers.ToolCall) {
|
|
toolCallsJSON, _ := json.Marshal(calls)
|
|
s.sendEvent("tool_use", string(toolCallsJSON))
|
|
}
|
|
|
|
func (s *sseSink) OnToolResult(result map[string]interface{}) {
|
|
resultJSON, _ := json.Marshal(result)
|
|
s.sendEvent("tool_result", string(resultJSON))
|
|
}
|
|
|
|
func (s *sseSink) OnFinish(reason string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
|
|
reason, s.model))
|
|
}
|
|
|
|
func (s *sseSink) OnDone() {
|
|
s.sendData("[DONE]")
|
|
}
|
|
|
|
func (s *sseSink) OnMaxRounds() {
|
|
log.Printf("⚠️ Tool loop hit max iterations for model %s", s.model)
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q}`,
|
|
escapeJSON("[Tool execution limit reached]"), s.model))
|
|
}
|
|
|
|
// sseModelSink extends sseSink with model_display attribution for
|
|
// multi-model streaming. Does NOT send [DONE] — the multiModelStream
|
|
// caller sends it after all models finish.
|
|
type sseModelSink struct {
|
|
sseSink
|
|
displayName string
|
|
}
|
|
|
|
func newSSEModelSink(c *gin.Context, model, displayName string) *sseModelSink {
|
|
base := newSSESink(c, model)
|
|
return &sseModelSink{
|
|
sseSink: *base,
|
|
displayName: displayName,
|
|
}
|
|
}
|
|
|
|
func (s *sseModelSink) OnDelta(delta string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
|
|
delta, s.model, s.displayName))
|
|
}
|
|
|
|
func (s *sseModelSink) OnReasoning(delta string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
|
|
delta, s.model, s.displayName))
|
|
}
|
|
|
|
func (s *sseModelSink) OnFinish(reason string) {
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q,"model_display":%q}`,
|
|
reason, s.model, s.displayName))
|
|
}
|
|
|
|
// OnDone is a no-op for multi-model — caller sends [DONE] after all models.
|
|
func (s *sseModelSink) OnDone() {}
|
|
|
|
func (s *sseModelSink) OnMaxRounds() {
|
|
log.Printf("⚠️ Multi-model tool loop hit max iterations for model %s", s.displayName)
|
|
s.sendData(fmt.Sprintf(
|
|
`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q,"model_display":%q}`,
|
|
escapeJSON("[Tool execution limit reached]"), s.model, s.displayName))
|
|
}
|
|
|
|
// accumSink accumulates silently — used by syncCompletion where the
|
|
// caller formats the HTTP response after the loop returns.
|
|
type accumSink struct{}
|
|
|
|
func (accumSink) OnDelta(string) {}
|
|
func (accumSink) OnReasoning(string) {}
|
|
func (accumSink) OnError(error) {}
|
|
func (accumSink) OnToolUse([]providers.ToolCall) {}
|
|
func (accumSink) OnToolResult(map[string]interface{}) {}
|
|
func (accumSink) OnFinish(string) {}
|
|
func (accumSink) OnDone() {}
|
|
func (accumSink) OnMaxRounds() {}
|
|
|
|
// HeadlessSink is used by the task scheduler — accumulates and logs.
|
|
// No client connection, no SSE output.
|
|
type HeadlessSink struct {
|
|
taskID string
|
|
}
|
|
|
|
func NewHeadlessSink(taskID string) *HeadlessSink {
|
|
return &HeadlessSink{taskID: taskID}
|
|
}
|
|
|
|
func (s *HeadlessSink) OnDelta(string) {}
|
|
func (s *HeadlessSink) OnReasoning(string) {}
|
|
|
|
func (s *HeadlessSink) OnError(err error) {
|
|
log.Printf("[task:%s] Error: %v", s.taskID, err)
|
|
}
|
|
|
|
func (s *HeadlessSink) OnToolUse(calls []providers.ToolCall) {
|
|
names := make([]string, len(calls))
|
|
for i, tc := range calls {
|
|
names[i] = tc.Function.Name
|
|
}
|
|
log.Printf("[task:%s] Tool calls: %v", s.taskID, names)
|
|
}
|
|
|
|
func (s *HeadlessSink) OnToolResult(result map[string]interface{}) {
|
|
name, _ := result["name"].(string)
|
|
isErr, _ := result["is_error"].(bool)
|
|
if isErr {
|
|
log.Printf("[task:%s] Tool %s failed: %v", s.taskID, name, result["content"])
|
|
}
|
|
}
|
|
|
|
func (s *HeadlessSink) OnFinish(reason string) {
|
|
log.Printf("[task:%s] Finished: %s", s.taskID, reason)
|
|
}
|
|
|
|
func (s *HeadlessSink) OnDone() {}
|
|
|
|
func (s *HeadlessSink) OnMaxRounds() {
|
|
log.Printf("[task:%s] Hit max tool iterations", s.taskID)
|
|
}
|