diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 7ce2a37a6..1c3af0552 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -686,7 +686,7 @@ func (al *AgentLoop) runAgentLoop( agent.Sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) // 4. Run LLM iteration loop - finalContent, iteration, err := al.runLLMIteration(ctx, agent, messages, opts) + finalContent, iteration, streamSent, err := al.runLLMIteration(ctx, agent, messages, opts) if err != nil { return "", err } @@ -708,8 +708,8 @@ func (al *AgentLoop) runAgentLoop( al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) } - // 8. Optional: send response via bus - if opts.SendResponse { + // 8. Optional: send response via bus (skip if streaming already delivered it) + if opts.SendResponse && !streamSent { al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, ChatID: opts.ChatID, @@ -727,6 +727,11 @@ func (al *AgentLoop) runAgentLoop( "final_length": len(finalContent), }) + // When streaming already delivered the message, return empty so the caller + // (Run loop) doesn't publish a duplicate via PublishOutbound. + if streamSent { + return "", nil + } return finalContent, nil } @@ -787,14 +792,25 @@ func (al *AgentLoop) handleReasoning( } // runLLMIteration executes the LLM call loop with tool handling. +// Returns (finalContent, iteration, streamed, error). +// When streamed is true, the response was already delivered to the user via +// streaming (sendMessageDraft + sendMessage) and PublishOutbound should be skipped. func (al *AgentLoop) runLLMIteration( ctx context.Context, agent *AgentInstance, messages []providers.Message, opts processOptions, -) (string, int, error) { +) (string, int, bool, error) { iteration := 0 var finalContent string + var streamed bool + + // Check if both the provider and channel support streaming + streamProvider, providerCanStream := agent.Provider.(providers.StreamingProvider) + var streamer bus.Streamer + if providerCanStream && !constants.IsInternalChannel(opts.Channel) { + streamer, _ = al.bus.GetStreamer(ctx, opts.Channel, opts.ChatID) + } for iteration < agent.MaxIterations { iteration++ @@ -834,22 +850,30 @@ func (al *AgentLoop) runLLMIteration( var response *providers.LLMResponse var err error + llmOpts := map[string]any{ + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + "prompt_cache_key": agent.ID, + } + callLLM := func() (*providers.LLMResponse, error) { + // Use streaming when available (streamer obtained, provider supports it) + if streamer != nil && streamProvider != nil { + return streamProvider.ChatStream( + ctx, messages, providerToolDefs, agent.Model, llmOpts, + func(accumulated string) { + streamer.Update(ctx, accumulated) + }, + ) + } + if len(agent.Candidates) > 1 && al.fallback != nil { fbResult, fbErr := al.fallback.Execute( ctx, agent.Candidates, func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { return agent.Provider.Chat( - ctx, - messages, - providerToolDefs, - model, - map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - "prompt_cache_key": agent.ID, - }, + ctx, messages, providerToolDefs, model, llmOpts, ) }, ) @@ -866,11 +890,7 @@ func (al *AgentLoop) runLLMIteration( } return fbResult.Response, nil } - return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - "prompt_cache_key": agent.ID, - }) + return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, llmOpts) } // Retry loop for context/token errors @@ -949,7 +969,7 @@ func (al *AgentLoop) runLLMIteration( "iteration": iteration, "error": err.Error(), }) - return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err) + return "", iteration, false, fmt.Errorf("LLM call failed after retries: %w", err) } go al.handleReasoning( @@ -972,15 +992,33 @@ func (al *AgentLoop) runLLMIteration( // Check if no tool calls - we're done if len(response.ToolCalls) == 0 { finalContent = response.Content + + // If we were streaming, finalize the message (sends the permanent message) + if streamer != nil { + if err := streamer.Finalize(ctx, finalContent); err != nil { + logger.WarnCF("agent", "Stream finalize failed", map[string]any{ + "error": err.Error(), + }) + } else { + streamed = true + } + } + logger.InfoCF("agent", "LLM response without tool calls (direct answer)", map[string]any{ "agent_id": agent.ID, "iteration": iteration, "content_chars": len(finalContent), + "streamed": streamed, }) break } + // Tool calls detected — cancel any active stream (draft auto-expires) + if streamer != nil { + streamer.Cancel(ctx) + } + normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) for _, tc := range response.ToolCalls { normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) @@ -1136,7 +1174,7 @@ func (al *AgentLoop) runLLMIteration( } } - return finalContent, iteration, nil + return finalContent, iteration, streamed, nil } // updateToolContexts updates the context for tools that need channel/chatID info. diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go index f5ff9587d..eff4eb94b 100644 --- a/pkg/bus/bus.go +++ b/pkg/bus/bus.go @@ -13,12 +13,29 @@ var ErrBusClosed = errors.New("message bus closed") const defaultBusBufferSize = 64 +// StreamDelegate is implemented by the channel Manager to provide streaming +// capabilities to the agent loop without tight coupling. +type StreamDelegate interface { + // GetStreamer returns a Streamer for the given channel+chatID if the channel + // supports streaming. Returns nil, false if streaming is unavailable. + GetStreamer(ctx context.Context, channel, chatID string) (Streamer, bool) +} + +// Streamer pushes incremental content to a streaming-capable channel. +// Defined here so the agent loop can use it without importing pkg/channels. +type Streamer interface { + Update(ctx context.Context, content string) error + Finalize(ctx context.Context, content string) error + Cancel(ctx context.Context) +} + type MessageBus struct { - inbound chan InboundMessage - outbound chan OutboundMessage - outboundMedia chan OutboundMediaMessage - done chan struct{} - closed atomic.Bool + inbound chan InboundMessage + outbound chan OutboundMessage + outboundMedia chan OutboundMediaMessage + done chan struct{} + closed atomic.Bool + streamDelegate atomic.Value // stores StreamDelegate } func NewMessageBus() *MessageBus { @@ -114,6 +131,19 @@ func (mb *MessageBus) SubscribeOutboundMedia(ctx context.Context) (OutboundMedia } } +// SetStreamDelegate registers a StreamDelegate (typically the channel Manager). +func (mb *MessageBus) SetStreamDelegate(d StreamDelegate) { + mb.streamDelegate.Store(d) +} + +// GetStreamer returns a Streamer for the given channel+chatID via the delegate. +func (mb *MessageBus) GetStreamer(ctx context.Context, channel, chatID string) (Streamer, bool) { + if d, ok := mb.streamDelegate.Load().(StreamDelegate); ok && d != nil { + return d.GetStreamer(ctx, channel, chatID) + } + return nil, false +} + func (mb *MessageBus) Close() { if mb.closed.CompareAndSwap(false, true) { close(mb.done) diff --git a/pkg/channels/base.go b/pkg/channels/base.go index 063a66523..1f4b998bc 100644 --- a/pkg/channels/base.go +++ b/pkg/channels/base.go @@ -271,20 +271,24 @@ func (c *BaseChannel) HandleMessage( // Auto-trigger typing indicator, message reaction, and placeholder before publishing. // Each capability is independent — all three may fire for the same message. + // Note: even when streaming is available, we still show typing + placeholder on inbound. + // If streaming actually activates, preSend will skip the placeholder edit (streamActive map) + // and the typing stop will still be called. This avoids the problem of compile-time interface + // checks incorrectly skipping indicators when streaming may not work at runtime. if c.owner != nil && c.placeholderRecorder != nil { - // Typing — independent pipeline + // Typing if tc, ok := c.owner.(TypingCapable); ok { if stop, err := tc.StartTyping(ctx, chatID); err == nil { c.placeholderRecorder.RecordTypingStop(c.name, chatID, stop) } } - // Reaction — independent pipeline + // Reaction if rc, ok := c.owner.(ReactionCapable); ok && messageID != "" { if undo, err := rc.ReactToMessage(ctx, chatID, messageID); err == nil { c.placeholderRecorder.RecordReactionUndo(c.name, chatID, undo) } } - // Placeholder — independent pipeline + // Placeholder if pc, ok := c.owner.(PlaceholderCapable); ok { if phID, err := pc.SendPlaceholder(ctx, chatID); err == nil && phID != "" { c.placeholderRecorder.RecordPlaceholder(c.name, chatID, phID) diff --git a/pkg/channels/interfaces.go b/pkg/channels/interfaces.go index 74caeeac5..580157ce0 100644 --- a/pkg/channels/interfaces.go +++ b/pkg/channels/interfaces.go @@ -31,6 +31,24 @@ type PlaceholderCapable interface { SendPlaceholder(ctx context.Context, chatID string) (messageID string, err error) } +// StreamingCapable — channels that can show partial LLM output in real-time. +// The channel SHOULD gracefully degrade if the platform rejects streaming +// (e.g. Telegram bot without forum mode). In that case, Update becomes a no-op +// and Finalize still delivers the final message. +type StreamingCapable interface { + BeginStream(ctx context.Context, chatID string) (Streamer, error) +} + +// Streamer pushes incremental content to a streaming-capable channel. +type Streamer interface { + // Update sends accumulated partial content to the user. + Update(ctx context.Context, content string) error + // Finalize commits the final message. After this, the Streamer is done. + Finalize(ctx context.Context, content string) error + // Cancel aborts the stream (e.g. when tool calls are detected mid-stream). + Cancel(ctx context.Context) +} + // PlaceholderRecorder is injected into channels by Manager. // Channels call these methods on inbound to register typing/placeholder state. // Manager uses the registered state on outbound to stop typing and edit placeholders. diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index fdd6d0c1f..ba56bd789 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -86,6 +86,7 @@ type Manager struct { placeholders sync.Map // "channel:chatID" → placeholderID (string) typingStops sync.Map // "channel:chatID" → func() reactionUndos sync.Map // "channel:chatID" → reactionEntry + streamActive sync.Map // "channel:chatID" → true (set when streamer.Finalize sent the message) } type asyncTask struct { @@ -114,7 +115,7 @@ func (m *Manager) RecordReactionUndo(channel, chatID string, undo func()) { } // preSend handles typing stop, reaction undo, and placeholder editing before sending a message. -// Returns true if the message was edited into a placeholder (skip Send). +// Returns true if the message was already delivered (skip Send). func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMessage, ch Channel) bool { key := name + ":" + msg.ChatID @@ -132,7 +133,14 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess } } - // 3. Try editing placeholder + // 3. If a stream already finalized this message, skip both placeholder and send + if _, loaded := m.streamActive.LoadAndDelete(key); loaded { + // Also clean up any stale placeholder + m.placeholders.Delete(key) + return true + } + + // 4. Try editing placeholder if v, loaded := m.placeholders.LoadAndDelete(key); loaded { if entry, ok := v.(placeholderEntry); ok && entry.id != "" { if editor, ok := ch.(MessageEditor); ok { @@ -156,6 +164,9 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi mediaStore: store, } + // Register as streaming delegate so the agent loop can obtain streamers + messageBus.SetStreamDelegate(m) + if err := m.initChannels(); err != nil { return nil, err } @@ -163,6 +174,59 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi return m, nil } +// GetStreamer implements bus.StreamDelegate. +// It checks if the named channel supports streaming and returns a Streamer. +func (m *Manager) GetStreamer(ctx context.Context, channelName, chatID string) (bus.Streamer, bool) { + m.mu.RLock() + ch, exists := m.channels[channelName] + m.mu.RUnlock() + + if !exists { + return nil, false + } + + sc, ok := ch.(StreamingCapable) + if !ok { + return nil, false + } + + streamer, err := sc.BeginStream(ctx, chatID) + if err != nil { + logger.DebugCF("channels", "Streaming unavailable, falling back to placeholder", map[string]any{ + "channel": channelName, + "error": err.Error(), + }) + return nil, false + } + + // Wrap the channels.Streamer to track finalization for preSend coordination + return &managerStreamer{ + inner: streamer, + manager: m, + key: channelName + ":" + chatID, + }, true +} + +// managerStreamer wraps a channels.Streamer to mark streamActive on Finalize. +type managerStreamer struct { + inner Streamer + manager *Manager + key string +} + +func (s *managerStreamer) Update(ctx context.Context, content string) error { + return s.inner.Update(ctx, content) +} + +func (s *managerStreamer) Finalize(ctx context.Context, content string) error { + s.manager.streamActive.Store(s.key, true) + return s.inner.Finalize(ctx, content) +} + +func (s *managerStreamer) Cancel(ctx context.Context) { + s.inner.Cancel(ctx) +} + // initChannel is a helper that looks up a factory by name and creates the channel. func (m *Manager) initChannel(name, displayName string) { f, ok := getFactory(name) diff --git a/pkg/channels/telegram/telegram.go b/pkg/channels/telegram/telegram.go index f328f32b8..a7df1ebc5 100644 --- a/pkg/channels/telegram/telegram.go +++ b/pkg/channels/telegram/telegram.go @@ -8,8 +8,10 @@ import ( "os" "regexp" "slices" + "math/rand" "strconv" "strings" + "sync" "time" "github.com/mymmrac/telego" @@ -767,3 +769,95 @@ func (c *TelegramChannel) stripBotMention(content string) string { content = re.ReplaceAllString(content, "") return strings.TrimSpace(content) } + +// BeginStream implements channels.StreamingCapable. +func (c *TelegramChannel) BeginStream(ctx context.Context, chatID string) (channels.Streamer, error) { + if !c.config.Channels.Telegram.Streaming.Enabled { + return nil, fmt.Errorf("streaming disabled in config") + } + + cid, err := parseChatID(chatID) + if err != nil { + return nil, err + } + + return &telegramStreamer{ + bot: c.bot, + chatID: cid, + draftID: rand.Intn(1<<31-1) + 1, // non-zero random draft ID + }, nil +} + +// telegramStreamer streams partial LLM output via Telegram's sendMessageDraft API. +// On first API error (e.g. bot lacks forum mode), it silently degrades: Update +// becomes a no-op, while Finalize still delivers the final message. +type telegramStreamer struct { + bot *telego.Bot + chatID int64 + draftID int + lastLen int + lastAt time.Time + failed bool + mu sync.Mutex +} + +const ( + streamThrottleInterval = 3 * time.Second + streamMinGrowth = 200 // minimum character growth to send an update +) + +func (s *telegramStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.failed { + return nil + } + + // Throttle: skip if not enough time or content has passed + now := time.Now() + growth := len(content) - s.lastLen + if s.lastLen > 0 && now.Sub(s.lastAt) < streamThrottleInterval && growth < streamMinGrowth { + return nil + } + + htmlContent := markdownToTelegramHTML(content) + + err := s.bot.SendMessageDraft(ctx, &telego.SendMessageDraftParams{ + ChatID: s.chatID, + DraftID: s.draftID, + Text: htmlContent, + ParseMode: telego.ModeHTML, + }) + if err != nil { + // First error → degrade silently (e.g. no forum mode) + logger.WarnCF("telegram", "sendMessageDraft failed, disabling streaming", map[string]any{ + "error": err.Error(), + }) + s.failed = true + return nil // don't propagate — Finalize will still deliver + } + + s.lastLen = len(content) + s.lastAt = now + return nil +} + +func (s *telegramStreamer) Finalize(ctx context.Context, content string) error { + htmlContent := markdownToTelegramHTML(content) + tgMsg := tu.Message(tu.ID(s.chatID), htmlContent) + tgMsg.ParseMode = telego.ModeHTML + + if _, err := s.bot.SendMessage(ctx, tgMsg); err != nil { + // Fallback to plain text + tgMsg.ParseMode = "" + if _, err = s.bot.SendMessage(ctx, tgMsg); err != nil { + return fmt.Errorf("telegram finalize: %w", err) + } + } + return nil +} + +func (s *telegramStreamer) Cancel(ctx context.Context) { + // Draft auto-expires on Telegram's side; nothing to clean up. +} diff --git a/pkg/config/config.go b/pkg/config/config.go index f40e05e1c..9d852b974 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -237,6 +237,10 @@ type PlaceholderConfig struct { Text string `json:"text,omitempty"` } +type StreamingConfig struct { + Enabled bool `json:"enabled,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED"` +} + type WhatsAppConfig struct { Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WHATSAPP_ENABLED"` BridgeURL string `json:"bridge_url" env:"PICOCLAW_CHANNELS_WHATSAPP_BRIDGE_URL"` @@ -255,6 +259,7 @@ type TelegramConfig struct { GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` Typing TypingConfig `json:"typing,omitempty"` Placeholder PlaceholderConfig `json:"placeholder,omitempty"` + Streaming StreamingConfig `json:"streaming,omitempty"` ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_TELEGRAM_REASONING_CHANNEL_ID"` } diff --git a/pkg/providers/anthropic/provider.go b/pkg/providers/anthropic/provider.go index 1bb15f771..23b1f9561 100644 --- a/pkg/providers/anthropic/provider.go +++ b/pkg/providers/anthropic/provider.go @@ -9,6 +9,7 @@ import ( "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/option" + "github.com/anthropics/anthropic-sdk-go/packages/ssestream" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -93,6 +94,136 @@ func (p *Provider) Chat( return parseResponse(resp), nil } +// ChatStream implements providers.StreamingProvider. +// It streams text deltas to onChunk (with accumulated text) while building +// the same LLMResponse returned by Chat for tool-call compatibility. +func (p *Provider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), +) (*LLMResponse, error) { + var opts []option.RequestOption + if p.tokenSource != nil { + tok, err := p.tokenSource() + if err != nil { + return nil, fmt.Errorf("refreshing token: %w", err) + } + opts = append(opts, option.WithAuthToken(tok)) + } + + params, err := buildParams(messages, tools, model, options) + if err != nil { + return nil, err + } + + stream := p.client.Messages.NewStreaming(ctx, params, opts...) + return parseStream(stream, onChunk) +} + +// parseStream consumes a streaming response and builds an LLMResponse. +func parseStream( + stream *ssestream.Stream[anthropic.MessageStreamEventUnion], + onChunk func(accumulated string), +) (*LLMResponse, error) { + defer stream.Close() + + var textContent strings.Builder + var toolCalls []ToolCall + var stopReason anthropic.StopReason + var inputTokens, outputTokens int64 + + // Track tool_use blocks being assembled (by content block index) + type toolBlock struct { + id string + name string + inputJSON strings.Builder + } + activeTools := map[int64]*toolBlock{} + + for stream.Next() { + event := stream.Current() + + switch event.Type { + case "message_start": + if event.Message.Usage.InputTokens > 0 { + inputTokens = event.Message.Usage.InputTokens + } + + case "content_block_start": + cb := event.ContentBlock + if cb.Type == "tool_use" { + activeTools[event.Index] = &toolBlock{ + id: cb.ID, + name: cb.Name, + } + } + + case "content_block_delta": + delta := event.Delta + switch delta.Type { + case "text_delta": + textContent.WriteString(delta.Text) + if onChunk != nil { + onChunk(textContent.String()) + } + case "input_json_delta": + if tb, ok := activeTools[event.Index]; ok { + tb.inputJSON.WriteString(delta.PartialJSON) + } + } + + case "content_block_stop": + if tb, ok := activeTools[event.Index]; ok { + var args map[string]any + if err := json.Unmarshal([]byte(tb.inputJSON.String()), &args); err != nil { + log.Printf("anthropic stream: failed to decode tool call input for %q: %v", tb.name, err) + args = map[string]any{"raw": tb.inputJSON.String()} + } + toolCalls = append(toolCalls, ToolCall{ + ID: tb.id, + Name: tb.name, + Arguments: args, + }) + delete(activeTools, event.Index) + } + + case "message_delta": + stopReason = event.Delta.StopReason + if event.Usage.OutputTokens > 0 { + outputTokens = event.Usage.OutputTokens + } + } + } + + if err := stream.Err(); err != nil { + return nil, fmt.Errorf("claude streaming API call: %w", err) + } + + finishReason := "stop" + switch stopReason { + case anthropic.StopReasonToolUse: + finishReason = "tool_calls" + case anthropic.StopReasonMaxTokens: + finishReason = "length" + case anthropic.StopReasonEndTurn: + finishReason = "stop" + } + + return &LLMResponse{ + Content: textContent.String(), + ToolCalls: toolCalls, + FinishReason: finishReason, + Usage: &UsageInfo{ + PromptTokens: int(inputTokens), + CompletionTokens: int(outputTokens), + TotalTokens: int(inputTokens + outputTokens), + }, + }, nil +} + func (p *Provider) GetDefaultModel() string { return "claude-sonnet-4.6" } diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go index 5c328f418..58b410aab 100644 --- a/pkg/providers/http_provider.go +++ b/pkg/providers/http_provider.go @@ -52,6 +52,19 @@ func (p *HTTPProvider) Chat( return p.delegate.Chat(ctx, messages, tools, model, options) } +// ChatStream implements providers.StreamingProvider by delegating to the +// OpenAI-compatible streaming endpoint (SSE with stream: true). +func (p *HTTPProvider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), +) (*LLMResponse, error) { + return p.delegate.ChatStream(ctx, messages, tools, model, options, onChunk) +} + func (p *HTTPProvider) GetDefaultModel() string { return "" } diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index ff9109e96..ce6b1f6d7 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -1,6 +1,7 @@ package openai_compat import ( + "bufio" "bytes" "context" "encoding/json" @@ -101,17 +102,8 @@ func NewProviderWithMaxTokensFieldAndTimeout( ) } -func (p *Provider) Chat( - ctx context.Context, - messages []Message, - tools []ToolDefinition, - model string, - options map[string]any, -) (*LLMResponse, error) { - if p.apiBase == "" { - return nil, fmt.Errorf("API base not configured") - } - +// buildRequestBody constructs the common request body for Chat and ChatStream. +func (p *Provider) buildRequestBody(messages []Message, tools []ToolDefinition, model string, options map[string]any) map[string]any { model = normalizeModel(model, p.apiBase) requestBody := map[string]any{ @@ -125,10 +117,8 @@ func (p *Provider) Chat( } if maxTokens, ok := asInt(options["max_tokens"]); ok { - // Use configured maxTokensField if specified, otherwise fallback to model-based detection fieldName := p.maxTokensField if fieldName == "" { - // Fallback: detect from model name for backward compatibility lowerModel := strings.ToLower(model) if strings.Contains(lowerModel, "glm") || strings.Contains(lowerModel, "o1") || strings.Contains(lowerModel, "gpt-5") { @@ -142,7 +132,6 @@ func (p *Provider) Chat( if temperature, ok := asFloat(options["temperature"]); ok { lowerModel := strings.ToLower(model) - // Kimi k2 models only support temperature=1. if strings.Contains(lowerModel, "kimi") && strings.Contains(lowerModel, "k2") { requestBody["temperature"] = 1.0 } else { @@ -150,18 +139,28 @@ func (p *Provider) Chat( } } - // Prompt caching: pass a stable cache key so OpenAI can bucket requests - // with the same key and reuse prefix KV cache across calls. - // The key is typically the agent ID — stable per agent, shared across requests. - // See: https://platform.openai.com/docs/guides/prompt-caching - // Prompt caching is only supported by OpenAI-native endpoints. - // Gemini and other providers reject unknown fields, so skip for non-OpenAI APIs. if cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" { if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") { requestBody["prompt_cache_key"] = cacheKey } } + return requestBody +} + +func (p *Provider) Chat( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, +) (*LLMResponse, error) { + if p.apiBase == "" { + return nil, fmt.Errorf("API base not configured") + } + + requestBody := p.buildRequestBody(messages, tools, model, options) + jsonData, err := json.Marshal(requestBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) @@ -195,6 +194,186 @@ func (p *Provider) Chat( return parseResponse(body) } +// ChatStream implements streaming via OpenAI-compatible SSE (stream: true). +// onChunk receives the accumulated text so far on each text delta. +func (p *Provider) ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), +) (*LLMResponse, error) { + if p.apiBase == "" { + return nil, fmt.Errorf("API base not configured") + } + + requestBody := p.buildRequestBody(messages, tools, model, options) + requestBody["stream"] = true + + jsonData, err := json.Marshal(requestBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", p.apiBase+"/chat/completions", bytes.NewReader(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if p.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+p.apiKey) + } + + // Use a client without Timeout for streaming — the http.Client.Timeout covers + // the entire request lifecycle including body reads, which would kill long streams. + // Context cancellation still provides the safety net. + streamClient := &http.Client{Transport: p.httpClient.Transport} + resp, err := streamClient.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body)) + } + + return parseStreamResponse(resp.Body, onChunk) +} + +// parseStreamResponse parses an OpenAI-compatible SSE stream. +func parseStreamResponse(reader io.Reader, onChunk func(accumulated string)) (*LLMResponse, error) { + var textContent strings.Builder + var finishReason string + var usage *UsageInfo + + // Tool call assembly: OpenAI streams tool calls as incremental deltas + type toolAccum struct { + id string + name string + argsJSON strings.Builder + } + activeTools := map[int]*toolAccum{} + + scanner := bufio.NewScanner(reader) + for scanner.Scan() { + line := scanner.Text() + + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + break + } + + var chunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + ToolCalls []struct { + Index int `json:"index"` + ID string `json:"id"` + Function *struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"delta"` + FinishReason *string `json:"finish_reason"` + } `json:"choices"` + Usage *UsageInfo `json:"usage"` + } + + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue // skip malformed chunks + } + + if chunk.Usage != nil { + usage = chunk.Usage + } + + if len(chunk.Choices) == 0 { + continue + } + + choice := chunk.Choices[0] + + // Accumulate text content + if choice.Delta.Content != "" { + textContent.WriteString(choice.Delta.Content) + if onChunk != nil { + onChunk(textContent.String()) + } + } + + // Accumulate tool call deltas + for _, tc := range choice.Delta.ToolCalls { + acc, ok := activeTools[tc.Index] + if !ok { + acc = &toolAccum{} + activeTools[tc.Index] = acc + } + if tc.ID != "" { + acc.id = tc.ID + } + if tc.Function != nil { + if tc.Function.Name != "" { + acc.name = tc.Function.Name + } + if tc.Function.Arguments != "" { + acc.argsJSON.WriteString(tc.Function.Arguments) + } + } + } + + if choice.FinishReason != nil { + finishReason = *choice.FinishReason + } + } + + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("streaming read error: %w", err) + } + + // Assemble tool calls from accumulated deltas + var toolCalls []ToolCall + for i := 0; i < len(activeTools); i++ { + acc, ok := activeTools[i] + if !ok { + continue + } + args := make(map[string]any) + raw := acc.argsJSON.String() + if raw != "" { + if err := json.Unmarshal([]byte(raw), &args); err != nil { + log.Printf("openai_compat stream: failed to decode tool call arguments for %q: %v", acc.name, err) + args["raw"] = raw + } + } + toolCalls = append(toolCalls, ToolCall{ + ID: acc.id, + Name: acc.name, + Arguments: args, + }) + } + + if finishReason == "" { + finishReason = "stop" + } + + return &LLMResponse{ + Content: textContent.String(), + ToolCalls: toolCalls, + FinishReason: finishReason, + Usage: usage, + }, nil +} + func parseResponse(body []byte) (*LLMResponse, error) { var apiResponse struct { Choices []struct { diff --git a/pkg/providers/types.go b/pkg/providers/types.go index f0c168bc6..038ac69dc 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -37,6 +37,20 @@ type StatefulProvider interface { Close() } +// StreamingProvider is an optional interface for providers that support token streaming. +// onChunk receives the accumulated text so far (not individual deltas). +// The returned LLMResponse is the same complete response for compatibility with tool-call handling. +type StreamingProvider interface { + ChatStream( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), + ) (*LLMResponse, error) +} + // FailoverReason classifies why an LLM request failed for fallback decisions. type FailoverReason string