feat(telegram): stream LLM responses in real-time via sendMessageDraft

Implements real-time token streaming to Telegram using the sendMessageDraft
API (telego v1.6.0). Instead of showing only a "Thinking..." placeholder
until the full response arrives, users now see partial LLM output appear
in the chat as it's generated.

The streaming pipeline threads through all layers:

- StreamingProvider interface (providers/types.go): opt-in ChatStream()
  method that receives an onChunk callback with accumulated text
- OpenAI-compatible SSE streaming (openai_compat/provider.go): parses
  SSE events with stream:true, handles text deltas and tool call assembly
- Anthropic native streaming (anthropic/provider.go): uses SDK's
  NewStreaming() for direct Anthropic API connections
- HTTPProvider delegation (http_provider.go): delegates ChatStream to
  the underlying openai_compat provider
- StreamingCapable + Streamer interfaces (channels/interfaces.go):
  opt-in channel capability like TypingCapable/PlaceholderCapable
- Telegram streamer (telegram/telegram.go): BeginStream returns a
  telegramStreamer that throttles sendMessageDraft calls (3s/200 chars)
  with graceful degradation on API errors
- StreamDelegate bridge (bus/bus.go): decouples agent loop from channel
  manager without tight imports
- Manager integration (manager.go): implements StreamDelegate, tracks
  streamActive state, coordinates with placeholder editing
- Agent loop (loop.go): uses ChatStream when both provider and channel
  support streaming, cancels stream on tool calls, skips PublishOutbound
  when Finalize already delivered the message

Graceful degradation:
- Bots without forum/topics mode: first sendMessageDraft error sets
  failed=true, subsequent Updates become no-ops, Finalize still delivers
  via SendMessage. User sees normal non-streaming behavior.
- Non-streaming providers: type assertion fails, falls back to Chat()
- Config opt-out: streaming.enabled (default true) in telegram config

Closes #1098
This commit is contained in:
Amir Mamaghani 2026-03-04 18:37:04 +01:00
parent 858e51da62
commit 468723be85
11 changed files with 640 additions and 50 deletions

View file

@ -623,7 +623,7 @@ func (al *AgentLoop) runAgentLoop(
agent.Sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) agent.Sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage)
// 4. Run LLM iteration loop // 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 { if err != nil {
return "", err return "", err
} }
@ -645,8 +645,8 @@ func (al *AgentLoop) runAgentLoop(
al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID)
} }
// 8. Optional: send response via bus // 8. Optional: send response via bus (skip if streaming already delivered it)
if opts.SendResponse { if opts.SendResponse && !streamSent {
al.bus.PublishOutbound(ctx, bus.OutboundMessage{ al.bus.PublishOutbound(ctx, bus.OutboundMessage{
Channel: opts.Channel, Channel: opts.Channel,
ChatID: opts.ChatID, ChatID: opts.ChatID,
@ -664,6 +664,11 @@ func (al *AgentLoop) runAgentLoop(
"final_length": len(finalContent), "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 return finalContent, nil
} }
@ -724,14 +729,25 @@ func (al *AgentLoop) handleReasoning(
} }
// runLLMIteration executes the LLM call loop with tool handling. // 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( func (al *AgentLoop) runLLMIteration(
ctx context.Context, ctx context.Context,
agent *AgentInstance, agent *AgentInstance,
messages []providers.Message, messages []providers.Message,
opts processOptions, opts processOptions,
) (string, int, error) { ) (string, int, bool, error) {
iteration := 0 iteration := 0
var finalContent string 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 { for iteration < agent.MaxIterations {
iteration++ iteration++
@ -771,22 +787,30 @@ func (al *AgentLoop) runLLMIteration(
var response *providers.LLMResponse var response *providers.LLMResponse
var err error var err error
llmOpts := map[string]any{
"max_tokens": agent.MaxTokens,
"temperature": agent.Temperature,
"prompt_cache_key": agent.ID,
}
callLLM := func() (*providers.LLMResponse, error) { 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 { if len(agent.Candidates) > 1 && al.fallback != nil {
fbResult, fbErr := al.fallback.Execute( fbResult, fbErr := al.fallback.Execute(
ctx, ctx,
agent.Candidates, agent.Candidates,
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
return agent.Provider.Chat( return agent.Provider.Chat(
ctx, ctx, messages, providerToolDefs, model, llmOpts,
messages,
providerToolDefs,
model,
map[string]any{
"max_tokens": agent.MaxTokens,
"temperature": agent.Temperature,
"prompt_cache_key": agent.ID,
},
) )
}, },
) )
@ -803,11 +827,7 @@ func (al *AgentLoop) runLLMIteration(
} }
return fbResult.Response, nil return fbResult.Response, nil
} }
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, llmOpts)
"max_tokens": agent.MaxTokens,
"temperature": agent.Temperature,
"prompt_cache_key": agent.ID,
})
} }
// Retry loop for context/token errors // Retry loop for context/token errors
@ -886,7 +906,7 @@ func (al *AgentLoop) runLLMIteration(
"iteration": iteration, "iteration": iteration,
"error": err.Error(), "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( go al.handleReasoning(
@ -909,15 +929,33 @@ func (al *AgentLoop) runLLMIteration(
// Check if no tool calls - we're done // Check if no tool calls - we're done
if len(response.ToolCalls) == 0 { if len(response.ToolCalls) == 0 {
finalContent = response.Content 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)", logger.InfoCF("agent", "LLM response without tool calls (direct answer)",
map[string]any{ map[string]any{
"agent_id": agent.ID, "agent_id": agent.ID,
"iteration": iteration, "iteration": iteration,
"content_chars": len(finalContent), "content_chars": len(finalContent),
"streamed": streamed,
}) })
break 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)) normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls))
for _, tc := range response.ToolCalls { for _, tc := range response.ToolCalls {
normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc))
@ -1073,7 +1111,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. // updateToolContexts updates the context for tools that need channel/chatID info.

View file

@ -13,12 +13,29 @@ var ErrBusClosed = errors.New("message bus closed")
const defaultBusBufferSize = 64 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 { type MessageBus struct {
inbound chan InboundMessage inbound chan InboundMessage
outbound chan OutboundMessage outbound chan OutboundMessage
outboundMedia chan OutboundMediaMessage outboundMedia chan OutboundMediaMessage
done chan struct{} done chan struct{}
closed atomic.Bool closed atomic.Bool
streamDelegate atomic.Value // stores StreamDelegate
} }
func NewMessageBus() *MessageBus { 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() { func (mb *MessageBus) Close() {
if mb.closed.CompareAndSwap(false, true) { if mb.closed.CompareAndSwap(false, true) {
close(mb.done) close(mb.done)

View file

@ -271,20 +271,24 @@ func (c *BaseChannel) HandleMessage(
// Auto-trigger typing indicator, message reaction, and placeholder before publishing. // Auto-trigger typing indicator, message reaction, and placeholder before publishing.
// Each capability is independent — all three may fire for the same message. // 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 { if c.owner != nil && c.placeholderRecorder != nil {
// Typing — independent pipeline // Typing
if tc, ok := c.owner.(TypingCapable); ok { if tc, ok := c.owner.(TypingCapable); ok {
if stop, err := tc.StartTyping(ctx, chatID); err == nil { if stop, err := tc.StartTyping(ctx, chatID); err == nil {
c.placeholderRecorder.RecordTypingStop(c.name, chatID, stop) c.placeholderRecorder.RecordTypingStop(c.name, chatID, stop)
} }
} }
// Reaction — independent pipeline // Reaction
if rc, ok := c.owner.(ReactionCapable); ok && messageID != "" { if rc, ok := c.owner.(ReactionCapable); ok && messageID != "" {
if undo, err := rc.ReactToMessage(ctx, chatID, messageID); err == nil { if undo, err := rc.ReactToMessage(ctx, chatID, messageID); err == nil {
c.placeholderRecorder.RecordReactionUndo(c.name, chatID, undo) c.placeholderRecorder.RecordReactionUndo(c.name, chatID, undo)
} }
} }
// Placeholder — independent pipeline // Placeholder
if pc, ok := c.owner.(PlaceholderCapable); ok { if pc, ok := c.owner.(PlaceholderCapable); ok {
if phID, err := pc.SendPlaceholder(ctx, chatID); err == nil && phID != "" { if phID, err := pc.SendPlaceholder(ctx, chatID); err == nil && phID != "" {
c.placeholderRecorder.RecordPlaceholder(c.name, chatID, phID) c.placeholderRecorder.RecordPlaceholder(c.name, chatID, phID)

View file

@ -31,6 +31,24 @@ type PlaceholderCapable interface {
SendPlaceholder(ctx context.Context, chatID string) (messageID string, err error) 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. // PlaceholderRecorder is injected into channels by Manager.
// Channels call these methods on inbound to register typing/placeholder state. // Channels call these methods on inbound to register typing/placeholder state.
// Manager uses the registered state on outbound to stop typing and edit placeholders. // Manager uses the registered state on outbound to stop typing and edit placeholders.

View file

@ -86,6 +86,7 @@ type Manager struct {
placeholders sync.Map // "channel:chatID" → placeholderID (string) placeholders sync.Map // "channel:chatID" → placeholderID (string)
typingStops sync.Map // "channel:chatID" → func() typingStops sync.Map // "channel:chatID" → func()
reactionUndos sync.Map // "channel:chatID" → reactionEntry reactionUndos sync.Map // "channel:chatID" → reactionEntry
streamActive sync.Map // "channel:chatID" → true (set when streamer.Finalize sent the message)
} }
type asyncTask struct { 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. // 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 { func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMessage, ch Channel) bool {
key := name + ":" + msg.ChatID 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 v, loaded := m.placeholders.LoadAndDelete(key); loaded {
if entry, ok := v.(placeholderEntry); ok && entry.id != "" { if entry, ok := v.(placeholderEntry); ok && entry.id != "" {
if editor, ok := ch.(MessageEditor); ok { if editor, ok := ch.(MessageEditor); ok {
@ -156,6 +164,9 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi
mediaStore: store, mediaStore: store,
} }
// Register as streaming delegate so the agent loop can obtain streamers
messageBus.SetStreamDelegate(m)
if err := m.initChannels(); err != nil { if err := m.initChannels(); err != nil {
return nil, err return nil, err
} }
@ -163,6 +174,59 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi
return m, nil 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. // initChannel is a helper that looks up a factory by name and creates the channel.
func (m *Manager) initChannel(name, displayName string) { func (m *Manager) initChannel(name, displayName string) {
f, ok := getFactory(name) f, ok := getFactory(name)

View file

@ -8,8 +8,10 @@ import (
"os" "os"
"regexp" "regexp"
"slices" "slices"
"math/rand"
"strconv" "strconv"
"strings" "strings"
"sync"
"time" "time"
"github.com/mymmrac/telego" "github.com/mymmrac/telego"
@ -767,3 +769,95 @@ func (c *TelegramChannel) stripBotMention(content string) string {
content = re.ReplaceAllString(content, "") content = re.ReplaceAllString(content, "")
return strings.TrimSpace(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.
}

View file

@ -237,6 +237,10 @@ type PlaceholderConfig struct {
Text string `json:"text,omitempty"` Text string `json:"text,omitempty"`
} }
type StreamingConfig struct {
Enabled bool `json:"enabled,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED"`
}
type WhatsAppConfig struct { type WhatsAppConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WHATSAPP_ENABLED"` Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WHATSAPP_ENABLED"`
BridgeURL string `json:"bridge_url" env:"PICOCLAW_CHANNELS_WHATSAPP_BRIDGE_URL"` BridgeURL string `json:"bridge_url" env:"PICOCLAW_CHANNELS_WHATSAPP_BRIDGE_URL"`
@ -255,6 +259,7 @@ type TelegramConfig struct {
GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"`
Typing TypingConfig `json:"typing,omitempty"` Typing TypingConfig `json:"typing,omitempty"`
Placeholder PlaceholderConfig `json:"placeholder,omitempty"` Placeholder PlaceholderConfig `json:"placeholder,omitempty"`
Streaming StreamingConfig `json:"streaming,omitempty"`
ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_TELEGRAM_REASONING_CHANNEL_ID"` ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_TELEGRAM_REASONING_CHANNEL_ID"`
} }

View file

@ -9,6 +9,7 @@ import (
"github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/option" "github.com/anthropics/anthropic-sdk-go/option"
"github.com/anthropics/anthropic-sdk-go/packages/ssestream"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
) )
@ -93,6 +94,136 @@ func (p *Provider) Chat(
return parseResponse(resp), nil 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 { func (p *Provider) GetDefaultModel() string {
return "claude-sonnet-4.6" return "claude-sonnet-4.6"
} }

View file

@ -52,6 +52,19 @@ func (p *HTTPProvider) Chat(
return p.delegate.Chat(ctx, messages, tools, model, options) 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 { func (p *HTTPProvider) GetDefaultModel() string {
return "" return ""
} }

View file

@ -1,6 +1,7 @@
package openai_compat package openai_compat
import ( import (
"bufio"
"bytes" "bytes"
"context" "context"
"encoding/json" "encoding/json"
@ -101,17 +102,8 @@ func NewProviderWithMaxTokensFieldAndTimeout(
) )
} }
func (p *Provider) Chat( // buildRequestBody constructs the common request body for Chat and ChatStream.
ctx context.Context, func (p *Provider) buildRequestBody(messages []Message, tools []ToolDefinition, model string, options map[string]any) map[string]any {
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (*LLMResponse, error) {
if p.apiBase == "" {
return nil, fmt.Errorf("API base not configured")
}
model = normalizeModel(model, p.apiBase) model = normalizeModel(model, p.apiBase)
requestBody := map[string]any{ requestBody := map[string]any{
@ -125,10 +117,8 @@ func (p *Provider) Chat(
} }
if maxTokens, ok := asInt(options["max_tokens"]); ok { if maxTokens, ok := asInt(options["max_tokens"]); ok {
// Use configured maxTokensField if specified, otherwise fallback to model-based detection
fieldName := p.maxTokensField fieldName := p.maxTokensField
if fieldName == "" { if fieldName == "" {
// Fallback: detect from model name for backward compatibility
lowerModel := strings.ToLower(model) lowerModel := strings.ToLower(model)
if strings.Contains(lowerModel, "glm") || strings.Contains(lowerModel, "o1") || if strings.Contains(lowerModel, "glm") || strings.Contains(lowerModel, "o1") ||
strings.Contains(lowerModel, "gpt-5") { strings.Contains(lowerModel, "gpt-5") {
@ -142,7 +132,6 @@ func (p *Provider) Chat(
if temperature, ok := asFloat(options["temperature"]); ok { if temperature, ok := asFloat(options["temperature"]); ok {
lowerModel := strings.ToLower(model) lowerModel := strings.ToLower(model)
// Kimi k2 models only support temperature=1.
if strings.Contains(lowerModel, "kimi") && strings.Contains(lowerModel, "k2") { if strings.Contains(lowerModel, "kimi") && strings.Contains(lowerModel, "k2") {
requestBody["temperature"] = 1.0 requestBody["temperature"] = 1.0
} else { } 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 cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" {
if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") { if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") {
requestBody["prompt_cache_key"] = cacheKey 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) jsonData, err := json.Marshal(requestBody)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err) return nil, fmt.Errorf("failed to marshal request: %w", err)
@ -195,6 +194,186 @@ func (p *Provider) Chat(
return parseResponse(body) 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) { func parseResponse(body []byte) (*LLMResponse, error) {
var apiResponse struct { var apiResponse struct {
Choices []struct { Choices []struct {

View file

@ -37,6 +37,20 @@ type StatefulProvider interface {
Close() 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. // FailoverReason classifies why an LLM request failed for fallback decisions.
type FailoverReason string type FailoverReason string