feat: add streaming pipeline with early repetition detection

Introduce StreamingProvider interface and buffered-channel SSE pipeline
so that repetition loops (e.g. 73K-char <think> incidents) are detected
mid-stream, cancelling the HTTP request early to save tokens and time.

- Add StreamEvent/StreamToolCallDelta types (protocoltypes)
- Add StreamingProvider interface (opt-in via type assertion)
- Refactor openai_compat: extract buildHTTPRequest, add ChatStream,
  readSSEIntoChannel, AccumulateStream, CanStream
- Forward ChatStream/CanStream through HTTPProvider
- Add consumeStreamWithRepetitionDetection in agent loop with periodic
  n-gram check every 1000 runes during streaming
- Non-streaming providers are completely unaffected

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
dj-oyu 2026-02-22 15:58:39 +09:00
parent 32d0147a08
commit 0545aa260f
7 changed files with 778 additions and 16 deletions

View file

@ -26,6 +26,7 @@ import (
"github.com/sipeed/picoclaw/pkg/constants"
"github.com/sipeed/picoclaw/pkg/logger"
"github.com/sipeed/picoclaw/pkg/providers"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
"github.com/sipeed/picoclaw/pkg/routing"
"github.com/sipeed/picoclaw/pkg/session"
"github.com/sipeed/picoclaw/pkg/skills"
@ -1377,6 +1378,97 @@ func buildRichStatus(task *activeTask, isBackground bool, workspace string) stri
}
// runLLMIteration executes the LLM call loop with tool handling.
// consumeStreamWithRepetitionDetection reads StreamEvents from ch, accumulates
// content and tool calls, and runs repetition detection every checkInterval runes.
// If repetition is detected, cancelFn is called to abort the HTTP request and
// the function returns the partial response with detected=true.
func consumeStreamWithRepetitionDetection(
ch <-chan protocoltypes.StreamEvent,
cancelFn context.CancelFunc,
checkInterval int,
) (*providers.LLMResponse, bool, error) {
var content strings.Builder
var toolCalls []streamToolCallAcc
var finishReason string
var usage *providers.UsageInfo
runesSinceLastCheck := 0
for ev := range ch {
if ev.Err != nil {
return nil, false, ev.Err
}
if ev.ContentDelta != "" {
content.WriteString(ev.ContentDelta)
runesSinceLastCheck += utf8.RuneCountInString(ev.ContentDelta)
}
if ev.FinishReason != "" {
finishReason = ev.FinishReason
}
if ev.Usage != nil {
usage = ev.Usage
}
for _, tc := range ev.ToolCallDeltas {
for len(toolCalls) <= tc.Index {
toolCalls = append(toolCalls, streamToolCallAcc{})
}
if tc.ID != "" {
toolCalls[tc.Index].id = tc.ID
}
if tc.Name != "" {
toolCalls[tc.Index].name = tc.Name
}
toolCalls[tc.Index].args.WriteString(tc.ArgumentsDelta)
}
// Run repetition detection periodically on accumulated content.
if runesSinceLastCheck >= checkInterval && content.Len() > 2000 {
runesSinceLastCheck = 0
if utils.DetectRepetitionLoop(content.String()) {
cancelFn()
// Drain remaining events so the producer goroutine can exit.
for range ch {
}
resp := buildAccumulatedResponse(content.String(), toolCalls, finishReason, usage)
return resp, true, nil
}
}
}
resp := buildAccumulatedResponse(content.String(), toolCalls, finishReason, usage)
return resp, false, nil
}
// streamToolCallAcc accumulates streamed tool call fragments.
type streamToolCallAcc struct {
id string
name string
args strings.Builder
}
// buildAccumulatedResponse constructs an LLMResponse from accumulated stream data.
func buildAccumulatedResponse(content string, toolCalls []streamToolCallAcc, finishReason string, usage *providers.UsageInfo) *providers.LLMResponse {
resp := &providers.LLMResponse{
Content: content,
FinishReason: finishReason,
Usage: usage,
}
for _, tc := range toolCalls {
arguments := make(map[string]any)
argStr := tc.args.String()
if argStr != "" {
if err := json.Unmarshal([]byte(argStr), &arguments); err != nil {
arguments["raw"] = argStr
}
}
resp.ToolCalls = append(resp.ToolCalls, providers.ToolCall{
ID: tc.id,
Name: tc.name,
Arguments: arguments,
})
}
return resp
}
func (al *AgentLoop) runLLMIteration(
ctx context.Context,
agent *AgentInstance,
@ -1459,15 +1551,38 @@ func (al *AgentLoop) runLLMIteration(
var response *providers.LLMResponse
var err error
// doCall invokes a single LLM provider, using streaming with
// early repetition detection when the provider supports it.
opts_ := map[string]any{
"max_tokens": agent.MaxTokens,
"temperature": agent.Temperature,
}
doCall := func(ctx context.Context, p providers.LLMProvider, model string) (*providers.LLMResponse, error) {
if sp, ok := p.(providers.StreamingProvider); ok && sp.CanStream() {
streamCtx, streamCancel := context.WithCancel(ctx)
defer streamCancel()
ch, err := sp.ChatStream(streamCtx, messages, providerToolDefs, model, opts_)
if err != nil {
return nil, err
}
resp, repetition, err := consumeStreamWithRepetitionDetection(ch, streamCancel, 1000)
if err != nil {
return nil, err
}
if repetition {
resp.FinishReason = "repetition_detected"
}
return resp, nil
}
return p.Chat(ctx, messages, providerToolDefs, model, opts_)
}
callLLM := func() (*providers.LLMResponse, error) {
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) {
p := al.resolveProvider(provider, model, agent.Provider)
return p.Chat(ctx, messages, providerToolDefs, model, map[string]interface{}{
"max_tokens": agent.MaxTokens,
"temperature": agent.Temperature,
})
return doCall(ctx, p, model)
},
)
if fbErr != nil {
@ -1480,10 +1595,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,
})
return doCall(ctx, agent.Provider, agent.Model)
}
// Retry loop for context/token errors
@ -1548,7 +1660,10 @@ func (al *AgentLoop) runLLMIteration(
// Detect repetition loop on raw text (before stripping think
// blocks so loops inside <think> are caught). Skip when the
// provider already returned native tool calls.
if len(response.ToolCalls) == 0 && utils.DetectRepetitionLoop(response.Content) {
// Streaming providers may have already flagged repetition via
// FinishReason="repetition_detected" — honour that too.
if response.FinishReason == "repetition_detected" ||
(len(response.ToolCalls) == 0 && utils.DetectRepetitionLoop(response.Content)) {
logger.WarnCF("agent", "Repetition loop detected in LLM response, retrying",
map[string]any{
"agent_id": agent.ID,

View file

@ -12,6 +12,7 @@ import (
"github.com/sipeed/picoclaw/pkg/bus"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/providers"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
"github.com/sipeed/picoclaw/pkg/tools"
)
@ -2116,3 +2117,142 @@ func (m *nudgeCaptureMockProvider) Chat(
func (m *nudgeCaptureMockProvider) GetDefaultModel() string {
return "mock-nudge-model"
}
// --- consumeStreamWithRepetitionDetection tests ---
func TestConsumeStream_NormalCompletion(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 8)
go func() {
ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "}
ch <- protocoltypes.StreamEvent{ContentDelta: "world!"}
ch <- protocoltypes.StreamEvent{
FinishReason: "stop",
Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7},
}
close(ch)
}()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if detected {
t.Fatal("expected detected=false for normal content")
}
if resp.Content != "Hello world!" {
t.Errorf("Content = %q, want %q", resp.Content, "Hello world!")
}
if resp.FinishReason != "stop" {
t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop")
}
if resp.Usage == nil || resp.Usage.TotalTokens != 7 {
t.Errorf("Usage.TotalTokens = %v, want 7", resp.Usage)
}
_ = ctx // keep linter happy
}
func TestConsumeStream_DetectsRepetition(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 64)
cancelCalled := false
ctx, cancel := context.WithCancel(context.Background())
wrappedCancel := func() {
cancelCalled = true
cancel()
}
// Send enough repetitive content to trigger detection.
// The pattern "abcdefghij" repeated many times will have very low n-gram uniqueness.
repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk
go func() {
// Send 6 chunks of repetitive content = 3000 chars total,
// each with 500 runes. The check triggers after every 1000 runes
// when content > 2000 chars.
for i := 0; i < 6; i++ {
ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk}
}
// Send more data that should be ignored after detection.
for i := 0; i < 10; i++ {
ch <- protocoltypes.StreamEvent{ContentDelta: "more data"}
}
close(ch)
}()
resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if !detected {
t.Fatal("expected repetition detection to trigger")
}
if !cancelCalled {
t.Error("expected cancelFn to be called")
}
// The response should be shorter than the full 3000+ chars
// because detection triggers early.
if len(resp.Content) >= 3000+10*len("more data") {
t.Errorf("Content length = %d, expected less than full output", len(resp.Content))
}
_ = ctx
}
func TestConsumeStream_ToolCallAccumulation(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 8)
go func() {
ch <- protocoltypes.StreamEvent{
ToolCallDeltas: []protocoltypes.StreamToolCallDelta{
{Index: 0, ID: "call_1", Name: "test_fn", ArgumentsDelta: `{"ke`},
},
}
ch <- protocoltypes.StreamEvent{
ToolCallDeltas: []protocoltypes.StreamToolCallDelta{
{Index: 0, ArgumentsDelta: `y":"val"}`},
},
}
ch <- protocoltypes.StreamEvent{FinishReason: "tool_calls"}
close(ch)
}()
_, cancel := context.WithCancel(context.Background())
defer cancel()
resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if detected {
t.Fatal("expected no repetition detection for tool calls")
}
if len(resp.ToolCalls) != 1 {
t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls))
}
if resp.ToolCalls[0].Name != "test_fn" {
t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_fn")
}
if resp.ToolCalls[0].Arguments["key"] != "val" {
t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "val")
}
}
func TestConsumeStream_StreamError(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 4)
go func() {
ch <- protocoltypes.StreamEvent{ContentDelta: "partial"}
ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("read error")}
close(ch)
}()
_, cancel := context.WithCancel(context.Background())
defer cancel()
_, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "read error") {
t.Errorf("error = %q, want to contain %q", err.Error(), "read error")
}
}

View file

@ -54,3 +54,19 @@ func (p *HTTPProvider) Chat(ctx context.Context, messages []Message, tools []Too
func (p *HTTPProvider) GetDefaultModel() string {
return ""
}
// CanStream returns true when the underlying provider uses SSE streaming.
func (p *HTTPProvider) CanStream() bool {
return p.delegate.CanStream()
}
// ChatStream opens an SSE stream and returns a channel of StreamEvent.
func (p *HTTPProvider) ChatStream(
ctx context.Context,
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (<-chan StreamEvent, error) {
return p.delegate.ChatStream(ctx, messages, tools, model, options)
}

View file

@ -89,13 +89,18 @@ func NewProviderWithOptions(apiKey, apiBase, proxy string, opts Options) *Provid
}
}
func (p *Provider) Chat(
// streamBufferSize is the channel buffer size for ChatStream events.
const streamBufferSize = 32
// buildHTTPRequest constructs a ready-to-send *http.Request for the chat API.
func (p *Provider) buildHTTPRequest(
ctx context.Context,
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (*LLMResponse, error) {
stream bool,
) (*http.Request, error) {
if p.apiBase == "" {
return nil, fmt.Errorf("API base not configured")
}
@ -138,7 +143,7 @@ func (p *Provider) Chat(
}
}
if p.stream {
if stream {
requestBody["stream"] = true
}
@ -157,6 +162,31 @@ func (p *Provider) Chat(
req.Header.Set("Authorization", "Bearer "+p.apiKey)
}
return req, nil
}
func (p *Provider) Chat(
ctx context.Context,
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (*LLMResponse, error) {
// When streaming is enabled, delegate to ChatStream + AccumulateStream
// so that the SSE→channel path is always exercised.
if p.stream {
ch, err := p.ChatStream(ctx, messages, tools, model, options)
if err != nil {
return nil, err
}
return AccumulateStream(ch)
}
req, err := p.buildHTTPRequest(ctx, messages, tools, model, options, false)
if err != nil {
return nil, err
}
resp, err := p.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
@ -168,10 +198,6 @@ func (p *Provider) Chat(
return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body))
}
if p.stream {
return parseStreamResponse(resp.Body)
}
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
@ -180,6 +206,174 @@ func (p *Provider) Chat(
return parseResponse(body)
}
// CanStream returns true when this provider is configured for SSE streaming.
func (p *Provider) CanStream() bool {
return p.stream
}
// ChatStream opens an SSE connection and returns a channel of StreamEvent.
// The channel is closed when the stream ends or an error occurs.
// Cancelling ctx will abort the HTTP request and close the channel.
func (p *Provider) ChatStream(
ctx context.Context,
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (<-chan protocoltypes.StreamEvent, error) {
req, err := p.buildHTTPRequest(ctx, messages, tools, model, options, true)
if err != nil {
return nil, err
}
resp, err := p.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("API request failed:\n Status: %d\n Body: %s", resp.StatusCode, string(body))
}
ch := make(chan protocoltypes.StreamEvent, streamBufferSize)
go func() {
defer resp.Body.Close()
defer close(ch)
readSSEIntoChannel(ctx, resp.Body, ch)
}()
return ch, nil
}
// readSSEIntoChannel reads SSE lines from r and sends StreamEvent values on ch.
// It returns when the stream ends, an error occurs, or ctx is cancelled.
func readSSEIntoChannel(ctx context.Context, r io.Reader, ch chan<- protocoltypes.StreamEvent) {
scanner := bufio.NewScanner(r)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
for scanner.Scan() {
// Check for context cancellation between lines.
select {
case <-ctx.Done():
return
default:
}
line := scanner.Text()
if !strings.HasPrefix(line, "data: ") {
continue
}
data := strings.TrimPrefix(line, "data: ")
if data == "[DONE]" {
return
}
var chunk streamChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue // skip malformed chunks
}
ev := protocoltypes.StreamEvent{}
if chunk.Usage != nil {
ev.Usage = chunk.Usage
}
if len(chunk.Choices) > 0 {
choice := chunk.Choices[0]
ev.ContentDelta = choice.Delta.Content
if choice.FinishReason != "" {
ev.FinishReason = choice.FinishReason
}
for _, tc := range choice.Delta.ToolCalls {
delta := protocoltypes.StreamToolCallDelta{
Index: tc.Index,
ID: tc.ID,
}
if tc.Function != nil {
delta.Name = tc.Function.Name
delta.ArgumentsDelta = tc.Function.Arguments
}
ev.ToolCallDeltas = append(ev.ToolCallDeltas, delta)
}
}
select {
case ch <- ev:
case <-ctx.Done():
return
}
}
if err := scanner.Err(); err != nil {
select {
case ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("reading stream: %w", err)}:
case <-ctx.Done():
}
}
}
// AccumulateStream drains a StreamEvent channel and returns a complete LLMResponse.
func AccumulateStream(ch <-chan protocoltypes.StreamEvent) (*LLMResponse, error) {
var content strings.Builder
var toolCalls []streamToolCallAcc
var finishReason string
var usage *UsageInfo
for ev := range ch {
if ev.Err != nil {
return nil, ev.Err
}
if ev.ContentDelta != "" {
content.WriteString(ev.ContentDelta)
}
if ev.FinishReason != "" {
finishReason = ev.FinishReason
}
if ev.Usage != nil {
usage = ev.Usage
}
for _, tc := range ev.ToolCallDeltas {
for len(toolCalls) <= tc.Index {
toolCalls = append(toolCalls, streamToolCallAcc{})
}
if tc.ID != "" {
toolCalls[tc.Index].ID = tc.ID
}
if tc.Name != "" {
toolCalls[tc.Index].Name = tc.Name
}
toolCalls[tc.Index].Arguments.WriteString(tc.ArgumentsDelta)
}
}
result := &LLMResponse{
Content: content.String(),
FinishReason: finishReason,
Usage: usage,
}
for _, tc := range toolCalls {
arguments := make(map[string]any)
argStr := tc.Arguments.String()
if argStr != "" {
if err := json.Unmarshal([]byte(argStr), &arguments); err != nil {
log.Printf("openai_compat: failed to decode streamed tool call arguments for %q: %v", tc.Name, err)
arguments["raw"] = argStr
}
}
result.ToolCalls = append(result.ToolCalls, ToolCall{
ID: tc.ID,
Name: tc.Name,
Arguments: arguments,
})
}
return result, nil
}
func parseResponse(body []byte) (*LLMResponse, error) {
var apiResponse struct {
Choices []struct {

View file

@ -1,12 +1,16 @@
package openai_compat
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
)
func TestProviderChat_UsesMaxCompletionTokensForGLM(t *testing.T) {
@ -405,3 +409,261 @@ func TestProviderChat_CustomEndpointPath(t *testing.T) {
t.Fatalf("endpoint path = %q, want %q", hitPath, "/text/chatcompletion_v2")
}
}
func TestReadSSEIntoChannel_TextAndToolCalls(t *testing.T) {
sseData := strings.Join([]string{
`data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`,
``,
`data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`,
``,
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"greet","arguments":"{\"n"}}]},"finish_reason":""}]}`,
``,
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ame\":\"Bob\"}"}}]},"finish_reason":""}]}`,
``,
`data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`,
``,
`data: [DONE]`,
``,
}, "\n")
ch := make(chan protocoltypes.StreamEvent, 32)
go func() {
defer close(ch)
readSSEIntoChannel(context.Background(), strings.NewReader(sseData), ch)
}()
var events []protocoltypes.StreamEvent
for ev := range ch {
events = append(events, ev)
}
if len(events) < 3 {
t.Fatalf("got %d events, want at least 3", len(events))
}
// Check content deltas
if events[0].ContentDelta != "Hello" {
t.Errorf("events[0].ContentDelta = %q, want %q", events[0].ContentDelta, "Hello")
}
if events[1].ContentDelta != " world" {
t.Errorf("events[1].ContentDelta = %q, want %q", events[1].ContentDelta, " world")
}
// Check tool call deltas
if len(events[2].ToolCallDeltas) != 1 || events[2].ToolCallDeltas[0].ID != "call_1" {
t.Errorf("events[2] should contain tool call with ID=call_1")
}
if events[2].ToolCallDeltas[0].Name != "greet" {
t.Errorf("events[2].ToolCallDeltas[0].Name = %q, want %q", events[2].ToolCallDeltas[0].Name, "greet")
}
// Check finish event
lastEv := events[len(events)-1]
if lastEv.FinishReason != "stop" {
t.Errorf("last event FinishReason = %q, want %q", lastEv.FinishReason, "stop")
}
if lastEv.Usage == nil || lastEv.Usage.TotalTokens != 7 {
t.Errorf("last event Usage.TotalTokens = %v, want 7", lastEv.Usage)
}
}
func TestReadSSEIntoChannel_ContextCancel(t *testing.T) {
// Simulate a slow SSE stream that gets cancelled.
ctx, cancel := context.WithCancel(context.Background())
// Create a reader that blocks after sending one chunk.
sseData := `data: {"choices":[{"delta":{"content":"first"},"finish_reason":""}]}` + "\n\n"
ch := make(chan protocoltypes.StreamEvent, 32)
go func() {
defer close(ch)
readSSEIntoChannel(ctx, strings.NewReader(sseData), ch)
}()
// Read the first event.
ev := <-ch
if ev.ContentDelta != "first" {
t.Fatalf("ContentDelta = %q, want %q", ev.ContentDelta, "first")
}
// Cancel the context; the channel should close.
cancel()
_, ok := <-ch
if ok {
t.Fatal("expected channel to be closed after context cancel")
}
}
func TestAccumulateStream_FullResponse(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 8)
go func() {
ch <- protocoltypes.StreamEvent{ContentDelta: "Hello"}
ch <- protocoltypes.StreamEvent{ContentDelta: " world"}
ch <- protocoltypes.StreamEvent{
ToolCallDeltas: []protocoltypes.StreamToolCallDelta{
{Index: 0, ID: "call_1", Name: "test_tool", ArgumentsDelta: `{"key"`},
},
}
ch <- protocoltypes.StreamEvent{
ToolCallDeltas: []protocoltypes.StreamToolCallDelta{
{Index: 0, ArgumentsDelta: `:"value"}`},
},
}
ch <- protocoltypes.StreamEvent{
FinishReason: "stop",
Usage: &UsageInfo{PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8},
}
close(ch)
}()
resp, err := AccumulateStream(ch)
if err != nil {
t.Fatalf("AccumulateStream() error = %v", err)
}
if resp.Content != "Hello world" {
t.Errorf("Content = %q, want %q", resp.Content, "Hello world")
}
if resp.FinishReason != "stop" {
t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop")
}
if resp.Usage == nil || resp.Usage.TotalTokens != 8 {
t.Errorf("Usage.TotalTokens = %v, want 8", resp.Usage)
}
if len(resp.ToolCalls) != 1 {
t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls))
}
if resp.ToolCalls[0].Name != "test_tool" {
t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_tool")
}
if resp.ToolCalls[0].Arguments["key"] != "value" {
t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "value")
}
}
func TestAccumulateStream_Error(t *testing.T) {
ch := make(chan protocoltypes.StreamEvent, 4)
go func() {
ch <- protocoltypes.StreamEvent{ContentDelta: "partial"}
ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("connection reset")}
close(ch)
}()
_, err := AccumulateStream(ch)
if err == nil {
t.Fatal("expected error, got nil")
}
if !strings.Contains(err.Error(), "connection reset") {
t.Fatalf("error = %q, want to contain %q", err.Error(), "connection reset")
}
}
func TestChatStream_EndToEnd(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
flusher, _ := w.(http.Flusher)
chunks := []string{
`data: {"choices":[{"delta":{"content":"stream"},"finish_reason":""}]}`,
`data: {"choices":[{"delta":{"content":"ed"},"finish_reason":""}]}`,
`data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`,
`data: [DONE]`,
}
for _, c := range chunks {
fmt.Fprintln(w, c)
fmt.Fprintln(w)
if flusher != nil {
flusher.Flush()
}
}
}))
defer server.Close()
p := NewProviderWithOptions("key", server.URL, "", Options{Stream: true})
ch, err := p.ChatStream(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil)
if err != nil {
t.Fatalf("ChatStream() error = %v", err)
}
resp, err := AccumulateStream(ch)
if err != nil {
t.Fatalf("AccumulateStream() error = %v", err)
}
if resp.Content != "streamed" {
t.Errorf("Content = %q, want %q", resp.Content, "streamed")
}
if resp.FinishReason != "stop" {
t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop")
}
if resp.Usage == nil || resp.Usage.TotalTokens != 3 {
t.Errorf("Usage.TotalTokens = %v, want 3", resp.Usage)
}
}
func TestChatStream_EarlyCancel(t *testing.T) {
serverDone := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer close(serverDone)
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
flusher, _ := w.(http.Flusher)
// Send many chunks; expect the client to cancel early.
for i := 0; i < 1000; i++ {
select {
case <-r.Context().Done():
return
default:
}
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"},\"finish_reason\":\"\"}]}\n\n")
if flusher != nil {
flusher.Flush()
}
}
}))
defer server.Close()
p := NewProviderWithOptions("key", server.URL, "", Options{Stream: true})
ctx, cancel := context.WithCancel(context.Background())
ch, err := p.ChatStream(ctx, []Message{{Role: "user", Content: "hi"}}, nil, "test", nil)
if err != nil {
t.Fatalf("ChatStream() error = %v", err)
}
// Read a few events, then cancel.
count := 0
for ev := range ch {
if ev.Err != nil {
break
}
count++
if count >= 5 {
cancel()
}
}
if count < 5 {
t.Errorf("expected at least 5 events before cancel, got %d", count)
}
// Server should have received the cancellation.
<-serverDone
}
func TestCanStream(t *testing.T) {
p1 := NewProvider("key", "https://example.com", "")
if p1.CanStream() {
t.Error("CanStream() = true for non-stream provider")
}
p2 := NewProviderWithOptions("key", "https://example.com", "", Options{Stream: true})
if !p2.CanStream() {
t.Error("CanStream() = false for stream provider")
}
}

View file

@ -54,3 +54,20 @@ type ToolFunctionDefinition struct {
Description string `json:"description"`
Parameters map[string]any `json:"parameters"`
}
// StreamEvent represents a single chunk from an SSE streaming response.
type StreamEvent struct {
ContentDelta string
ToolCallDeltas []StreamToolCallDelta
FinishReason string // set only on the final event
Usage *UsageInfo // set only on the final event
Err error // non-nil when the stream encountered an error
}
// StreamToolCallDelta carries an incremental piece of a streaming tool call.
type StreamToolCallDelta struct {
Index int
ID string // set on the first chunk for this tool call
Name string // set on the first chunk for this tool call
ArgumentsDelta string // JSON fragment (incremental)
}

View file

@ -17,6 +17,8 @@ type (
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
ExtraContent = protocoltypes.ExtraContent
GoogleExtra = protocoltypes.GoogleExtra
StreamEvent = protocoltypes.StreamEvent
StreamToolCallDelta = protocoltypes.StreamToolCallDelta
)
type LLMProvider interface {
@ -67,6 +69,22 @@ func (e *FailoverError) IsRetriable() bool {
return e.Reason != FailoverFormat
}
// StreamingProvider extends LLMProvider with SSE channel-based streaming.
// Use a type assertion to check if a provider supports streaming:
//
// if sp, ok := provider.(StreamingProvider); ok && sp.CanStream() { ... }
type StreamingProvider interface {
LLMProvider
CanStream() bool
ChatStream(
ctx context.Context,
messages []Message,
tools []ToolDefinition,
model string,
options map[string]any,
) (<-chan StreamEvent, error)
}
// ModelConfig holds primary model and fallback list.
type ModelConfig struct {
Primary string