From 0920e9c939dc674b749cb57245acc72c84a7c38d Mon Sep 17 00:00:00 2001 From: andev0x Date: Thu, 23 Jul 2026 20:23:08 +0700 Subject: [PATCH] fix(tokens): add stream options for usage inclusion - update openai.go to include stream options - modify openrouter.go to include stream options - adjust ui/stream.go to handle usage from stream options --- internal/llm/openai.go | 26 ++++++++++++++++---------- internal/providers/openai.go | 26 ++++++++++++++++---------- internal/providers/openrouter.go | 22 ++++++++++++---------- internal/ui/stream.go | 8 +++++--- 4 files changed, 49 insertions(+), 33 deletions(-) diff --git a/internal/llm/openai.go b/internal/llm/openai.go index d9ef664..9560c4c 100644 --- a/internal/llm/openai.go +++ b/internal/llm/openai.go @@ -35,12 +35,17 @@ type openAIMessage struct { Content string `json:"content"` } +type streamOptions struct { + IncludeUsage bool `json:"include_usage"` +} + type openAIReq struct { - Model string `json:"model"` - Messages []openAIMessage `json:"messages"` - Stream bool `json:"stream"` - MaxTokens int `json:"max_tokens,omitempty"` - Temperature float64 `json:"temperature,omitempty"` + Model string `json:"model"` + Messages []openAIMessage `json:"messages"` + Stream bool `json:"stream"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + StreamOptions *streamOptions `json:"stream_options,omitempty"` } type openAIResp struct { @@ -199,11 +204,12 @@ func (c *OpenAIClient) GenerateResponse(ctx context.Context, req PromptRequest) func (c *OpenAIClient) StreamResponse(ctx context.Context, req PromptRequest, handler StreamHandler) (LLMResponse, error) { body := openAIReq{ - Model: c.resolveModel(req.Model), - Messages: c.buildMessages(req), - Stream: true, - MaxTokens: req.MaxTokens, - Temperature: req.Temperature, + Model: c.resolveModel(req.Model), + Messages: c.buildMessages(req), + Stream: true, + MaxTokens: req.MaxTokens, + Temperature: req.Temperature, + StreamOptions: &streamOptions{IncludeUsage: true}, } if body.MaxTokens <= 0 { body.MaxTokens = 4096 diff --git a/internal/providers/openai.go b/internal/providers/openai.go index 6f08cd0..b1c36d2 100644 --- a/internal/providers/openai.go +++ b/internal/providers/openai.go @@ -37,9 +37,10 @@ type openaiMessage struct { } type openaiRequest struct { - Model string `json:"model"` - Messages []openaiMessage `json:"messages"` - Stream bool `json:"stream"` + Model string `json:"model"` + Messages []openaiMessage `json:"messages"` + Stream bool `json:"stream"` + StreamOptions *streamOptions `json:"stream_options,omitempty"` } type openaiResponse struct { @@ -63,6 +64,10 @@ type openaiDelta struct { Content string `json:"content,omitempty"` } +type streamOptions struct { + IncludeUsage bool `json:"include_usage"` +} + type openaiUsage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` @@ -155,9 +160,10 @@ func (p *OpenAIProvider) ExecuteStream(ctx context.Context, req ai.Request) (io. msgs := p.buildMessages(req) body := openaiRequest{ - Model: model, - Messages: msgs, - Stream: true, + Model: model, + Messages: msgs, + Stream: true, + StreamOptions: &streamOptions{IncludeUsage: true}, } payload, err := json.Marshal(body) @@ -243,14 +249,14 @@ func (s *openaiSSEReader) Read(p []byte) (int, error) { continue } - if len(chunk.Choices) == 0 { - continue - } - if chunk.Usage != nil { s.finalUsage = chunk.Usage } + if len(chunk.Choices) == 0 { + continue + } + if chunk.Choices[0].Delta != nil && chunk.Choices[0].Delta.Content != "" { n := copy(p, chunk.Choices[0].Delta.Content) return n, nil diff --git a/internal/providers/openrouter.go b/internal/providers/openrouter.go index af04d3f..438f7da 100644 --- a/internal/providers/openrouter.go +++ b/internal/providers/openrouter.go @@ -110,9 +110,10 @@ func (p *OpenRouterProvider) ExecuteStream(ctx context.Context, req ai.Request) msgs := p.buildMessages(req) body := openrouterRequest{ - Model: model, - Messages: msgs, - Stream: true, + Model: model, + Messages: msgs, + Stream: true, + StreamOptions: &streamOptions{IncludeUsage: true}, } payload, err := json.Marshal(body) @@ -161,9 +162,10 @@ type openrouterMessage struct { } type openrouterRequest struct { - Model string `json:"model"` - Messages []openrouterMessage `json:"messages"` - Stream bool `json:"stream"` + Model string `json:"model"` + Messages []openrouterMessage `json:"messages"` + Stream bool `json:"stream"` + StreamOptions *streamOptions `json:"stream_options,omitempty"` } type openrouterResponse struct { @@ -253,14 +255,14 @@ func (s *openrouterSSEReader) Read(p []byte) (int, error) { continue } - if len(chunk.Choices) == 0 { - continue - } - if chunk.Usage != nil { s.finalUsage = chunk.Usage } + if len(chunk.Choices) == 0 { + continue + } + if chunk.Choices[0].Delta != nil && chunk.Choices[0].Delta.Content != "" { n := copy(p, chunk.Choices[0].Delta.Content) return n, nil diff --git a/internal/ui/stream.go b/internal/ui/stream.go index 235dde0..c0d21d3 100644 --- a/internal/ui/stream.go +++ b/internal/ui/stream.go @@ -17,7 +17,6 @@ import ( "github.com/PizenLabs/izen/internal/modes" "github.com/PizenLabs/izen/internal/modes/plan" "github.com/PizenLabs/izen/internal/prompt" - "github.com/PizenLabs/izen/internal/providers" ) // debugLogPayload writes the exact outgoing LLM payload to @@ -185,8 +184,11 @@ func (m *model) streamCmd(content string) tea.Cmd { streamCh <- tokenMsg(chunk) } if err == io.EOF { - if sr, ok := rawStream.(*providers.StreamResult); ok { - tokIn, tokOut = sr.Usage() + type usageProvider interface { + Usage() (input, output int) + } + if up, ok := rawStream.(usageProvider); ok { + tokIn, tokOut = up.Usage() } if tokIn == 0 && tokOut == 0 { tokIn = len(content) / 4