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 <noreply@anthropic.com>
This commit is contained in:
parent
eea51e30d1
commit
838eef642a
1 changed files with 166 additions and 53 deletions
|
|
@ -2,12 +2,13 @@ package openai
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
|
||||||
"iter"
|
"iter"
|
||||||
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
|
@ -23,6 +24,7 @@ type Config struct {
|
||||||
// Client implements llm.LLMClient for OpenAI.
|
// Client implements llm.LLMClient for OpenAI.
|
||||||
type Client struct {
|
type Client struct {
|
||||||
apiKey string
|
apiKey string
|
||||||
|
model string
|
||||||
baseURL string
|
baseURL string
|
||||||
http *http.Client
|
http *http.Client
|
||||||
}
|
}
|
||||||
|
|
@ -40,6 +42,7 @@ func New(cfg Config) (*Client, error) {
|
||||||
|
|
||||||
return &Client{
|
return &Client{
|
||||||
apiKey: cfg.APIKey,
|
apiKey: cfg.APIKey,
|
||||||
|
model: cfg.Model,
|
||||||
baseURL: baseURL,
|
baseURL: baseURL,
|
||||||
http: http.DefaultClient,
|
http: http.DefaultClient,
|
||||||
}, nil
|
}, nil
|
||||||
|
|
@ -48,7 +51,7 @@ func New(cfg Config) (*Client, error) {
|
||||||
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
endpoint := c.baseURL + "/chat/completions"
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
payload, err := c.buildRequest(req)
|
payload, err := c.buildRequest(req, false)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
|
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 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] {
|
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
return func(yield func(llm.StreamChunk, error) bool) {
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
endpoint := c.baseURL + "/chat/completions"
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
payload, err := c.buildRequest(req)
|
payload, err := c.buildRequest(req, true)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
|
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
|
||||||
return
|
return
|
||||||
|
|
@ -111,7 +114,45 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
||||||
return
|
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)
|
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() {
|
for scanner.Scan() {
|
||||||
line := scanner.Text()
|
line := scanner.Text()
|
||||||
if !strings.HasPrefix(line, "data: ") {
|
if !strings.HasPrefix(line, "data: ") {
|
||||||
|
|
@ -128,13 +169,61 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
||||||
return
|
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 {
|
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{
|
chunk := llm.StreamChunk{
|
||||||
Delta: choice.Delta.Content,
|
Delta: choice.Delta.Content,
|
||||||
|
Usage: usage,
|
||||||
}
|
}
|
||||||
if choice.FinishReason != "" {
|
if choice.FinishReason != "" {
|
||||||
chunk.FinishReason = 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) {
|
if !yield(chunk, nil) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -160,7 +249,7 @@ func (c *Client) Capabilities() llm.ProviderCapabilities {
|
||||||
}
|
}
|
||||||
|
|
||||||
// buildRequest converts an llm.CompletionRequest to the OpenAI API format.
|
// 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
|
// Convert messages to OpenAI format
|
||||||
messages := make([]openaiMessage, len(req.Messages))
|
messages := make([]openaiMessage, len(req.Messages))
|
||||||
for i, m := range 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
|
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{
|
openaiReq := openaiChatRequest{
|
||||||
Model: req.Model,
|
Model: model,
|
||||||
Messages: messages,
|
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 {
|
if len(tools) > 0 {
|
||||||
openaiReq.Tools = tools
|
openaiReq.Tools = tools
|
||||||
|
|
@ -222,11 +326,14 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("marshaling request: %w", err)
|
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.
|
// 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]
|
choice := resp.Choices[0]
|
||||||
result := llm.CompletionResponse{
|
result := llm.CompletionResponse{
|
||||||
ID: resp.ID,
|
ID: resp.ID,
|
||||||
|
|
@ -249,7 +356,7 @@ func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
|
||||||
TotalTokens: resp.Usage.TotalTokens,
|
TotalTokens: resp.Usage.TotalTokens,
|
||||||
}
|
}
|
||||||
|
|
||||||
return result
|
return result, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// OpenAI API types
|
// OpenAI API types
|
||||||
|
|
@ -263,6 +370,11 @@ type openaiChatRequest struct {
|
||||||
MaxTokens *int `json:"max_tokens,omitempty"`
|
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||||
Stop []string `json:"stop,omitempty"`
|
Stop []string `json:"stop,omitempty"`
|
||||||
Stream bool `json:"stream"`
|
Stream bool `json:"stream"`
|
||||||
|
StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiStreamOptions struct {
|
||||||
|
IncludeUsage bool `json:"include_usage"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type openaiMessage struct {
|
type openaiMessage struct {
|
||||||
|
|
@ -319,6 +431,7 @@ type openaiUsage struct {
|
||||||
type openaiStreamEvent struct {
|
type openaiStreamEvent struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Choices []openaiStreamChoice `json:"choices"`
|
Choices []openaiStreamChoice `json:"choices"`
|
||||||
|
Usage *openaiUsage `json:"usage"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type openaiStreamChoice struct {
|
type openaiStreamChoice struct {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue