diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index e308e1cca..88d229ac1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 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, diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index fe66b6d94..567cd0aea 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -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") + } +} diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go index 46f921efd..7572a3837 100644 --- a/pkg/providers/http_provider.go +++ b/pkg/providers/http_provider.go @@ -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) +} diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 17ad2f3ed..2c49e1044 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -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 { diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 64927d85c..7340398e9 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -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") + } +} diff --git a/pkg/providers/protocoltypes/types.go b/pkg/providers/protocoltypes/types.go index 3a089ca47..65d7a2c05 100644 --- a/pkg/providers/protocoltypes/types.go +++ b/pkg/providers/protocoltypes/types.go @@ -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) +} diff --git a/pkg/providers/types.go b/pkg/providers/types.go index f711e7803..bdce83765 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -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