From c8565951997f9ab6657111ee69d1b68bc23b70a3 Mon Sep 17 00:00:00 2001 From: Luxferre Date: Thu, 27 Aug 2026 14:35:49 +0300 Subject: [PATCH] feat: add OpenAI-compatible prompt token caching details and stream_options support --- main.go | 127 ++++++++++++++++++++++++++++++++++++++----------- qorona_test.go | 75 +++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 29 deletions(-) diff --git a/main.go b/main.go index c164366..6401aa2 100644 --- a/main.go +++ b/main.go @@ -98,18 +98,23 @@ func (m *ChatMessage) GetContentString() string { return fmt.Sprintf("%v", m.Content) } +type StreamOptions struct { + IncludeUsage bool `json:"include_usage,omitempty"` +} + type ChatCompletionRequest struct { - Model string `json:"model"` - Messages []ChatMessage `json:"messages"` - Tools []Tool `json:"tools,omitempty"` - ToolChoice interface{} `json:"tool_choice,omitempty"` - Stream bool `json:"stream"` - MaxTokens int `json:"max_tokens"` - MaxCompletionTokens int `json:"max_completion_tokens"` - Temperature *float64 `json:"temperature,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - ReasoningEffort string `json:"reasoning_effort,omitempty"` - Thinking interface{} `json:"thinking,omitempty"` + Model string `json:"model"` + Messages []ChatMessage `json:"messages"` + Tools []Tool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Stream bool `json:"stream"` + StreamOptions *StreamOptions `json:"stream_options,omitempty"` + MaxTokens int `json:"max_tokens"` + MaxCompletionTokens int `json:"max_completion_tokens"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + ReasoningEffort string `json:"reasoning_effort,omitempty"` + Thinking interface{} `json:"thinking,omitempty"` } type ChatCompletionResponseChoice struct { @@ -118,10 +123,20 @@ type ChatCompletionResponseChoice struct { FinishReason string `json:"finish_reason"` } +type PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` +} + +type CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens,omitempty"` +} + type Usage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"` + CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"` } type ChatCompletionResponse struct { @@ -152,6 +167,7 @@ type StreamResponse struct { Created int64 `json:"created"` Model string `json:"model"` Choices []StreamChoice `json:"choices"` + Usage *Usage `json:"usage,omitempty"` } // --------------------------------------------------------------------------- @@ -402,6 +418,23 @@ type FinalOutput struct { FinishReason string PromptTokens int CompletionTokens int + CachedTokens int + ReasoningTokens int +} + +func BuildUsage(promptTokens, completionTokens, cachedTokens, reasoningTokens int) Usage { + totalTokens := promptTokens + completionTokens + return Usage{ + PromptTokens: promptTokens, + CompletionTokens: completionTokens, + TotalTokens: totalTokens, + PromptTokensDetails: &PromptTokensDetails{ + CachedTokens: cachedTokens, + }, + CompletionTokensDetails: &CompletionTokensDetails{ + ReasoningTokens: reasoningTokens, + }, + } } func WriteCompletionResponse(w http.ResponseWriter, completionID string, created int64, model string, out FinalOutput) { @@ -409,7 +442,6 @@ func WriteCompletionResponse(w http.ResponseWriter, completionID string, created if finish == "" { finish = "stop" } - totalTokens := out.PromptTokens + out.CompletionTokens resp := ChatCompletionResponse{ ID: completionID, Object: "chat.completion", @@ -427,11 +459,7 @@ func WriteCompletionResponse(w http.ResponseWriter, completionID string, created FinishReason: finish, }, }, - Usage: Usage{ - PromptTokens: out.PromptTokens, - CompletionTokens: out.CompletionTokens, - TotalTokens: totalTokens, - }, + Usage: BuildUsage(out.PromptTokens, out.CompletionTokens, out.CachedTokens, out.ReasoningTokens), } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(resp) @@ -472,6 +500,22 @@ func (s *Streamer) Finish(reason string) { sendStreamChunk(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{}, &reason) } +func (s *Streamer) EmitUsage(usage Usage) { + chunk := StreamResponse{ + ID: s.id, + Object: "chat.completion.chunk", + Created: s.created, + Model: s.model, + Choices: []StreamChoice{}, + Usage: &usage, + } + b, _ := json.Marshal(chunk) + fmt.Fprintf(s.w, "data: %s\n\n", b) + if s.flusher != nil { + s.flusher.Flush() + } +} + func (s *Streamer) Done() { fmt.Fprintf(s.w, "data: [DONE]\n\n") if s.flusher != nil { @@ -1397,10 +1441,14 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re Cursor int `json:"cursor"` Done bool `json:"done"` FinalStatus struct { - Status string `json:"status"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - Error string `json:"error"` + Status string `json:"status"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CachedTokens int `json:"cached_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + Error string `json:"error"` } `json:"final_status"` } `json:"result"` } @@ -1418,8 +1466,19 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re finishReason = "tool_calls" } streamer.Finish(finishReason) + + promptTokens := sResult.Result.FinalStatus.PromptTokens + completionTokens := sResult.Result.FinalStatus.CompletionTokens + cachedTokens := sResult.Result.FinalStatus.CachedTokens + if cachedTokens == 0 { + cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens + } + + if req.StreamOptions != nil && req.StreamOptions.IncludeUsage { + streamer.EmitUsage(BuildUsage(promptTokens, completionTokens, cachedTokens, 0)) + } streamer.Done() - log.Printf("[Stream] Completed run_id=%s finish_reason=%s", runID, finishReason) + log.Printf("[Stream] Completed run_id=%s finish_reason=%s tokens=%d/%d", runID, finishReason, promptTokens, completionTokens) return } } @@ -1523,10 +1582,14 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re Cursor int `json:"cursor"` Done bool `json:"done"` FinalStatus struct { - Status string `json:"status"` - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - Error string `json:"error"` + Status string `json:"status"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CachedTokens int `json:"cached_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + Error string `json:"error"` } `json:"final_status"` } `json:"result"` } @@ -1539,6 +1602,10 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re sResp.Body.Close() promptTokens = sResult.Result.FinalStatus.PromptTokens completionTokens = sResult.Result.FinalStatus.CompletionTokens + cachedTokens := sResult.Result.FinalStatus.CachedTokens + if cachedTokens == 0 { + cachedTokens = sResult.Result.FinalStatus.PromptTokensDetails.CachedTokens + } fullContent := accumulatedContent.String() rawThinking := accumulatedThinking.String() @@ -1572,6 +1639,8 @@ func (gw *QoronaGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re FinishReason: finishReason, PromptTokens: promptTokens, CompletionTokens: completionTokens, + CachedTokens: cachedTokens, + ReasoningTokens: 0, } WriteCompletionResponse(w, completionID, createdTime, modelName, out) diff --git a/qorona_test.go b/qorona_test.go index 2f838e6..29e2c87 100644 --- a/qorona_test.go +++ b/qorona_test.go @@ -117,3 +117,78 @@ func TestModelsHandler(t *testing.T) { t.Fatalf("model not found in models list") } } + +func TestBuildUsage(t *testing.T) { + u := BuildUsage(100, 50, 80, 20) + if u.PromptTokens != 100 { + t.Fatalf("expected prompt_tokens 100, got %d", u.PromptTokens) + } + if u.CompletionTokens != 50 { + t.Fatalf("expected completion_tokens 50, got %d", u.CompletionTokens) + } + if u.TotalTokens != 150 { + t.Fatalf("expected total_tokens 150, got %d", u.TotalTokens) + } + if u.PromptTokensDetails == nil || u.PromptTokensDetails.CachedTokens != 80 { + t.Fatalf("expected cached_tokens 80, got %+v", u.PromptTokensDetails) + } + if u.CompletionTokensDetails == nil || u.CompletionTokensDetails.ReasoningTokens != 20 { + t.Fatalf("expected reasoning_tokens 20, got %+v", u.CompletionTokensDetails) + } + + b, err := json.Marshal(u) + if err != nil { + t.Fatalf("failed to marshal usage: %v", err) + } + if !strings.Contains(string(b), `"cached_tokens":80`) { + t.Fatalf("expected json to contain cached_tokens: %s", string(b)) + } +} + +func TestWriteCompletionResponseUsage(t *testing.T) { + w := httptest.NewRecorder() + out := FinalOutput{ + Content: "Hello", + ReasoningContent: "Thinking...", + FinishReason: "stop", + PromptTokens: 900, + CompletionTokens: 25, + CachedTokens: 850, + ReasoningTokens: 10, + } + + WriteCompletionResponse(w, "chatcmpl-test", 1234567890, "qwen-3.8-max", out) + if w.Code != http.StatusOK { + t.Fatalf("expected 200 OK, got %d", w.Code) + } + + var resp ChatCompletionResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + + if resp.Usage.PromptTokens != 900 { + t.Fatalf("expected prompt_tokens 900, got %d", resp.Usage.PromptTokens) + } + if resp.Usage.PromptTokensDetails == nil || resp.Usage.PromptTokensDetails.CachedTokens != 850 { + t.Fatalf("expected cached_tokens 850, got %+v", resp.Usage.PromptTokensDetails) + } +} + +func TestStreamerUsage(t *testing.T) { + w := httptest.NewRecorder() + streamer := NewStreamer(w, nil, "chatcmpl-stream-test", 1234567890, "qwen-3.8-max") + usage := BuildUsage(900, 30, 850, 0) + + streamer.EmitUsage(usage) + streamer.Done() + + body := w.Body.String() + if !strings.Contains(body, `"cached_tokens":850`) { + t.Fatalf("expected streamed chunk to contain cached_tokens: %s", body) + } + if !strings.Contains(body, `data: [DONE]`) { + t.Fatalf("expected streamed body to contain [DONE]: %s", body) + } +} +