feat: add OpenAI-compatible prompt token caching details and stream_options support
This commit is contained in:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user