From 838eef642aca5433d7c64c53bc294de9987d90b0 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Sun, 12 Jul 2026 16:14:01 -0700 Subject: [PATCH] fix(openai): make streaming actually work; honor configured model The Stream() path was broken end to end: - requests always went out with "stream": false, so the SSE parser found no data lines and every stream ended empty - Config.Model was discarded at construction, and the agent loop never sets req.Model, so requests carried an empty model (hard API error) - tool-call deltas were ignored entirely: the agent never executed tools over a stream with this provider (which also backs the ollama type) - usage was neither requested nor parsed, so token tracking stayed at 0 Now mirrors the proven llamacpp client: stream flag + stream_options .include_usage, per-index tool-call fragment accumulation flushed on finish_reason, usage passthrough, a 4MB SSE scanner buffer (64KB default kills the stream on large tool arguments), and an empty-choices guard in toResponse instead of a panic. Co-Authored-By: Claude Fable 5 --- pkg/llm/providers/openai/client.go | 219 ++++++++++++++++++++++------- 1 file changed, 166 insertions(+), 53 deletions(-) diff --git a/pkg/llm/providers/openai/client.go b/pkg/llm/providers/openai/client.go index 291e179..fbc0079 100644 --- a/pkg/llm/providers/openai/client.go +++ b/pkg/llm/providers/openai/client.go @@ -2,12 +2,13 @@ package openai import ( "bufio" + "bytes" "context" "encoding/json" "fmt" "io" - "net/http" "iter" + "net/http" "strings" "github.com/VictorVargas/rony-llm-agent/pkg/llm" @@ -15,14 +16,15 @@ import ( // Config holds the settings needed to create an OpenAI client. type Config struct { - APIKey string - Model string - BaseURL string // defaults to https://api.openai.com/v1 + APIKey string + Model string + BaseURL string // defaults to https://api.openai.com/v1 } // Client implements llm.LLMClient for OpenAI. type Client struct { apiKey string + model string baseURL string http *http.Client } @@ -40,6 +42,7 @@ func New(cfg Config) (*Client, error) { return &Client{ apiKey: cfg.APIKey, + model: cfg.Model, baseURL: baseURL, http: http.DefaultClient, }, nil @@ -48,7 +51,7 @@ func New(cfg Config) (*Client, error) { func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { endpoint := c.baseURL + "/chat/completions" - payload, err := c.buildRequest(req) + payload, err := c.buildRequest(req, false) if err != nil { return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err) } @@ -76,14 +79,14 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err) } - return c.toResponse(apiResp), nil + return c.toResponse(apiResp) } func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { return func(yield func(llm.StreamChunk, error) bool) { endpoint := c.baseURL + "/chat/completions" - payload, err := c.buildRequest(req) + payload, err := c.buildRequest(req, true) if err != nil { yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err)) return @@ -111,7 +114,45 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq return } + // toolCallAccum buffers one tool call's fragments as they stream in: + // the SSE format sends the id/name in the first delta for a given + // tool-call index and the (potentially large) arguments JSON in + // pieces across many subsequent deltas, so it can't be handed to a + // tool handler until it's fully assembled. Same accumulation the + // llamacpp client does — without it, tool calls made over a stream + // were silently dropped and the agent loop never executed them. + type toolCallAccum struct { + id string + name string + args strings.Builder + } + toolCallFrags := map[int]*toolCallAccum{} + var toolCallOrder []int + + flushToolCalls := func() []llm.ToolCall { + if len(toolCallOrder) == 0 { + return nil + } + calls := make([]llm.ToolCall, 0, len(toolCallOrder)) + for _, idx := range toolCallOrder { + frag := toolCallFrags[idx] + calls = append(calls, llm.ToolCall{ + ID: frag.id, + Name: frag.name, + Arguments: json.RawMessage(frag.args.String()), + }) + } + toolCallFrags = map[int]*toolCallAccum{} + toolCallOrder = nil + return calls + } + scanner := bufio.NewScanner(resp.Body) + // A single SSE line can exceed bufio.Scanner's 64KB default cap + // (e.g. a large tool-call arguments delta), which would kill the + // stream with "token too long" — same headroom the anthropic + // client already reserves. + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) for scanner.Scan() { line := scanner.Text() if !strings.HasPrefix(line, "data: ") { @@ -128,13 +169,61 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq return } + var usage llm.TokenUsage + if event.Usage != nil { + usage = llm.TokenUsage{ + InputTokens: event.Usage.PromptTokens, + OutputTokens: event.Usage.CompletionTokens, + TotalTokens: event.Usage.TotalTokens, + } + } + + if len(event.Choices) == 0 { + // The usage-only event (per stream_options.include_usage) + // carries no choices, so it needs its own chunk. + if event.Usage != nil { + if !yield(llm.StreamChunk{Usage: usage}, nil) { + return + } + } + continue + } + for _, choice := range event.Choices { + hasFragment := len(choice.Delta.ToolCalls) > 0 + for _, tc := range choice.Delta.ToolCalls { + frag, ok := toolCallFrags[tc.Index] + if !ok { + frag = &toolCallAccum{} + toolCallFrags[tc.Index] = frag + toolCallOrder = append(toolCallOrder, tc.Index) + } + if tc.ID != "" { + frag.id = tc.ID + } + if tc.Function.Name != "" { + frag.name = tc.Function.Name + } + frag.args.WriteString(tc.Function.Arguments) + } + chunk := llm.StreamChunk{ Delta: choice.Delta.Content, + Usage: usage, } if choice.FinishReason != "" { chunk.FinishReason = choice.FinishReason + chunk.ToolCalls = flushToolCalls() } + + // A fragment-only event (a piece of a tool call's streamed + // arguments, with nothing else in this delta) has nothing + // yet for the agent loop to act on: it was buffered above, + // so skip yielding an empty chunk for it. + if hasFragment && chunk.Delta == "" && chunk.FinishReason == "" { + continue + } + if !yield(chunk, nil) { return } @@ -160,7 +249,7 @@ func (c *Client) Capabilities() llm.ProviderCapabilities { } // buildRequest converts an llm.CompletionRequest to the OpenAI API format. -func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { +func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) { // Convert messages to OpenAI format messages := make([]openaiMessage, len(req.Messages)) for i, m := range req.Messages { @@ -195,10 +284,25 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { tools[i] = tool } + // The per-request model wins when set; otherwise fall back to the + // client's configured one (Config.Model used to be discarded entirely, + // so every request went out with an empty model — a hard API error on + // OpenAI, and the agent loop never sets req.Model). + model := req.Model + if model == "" { + model = c.model + } + openaiReq := openaiChatRequest{ - Model: req.Model, + Model: model, Messages: messages, - Stream: false, + Stream: stream, + } + if stream { + // Ask for a final SSE event carrying token usage (OpenAI-style + // streaming omits it otherwise), so callers can track real token + // counts per turn instead of always seeing zero. + openaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true} } if len(tools) > 0 { openaiReq.Tools = tools @@ -222,11 +326,14 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { if err != nil { return nil, fmt.Errorf("marshaling request: %w", err) } - return strings.NewReader(string(data)), nil + return bytes.NewReader(data), nil } // toResponse converts an OpenAI API response to our CompletionResponse. -func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse { +func (c *Client) toResponse(resp openaiChatResponse) (llm.CompletionResponse, error) { + if len(resp.Choices) == 0 { + return llm.CompletionResponse{}, fmt.Errorf("openai: response contained no choices") + } choice := resp.Choices[0] result := llm.CompletionResponse{ ID: resp.ID, @@ -249,63 +356,68 @@ func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse { TotalTokens: resp.Usage.TotalTokens, } - return result + return result, nil } // OpenAI API types type openaiChatRequest struct { - Model string `json:"model"` - Messages []openaiMessage `json:"messages"` - Tools []openaiTool `json:"tools,omitempty"` - ToolChoice interface{} `json:"tool_choice,omitempty"` - Temperature *float32 `json:"temperature,omitempty"` - MaxTokens *int `json:"max_tokens,omitempty"` - Stop []string `json:"stop,omitempty"` - Stream bool `json:"stream"` + Model string `json:"model"` + Messages []openaiMessage `json:"messages"` + Tools []openaiTool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Temperature *float32 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + Stop []string `json:"stop,omitempty"` + Stream bool `json:"stream"` + StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"` +} + +type openaiStreamOptions struct { + IncludeUsage bool `json:"include_usage"` } type openaiMessage struct { - Role string `json:"role"` - Content string `json:"content"` - ToolCallID string `json:"tool_call_id,omitempty"` - Name string `json:"name,omitempty"` - ToolCalls []openaiToolCall `json:"tool_calls,omitempty"` + Role string `json:"role"` + Content string `json:"content"` + ToolCallID string `json:"tool_call_id,omitempty"` + Name string `json:"name,omitempty"` + ToolCalls []openaiToolCall `json:"tool_calls,omitempty"` } type openaiTool struct { - Type string `json:"type"` - Function json.RawMessage `json:"function"` + Type string `json:"type"` + Function json.RawMessage `json:"function"` } type openaiChatResponse struct { - ID string `json:"id"` - Model string `json:"model"` - Choices []openaiChoice `json:"choices"` - Usage openaiUsage `json:"usage"` + ID string `json:"id"` + Model string `json:"model"` + Choices []openaiChoice `json:"choices"` + Usage openaiUsage `json:"usage"` } type openaiChoice struct { Index int `json:"index"` Message openaiMessageResult `json:"message"` - FinishReason string `json:"finish_reason"` + FinishReason string `json:"finish_reason"` } type openaiMessageResult struct { - Role string `json:"role"` - Content string `json:"content"` - ToolCalls []openaiToolCall `json:"tool_calls"` + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []openaiToolCall `json:"tool_calls"` } type openaiToolCall struct { - ID string `json:"id"` - Type string `json:"type"` - Function openaiFunction `json:"function"` + ID string `json:"id"` + Type string `json:"type"` + Function openaiFunction `json:"function"` } type openaiFunction struct { - Name string `json:"name"` - Arguments string `json:"arguments"` + Name string `json:"name"` + Arguments string `json:"arguments"` } type openaiUsage struct { @@ -317,30 +429,31 @@ type openaiUsage struct { // Stream event types type openaiStreamEvent struct { - ID string `json:"id"` - Choices []openaiStreamChoice `json:"choices"` + ID string `json:"id"` + Choices []openaiStreamChoice `json:"choices"` + Usage *openaiUsage `json:"usage"` } type openaiStreamChoice struct { - Index int `json:"index"` - Delta openaiStreamDelta `json:"delta"` - FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Delta openaiStreamDelta `json:"delta"` + FinishReason string `json:"finish_reason"` } type openaiStreamDelta struct { - Content string `json:"content"` - Role string `json:"role"` + Content string `json:"content"` + Role string `json:"role"` ToolCalls []openaiStreamToolCall `json:"tool_calls"` } type openaiStreamToolCall struct { - Index int `json:"index"` - ID string `json:"id"` - Type string `json:"type"` + Index int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` Function openaiStreamFunction `json:"function"` } type openaiStreamFunction struct { - Name string `json:"name"` - Arguments string `json:"arguments"` + Name string `json:"name"` + Arguments string `json:"arguments"` }