upstream-archive/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-07-23 09:10:53 -05:00
Code Issues Packages Projects Releases Wiki Activity

server: add cached eval metric to response

Browse Source 
Added cached prompt token counts to Ollama responses and compatibility usage fields.

This carries local `llama-server` `cache_n` and MLX cache hits through `/api/generate`, `/api/chat`, OpenAI-compatible endpoints, and Anthropic-compatible `/v1/messages`. Cloud responses are passed through as-is, so cache counts will show up there once Cloud starts returning them.
This commit is contained in:
Bruce MacDonald
2026-06-25 16:21:08 -07:00
parent 2a6b50421a
commit 6e65d95ef5
19 changed files with 360 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
anthropic/anthropic.go
View File

@@ -217,8 +217,30 @@ type MessagesResponse struct {
// Usage contains token usage information
type Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
InputTokens int `json:"input_tokens"`
CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
CacheReadInputTokens int `json:"cache_read_input_tokens"`
OutputTokens int `json:"output_tokens"`
}
func UsageFromMetrics(metrics api.Metrics) Usage {
cached := metrics.PromptEvalCachedCount
if cached > metrics.PromptEvalCount {
cached = metrics.PromptEvalCount
}
return Usage{
InputTokens: metrics.PromptEvalCount - cached,
CacheReadInputTokens: cached,
OutputTokens: metrics.EvalCount,
}
}
func (u *Usage) Add(other Usage) {
u.InputTokens += other.InputTokens
u.CacheCreationInputTokens += other.CacheCreationInputTokens
u.CacheReadInputTokens += other.CacheReadInputTokens
u.OutputTokens += other.OutputTokens
}
// Streaming event types
@@ -273,8 +295,10 @@ type MessageDelta struct {
// DeltaUsage contains cumulative token usage
type DeltaUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
InputTokens int `json:"input_tokens"`
CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
CacheReadInputTokens int `json:"cache_read_input_tokens"`
OutputTokens int `json:"output_tokens"`
}
// MessageStopEvent signals the end of the message
@@ -688,10 +712,7 @@ func ToMessagesResponse(id string, r api.ChatResponse) MessagesResponse {
Model: r.Model,
Content: content,
StopReason: stopReason,
Usage: Usage{
InputTokens: r.Metrics.PromptEvalCount,
OutputTokens: r.Metrics.EvalCount,
},
Usage: UsageFromMetrics(r.Metrics),
}
}
@@ -721,6 +742,8 @@ type StreamConverter struct {
firstWrite bool
contentIndex int
inputTokens int
cacheCreationTokens int
cacheReadTokens int
outputTokens int
estimatedInputTokens int // Estimated tokens from request (used when actual metrics are 0)
thinkingStarted bool
@@ -753,7 +776,12 @@ func (c *StreamConverter) Process(r api.ChatResponse) []StreamEvent {
c.firstWrite = false
// Use actual metrics if available, otherwise use estimate
c.inputTokens = r.Metrics.PromptEvalCount
if c.inputTokens == 0 && c.estimatedInputTokens > 0 {
c.cacheReadTokens = r.Metrics.PromptEvalCachedCount
if c.cacheReadTokens > c.inputTokens {
c.cacheReadTokens = c.inputTokens
}
c.inputTokens -= c.cacheReadTokens
if c.inputTokens == 0 && c.cacheReadTokens == 0 && c.estimatedInputTokens > 0 {
c.inputTokens = c.estimatedInputTokens
}
@@ -768,8 +796,10 @@ func (c *StreamConverter) Process(r api.ChatResponse) []StreamEvent {
Model: c.Model,
Content: []ContentBlock{},
Usage: Usage{
InputTokens: c.inputTokens,
OutputTokens: 0,
InputTokens: c.inputTokens,
CacheCreationInputTokens: c.cacheCreationTokens,
CacheReadInputTokens: c.cacheReadTokens,
OutputTokens: 0,
},
},
},
@@ -939,6 +969,11 @@ func (c *StreamConverter) Process(r api.ChatResponse) []StreamEvent {
}
c.inputTokens = r.Metrics.PromptEvalCount
c.cacheReadTokens = r.Metrics.PromptEvalCachedCount
if c.cacheReadTokens > c.inputTokens {
c.cacheReadTokens = c.inputTokens
}
c.inputTokens -= c.cacheReadTokens
c.outputTokens = r.Metrics.EvalCount
stopReason := mapStopReason(r.DoneReason, len(c.toolCallsSent) > 0)
@@ -950,8 +985,10 @@ func (c *StreamConverter) Process(r api.ChatResponse) []StreamEvent {
StopReason: stopReason,
},
Usage: DeltaUsage{
InputTokens: c.inputTokens,
OutputTokens: c.outputTokens,
InputTokens: c.inputTokens,
CacheCreationInputTokens: c.cacheCreationTokens,
CacheReadInputTokens: c.cacheReadTokens,
OutputTokens: c.outputTokens,
},
},
})

77
anthropic/anthropic_test.go
View File

@@ -778,6 +778,38 @@ func TestToMessagesResponse_Basic(t *testing.T) {
}
}
func TestToMessagesResponse_UsageIncludesCacheReadTokens(t *testing.T) {
resp := api.ChatResponse{
Model: "test-model",
Message: api.Message{
Role: "assistant",
Content: "Hello there!",
},
Done: true,
DoneReason: "stop",
Metrics: api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 5,
},
}
result := ToMessagesResponse("msg_123", resp)
if result.Usage.InputTokens != 6 {
t.Errorf("expected input_tokens 6, got %d", result.Usage.InputTokens)
}
if result.Usage.CacheReadInputTokens != 4 {
t.Errorf("expected cache_read_input_tokens 4, got %d", result.Usage.CacheReadInputTokens)
}
if result.Usage.CacheCreationInputTokens != 0 {
t.Errorf("expected cache_creation_input_tokens 0, got %d", result.Usage.CacheCreationInputTokens)
}
if result.Usage.OutputTokens != 5 {
t.Errorf("expected output_tokens 5, got %d", result.Usage.OutputTokens)
}
}
func TestToMessagesResponse_WithToolCalls(t *testing.T) {
resp := api.ChatResponse{
Model: "test-model",
@@ -988,6 +1020,51 @@ func TestStreamConverter_Basic(t *testing.T) {
}
}
func TestStreamConverter_UsageIncludesCacheReadTokens(t *testing.T) {
conv := NewStreamConverter("msg_123", "test-model", 0)
events := conv.Process(api.ChatResponse{
Model: "test-model",
Message: api.Message{Role: "assistant"},
Done: true,
DoneReason: "stop",
Metrics: api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 5,
},
})
var found bool
for _, e := range events {
if e.Event != "message_delta" {
continue
}
data, ok := e.Data.(MessageDeltaEvent)
if !ok {
t.Fatalf("unexpected data type: %T", e.Data)
}
found = true
if data.Usage.InputTokens != 6 {
t.Errorf("expected input_tokens 6, got %d", data.Usage.InputTokens)
}
if data.Usage.CacheReadInputTokens != 4 {
t.Errorf("expected cache_read_input_tokens 4, got %d", data.Usage.CacheReadInputTokens)
}
if data.Usage.CacheCreationInputTokens != 0 {
t.Errorf("expected cache_creation_input_tokens 0, got %d", data.Usage.CacheCreationInputTokens)
}
if data.Usage.OutputTokens != 5 {
t.Errorf("expected output_tokens 5, got %d", data.Usage.OutputTokens)
}
}
if !found {
t.Fatal("expected message_delta event")
}
}
func TestStreamConverter_WithToolCalls(t *testing.T) {
conv := NewStreamConverter("msg_123", "test-model", 0)

17
api/types.go
View File

@@ -569,12 +569,13 @@ type DebugInfo struct {
}
type Metrics struct {
TotalDuration time.Duration `json:"total_duration,omitempty"`
LoadDuration time.Duration `json:"load_duration,omitempty"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration,omitempty"`
EvalCount int `json:"eval_count,omitempty"`
EvalDuration time.Duration `json:"eval_duration,omitempty"`
TotalDuration time.Duration `json:"total_duration,omitempty"`
LoadDuration time.Duration `json:"load_duration,omitempty"`
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
PromptEvalCachedCount int `json:"prompt_eval_cached_count,omitempty"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration,omitempty"`
EvalCount int `json:"eval_count,omitempty"`
EvalDuration time.Duration `json:"eval_duration,omitempty"`
}
// Options specified in [GenerateRequest]. If you add a new option here, also
@@ -968,6 +969,10 @@ func (m *Metrics) Summary() {
fmt.Fprintf(os.Stderr, "prompt eval count: %d token(s)\n", m.PromptEvalCount)
}
if m.PromptEvalCachedCount > 0 {
fmt.Fprintf(os.Stderr, "prompt eval cached: %d token(s)\n", m.PromptEvalCachedCount)
}
if m.PromptEvalDuration > 0 {
fmt.Fprintf(os.Stderr, "prompt eval duration: %s\n", m.PromptEvalDuration)
fmt.Fprintf(os.Stderr, "prompt eval rate: %.2f tokens/s\n", float64(m.PromptEvalCount)/m.PromptEvalDuration.Seconds())

1
docs/api.md
View File

@@ -110,6 +110,7 @@ The final response in the stream also includes additional data about the generat
- `total_duration`: time spent generating the response
- `load_duration`: time spent in nanoseconds loading the model
- `prompt_eval_count`: number of tokens in the prompt
- `prompt_eval_cached_count`: number of prompt tokens read from the prompt cache, when available
- `prompt_eval_duration`: time spent in nanoseconds evaluating the prompt
- `eval_count`: number of tokens in the response
- `eval_duration`: time in nanoseconds spent generating the response

1
docs/api/usage.mdx
View File

@@ -7,6 +7,7 @@ Ollama's API responses include metrics that can be used for measuring performanc
* `total_duration`: How long the response took to generate
* `load_duration`: How long the model took to load
* `prompt_eval_count`: How many input tokens were processed
* `prompt_eval_cached_count`: How many input tokens were read from the prompt cache, when available
* `prompt_eval_duration`: How long it took to evaluate the prompt
* `eval_count`: How many output tokens were processes
* `eval_duration`: How long it took to generate the output tokens

9
docs/openapi.yaml
View File

@@ -147,6 +147,9 @@ components:
prompt_eval_count:
type: integer
description: Number of input tokens in the prompt
prompt_eval_cached_count:
type: integer
description: Number of input tokens read from the prompt cache
prompt_eval_duration:
type: integer
description: Time spent evaluating the prompt in nanoseconds
@@ -191,6 +194,9 @@ components:
prompt_eval_count:
type: integer
description: Number of input tokens in the prompt
prompt_eval_cached_count:
type: integer
description: Number of input tokens read from the prompt cache
prompt_eval_duration:
type: integer
description: Time spent evaluating the prompt in nanoseconds
@@ -352,6 +358,9 @@ components:
prompt_eval_count:
type: integer
description: Number of tokens in the prompt
prompt_eval_cached_count:
type: integer
description: Number of prompt tokens read from the prompt cache
prompt_eval_duration:
type: integer
description: Time spent evaluating the prompt in nanoseconds

16
llm/llama_server.go
View File

@@ -1625,13 +1625,14 @@ func (s *llamaServerRunner) Completion(ctx context.Context, req CompletionReques
}
finalResp = CompletionResponse{
Content: lsResp.Content,
Done: true,
DoneReason: doneReason,
PromptEvalCount: lsResp.Timings.promptEvalCount(),
PromptEvalDuration: time.Duration(lsResp.Timings.PromptMS * float64(time.Millisecond)),
EvalCount: lsResp.Timings.PredictN,
EvalDuration: time.Duration(lsResp.Timings.PredictMS * float64(time.Millisecond)),
Content: lsResp.Content,
Done: true,
DoneReason: doneReason,
PromptEvalCount: lsResp.Timings.promptEvalCount(),
PromptEvalCachedCount: lsResp.Timings.CacheN,
PromptEvalDuration: time.Duration(lsResp.Timings.PromptMS * float64(time.Millisecond)),
EvalCount: lsResp.Timings.PredictN,
EvalDuration: time.Duration(lsResp.Timings.PredictMS * float64(time.Millisecond)),
}
hasFinalResp = true
}
@@ -1928,6 +1929,7 @@ func (s *llamaServerRunner) Chat(ctx context.Context, req ChatRequest, fn func(C
resp.Done = true
resp.DoneReason = doneReason
resp.PromptEvalCount = lsResp.Timings.promptEvalCount()
resp.PromptEvalCachedCount = lsResp.Timings.CacheN
resp.PromptEvalDuration = time.Duration(lsResp.Timings.PromptMS * float64(time.Millisecond))
resp.EvalCount = lsResp.Timings.PredictN
resp.EvalDuration = time.Duration(lsResp.Timings.PredictMS * float64(time.Millisecond))

6
llm/llama_server_test.go
View File

@@ -291,6 +291,9 @@ func TestLlamaServerCompletionPromptEvalCountIncludesCache(t *testing.T) {
if responses[0].PromptEvalCount != 17 {
t.Errorf("PromptEvalCount = %d, want 17", responses[0].PromptEvalCount)
}
if responses[0].PromptEvalCachedCount != 12 {
t.Errorf("PromptEvalCachedCount = %d, want 12", responses[0].PromptEvalCachedCount)
}
if responses[0].PromptEvalDuration != 10*time.Millisecond {
t.Errorf("PromptEvalDuration = %s, want 10ms", responses[0].PromptEvalDuration)
}
@@ -341,6 +344,9 @@ func TestLlamaServerChatPromptEvalCountIncludesCache(t *testing.T) {
if responses[1].PromptEvalCount != 17 {
t.Errorf("PromptEvalCount = %d, want 17", responses[1].PromptEvalCount)
}
if responses[1].PromptEvalCachedCount != 12 {
t.Errorf("PromptEvalCachedCount = %d, want 12", responses[1].PromptEvalCachedCount)
}
if responses[1].PromptEvalDuration != 10*time.Millisecond {
t.Errorf("PromptEvalDuration = %s, want 10ms", responses[1].PromptEvalDuration)
}

32
llm/server.go
View File

@@ -242,14 +242,15 @@ type ChatRequest struct {
}
type ChatResponse struct {
Message api.Message `json:"message"`
DoneReason DoneReason `json:"done_reason"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration"`
EvalCount int `json:"eval_count"`
EvalDuration time.Duration `json:"eval_duration"`
Logprobs []Logprob `json:"logprobs,omitempty"`
Message api.Message `json:"message"`
DoneReason DoneReason `json:"done_reason"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
PromptEvalCachedCount int `json:"prompt_eval_cached_count"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration"`
EvalCount int `json:"eval_count"`
EvalDuration time.Duration `json:"eval_duration"`
Logprobs []Logprob `json:"logprobs,omitempty"`
}
// DoneReason represents the reason why a completion response is done
@@ -285,13 +286,14 @@ type Logprob struct {
}
type CompletionResponse struct {
Content string `json:"content"`
DoneReason DoneReason `json:"done_reason"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration"`
EvalCount int `json:"eval_count"`
EvalDuration time.Duration `json:"eval_duration"`
Content string `json:"content"`
DoneReason DoneReason `json:"done_reason"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
PromptEvalCachedCount int `json:"prompt_eval_cached_count"`
PromptEvalDuration time.Duration `json:"prompt_eval_duration"`
EvalCount int `json:"eval_count"`
EvalDuration time.Duration `json:"eval_duration"`
// Logprobs contains log probability information if requested
Logprobs []Logprob `json:"logprobs,omitempty"`

75
middleware/anthropic.go
View File

@@ -101,13 +101,13 @@ type WebSearchAnthropicWriter struct {
terminalSent bool
observedPromptEvalCount int
observedEvalCount int
observedPromptEvalCount int
observedPromptEvalCachedCount int
observedEvalCount int
loopInFlight bool
loopBaseInputTok int
loopBaseOutputTok int
loopResultCh chan webSearchLoopResult
loopInFlight bool
loopBaseUsage anthropic.Usage
loopResultCh chan webSearchLoopResult
streamMessageStarted bool
streamHasOpenBlock bool
@@ -201,10 +201,7 @@ func (w *WebSearchAnthropicWriter) Write(data []byte) (int, error) {
loopCtx, cancel := w.startLoopContext()
defer cancel()
initialUsage := anthropic.Usage{
InputTokens: max(w.observedPromptEvalCount, chatResponse.Metrics.PromptEvalCount),
OutputTokens: max(w.observedEvalCount, chatResponse.Metrics.EvalCount),
}
initialUsage := w.usageWithObservedMetrics(chatResponse.Metrics)
logutil.Trace("anthropic middleware: starting sync web_search loop",
"tool_call", anthropic.TraceToolCall(webSearchCall),
"resp", anthropic.TraceChatResponse(chatResponse),
@@ -319,8 +316,7 @@ func (w *WebSearchAnthropicWriter) runWebSearchLoop(ctx context.Context, initial
"resp", anthropic.TraceChatResponse(followUpResponse),
)
usage.InputTokens += followUpResponse.Metrics.PromptEvalCount
usage.OutputTokens += followUpResponse.Metrics.EvalCount
usage.Add(anthropic.UsageFromMetrics(followUpResponse.Metrics))
nextToolCall, hasWebSearch, hasOtherTools := findWebSearchToolCall(followUpResponse.Message.ToolCalls)
if hasWebSearch && hasOtherTools {
@@ -380,12 +376,8 @@ func (w *WebSearchAnthropicWriter) startLoopWorker(initialResponse api.ChatRespo
return
}
initialUsage := anthropic.Usage{
InputTokens: max(w.observedPromptEvalCount, initialResponse.Metrics.PromptEvalCount),
OutputTokens: max(w.observedEvalCount, initialResponse.Metrics.EvalCount),
}
w.loopBaseInputTok = initialUsage.InputTokens
w.loopBaseOutputTok = initialUsage.OutputTokens
initialUsage := w.usageWithObservedMetrics(initialResponse.Metrics)
w.loopBaseUsage = initialUsage
w.loopResultCh = make(chan webSearchLoopResult, 1)
w.loopInFlight = true
logutil.Trace("anthropic middleware: loop worker started",
@@ -438,25 +430,42 @@ func (w *WebSearchAnthropicWriter) recordObservedUsage(metrics api.Metrics) {
if metrics.PromptEvalCount > w.observedPromptEvalCount {
w.observedPromptEvalCount = metrics.PromptEvalCount
}
if metrics.PromptEvalCachedCount > w.observedPromptEvalCachedCount {
w.observedPromptEvalCachedCount = metrics.PromptEvalCachedCount
}
if metrics.EvalCount > w.observedEvalCount {
w.observedEvalCount = metrics.EvalCount
}
}
func (w *WebSearchAnthropicWriter) applyObservedUsageDeltaToUsage(usage *anthropic.Usage) {
if deltaIn := w.observedPromptEvalCount - w.loopBaseInputTok; deltaIn > 0 {
usage.InputTokens += deltaIn
}
if deltaOut := w.observedEvalCount - w.loopBaseOutputTok; deltaOut > 0 {
usage.OutputTokens += deltaOut
}
observed := w.currentObservedUsage()
usage.InputTokens += observed.InputTokens - w.loopBaseUsage.InputTokens
usage.CacheCreationInputTokens += observed.CacheCreationInputTokens - w.loopBaseUsage.CacheCreationInputTokens
usage.CacheReadInputTokens += observed.CacheReadInputTokens - w.loopBaseUsage.CacheReadInputTokens
usage.OutputTokens += observed.OutputTokens - w.loopBaseUsage.OutputTokens
}
func (w *WebSearchAnthropicWriter) currentObservedUsage() anthropic.Usage {
return anthropic.Usage{
InputTokens: w.observedPromptEvalCount,
OutputTokens: w.observedEvalCount,
return anthropic.UsageFromMetrics(api.Metrics{
PromptEvalCount: w.observedPromptEvalCount,
PromptEvalCachedCount: w.observedPromptEvalCachedCount,
EvalCount: w.observedEvalCount,
})
}
func (w *WebSearchAnthropicWriter) usageWithObservedMetrics(metrics api.Metrics) anthropic.Usage {
if w.observedPromptEvalCount > metrics.PromptEvalCount {
metrics.PromptEvalCount = w.observedPromptEvalCount
}
if w.observedPromptEvalCachedCount > metrics.PromptEvalCachedCount {
metrics.PromptEvalCachedCount = w.observedPromptEvalCachedCount
}
if w.observedEvalCount > metrics.EvalCount {
metrics.EvalCount = w.observedEvalCount
}
return anthropic.UsageFromMetrics(metrics)
}
func (w *WebSearchAnthropicWriter) startLoopContext() (context.Context, context.CancelFunc) {
@@ -625,7 +634,7 @@ func (w *WebSearchAnthropicWriter) ensureStreamMessageStart(usage anthropic.Usag
}
inputTokens := usage.InputTokens
if inputTokens == 0 {
if inputTokens == 0 && usage.CacheCreationInputTokens == 0 && usage.CacheReadInputTokens == 0 {
inputTokens = w.estimatedInputTokens
}
@@ -638,7 +647,9 @@ func (w *WebSearchAnthropicWriter) ensureStreamMessageStart(usage anthropic.Usag
Model: w.req.Model,
Content: []anthropic.ContentBlock{},
Usage: anthropic.Usage{
InputTokens: inputTokens,
InputTokens: inputTokens,
CacheCreationInputTokens: usage.CacheCreationInputTokens,
CacheReadInputTokens: usage.CacheReadInputTokens,
},
},
}); err != nil {
@@ -751,8 +762,10 @@ func (w *WebSearchAnthropicWriter) writeTerminalResponse(response anthropic.Mess
StopReason: response.StopReason,
},
Usage: anthropic.DeltaUsage{
InputTokens: response.Usage.InputTokens,
OutputTokens: response.Usage.OutputTokens,
InputTokens: response.Usage.InputTokens,
CacheCreationInputTokens: response.Usage.CacheCreationInputTokens,
CacheReadInputTokens: response.Usage.CacheReadInputTokens,
OutputTokens: response.Usage.OutputTokens,
},
}); err != nil {
return err

13
middleware/anthropic_test.go
View File

@@ -2164,7 +2164,7 @@ func TestWebSearchStreamingUsageUsesObservedChunkMetrics(t *testing.T) {
Message: api.Message{Role: "assistant", Content: "After search."},
Done: true,
DoneReason: "stop",
Metrics: api.Metrics{PromptEvalCount: 20, EvalCount: 7},
Metrics: api.Metrics{PromptEvalCount: 20, PromptEvalCachedCount: 5, EvalCount: 7},
}
_ = json.NewEncoder(w).Encode(resp)
}))
@@ -2192,7 +2192,7 @@ func TestWebSearchStreamingUsageUsesObservedChunkMetrics(t *testing.T) {
Model: "test-model",
Message: api.Message{Role: "assistant", Content: "Preface "},
Done: false,
Metrics: api.Metrics{PromptEvalCount: 12, EvalCount: 4},
Metrics: api.Metrics{PromptEvalCount: 12, PromptEvalCachedCount: 4, EvalCount: 4},
},
{
Model: "test-model",
@@ -2216,7 +2216,7 @@ func TestWebSearchStreamingUsageUsesObservedChunkMetrics(t *testing.T) {
Message: api.Message{Role: "assistant"},
Done: true,
DoneReason: "stop",
Metrics: api.Metrics{PromptEvalCount: 12, EvalCount: 4},
Metrics: api.Metrics{PromptEvalCount: 12, PromptEvalCachedCount: 4, EvalCount: 4},
},
}
c.Writer.WriteHeader(http.StatusOK)
@@ -2259,8 +2259,11 @@ func TestWebSearchStreamingUsageUsesObservedChunkMetrics(t *testing.T) {
if !found {
t.Fatal("expected message_delta event")
}
if messageDelta.Usage.InputTokens != 32 {
t.Fatalf("expected aggregated input tokens 32 (12 passthrough + 20 followup), got %d", messageDelta.Usage.InputTokens)
if messageDelta.Usage.InputTokens != 23 {
t.Fatalf("expected aggregated fresh input tokens 23 ((12-4) passthrough + (20-5) followup), got %d", messageDelta.Usage.InputTokens)
}
if messageDelta.Usage.CacheReadInputTokens != 9 {
t.Fatalf("expected aggregated cache read tokens 9 (4 passthrough + 5 followup), got %d", messageDelta.Usage.CacheReadInputTokens)
}
if messageDelta.Usage.OutputTokens != 11 {
t.Fatalf("expected aggregated output tokens 11 (4 passthrough + 7 followup), got %d", messageDelta.Usage.OutputTokens)

25
openai/openai.go
View File

@@ -65,10 +65,15 @@ type CompleteChunkChoice struct {
Logprobs *ChoiceLogprobs `json:"logprobs,omitempty"`
}
type PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
}
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokens int `json:"prompt_tokens"`
PromptTokensDetails PromptTokensDetails `json:"prompt_tokens_details"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}
type ResponseFormat struct {
@@ -232,9 +237,10 @@ func NewError(code int, message string) ErrorResponse {
// ToUsage converts an api.ChatResponse to Usage
func ToUsage(r api.ChatResponse) Usage {
return Usage{
PromptTokens: r.Metrics.PromptEvalCount,
CompletionTokens: r.Metrics.EvalCount,
TotalTokens: r.Metrics.PromptEvalCount + r.Metrics.EvalCount,
PromptTokens: r.Metrics.PromptEvalCount,
PromptTokensDetails: PromptTokensDetails{CachedTokens: r.Metrics.PromptEvalCachedCount},
CompletionTokens: r.Metrics.EvalCount,
TotalTokens: r.Metrics.PromptEvalCount + r.Metrics.EvalCount,
}
}
@@ -355,9 +361,10 @@ func ToChunk(id string, r api.ChatResponse, toolCallSent bool) ChatCompletionChu
// ToUsageGenerate converts an api.GenerateResponse to Usage
func ToUsageGenerate(r api.GenerateResponse) Usage {
return Usage{
PromptTokens: r.Metrics.PromptEvalCount,
CompletionTokens: r.Metrics.EvalCount,
TotalTokens: r.Metrics.PromptEvalCount + r.Metrics.EvalCount,
PromptTokens: r.Metrics.PromptEvalCount,
PromptTokensDetails: PromptTokensDetails{CachedTokens: r.Metrics.PromptEvalCachedCount},
CompletionTokens: r.Metrics.EvalCount,
TotalTokens: r.Metrics.PromptEvalCount + r.Metrics.EvalCount,
}
}

9
openai/openai_test.go
View File

@@ -176,8 +176,9 @@ func TestFromCompleteRequest_Basic(t *testing.T) {
func TestToUsage(t *testing.T) {
resp := api.ChatResponse{
Metrics: api.Metrics{
PromptEvalCount: 10,
EvalCount: 20,
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 20,
},
}
@@ -187,6 +188,10 @@ func TestToUsage(t *testing.T) {
t.Errorf("expected PromptTokens 10, got %d", usage.PromptTokens)
}
if usage.PromptTokensDetails.CachedTokens != 4 {
t.Errorf("expected CachedTokens 4, got %d", usage.PromptTokensDetails.CachedTokens)
}
if usage.CompletionTokens != 20 {
t.Errorf("expected CompletionTokens 20, got %d", usage.CompletionTokens)
}

11
openai/responses.go
View File

@@ -856,11 +856,10 @@ func ToResponse(model, responseID, itemID string, chatResponse api.ChatResponse,
Temperature: derefFloat64(request.Temperature, 1.0),
Reasoning: reasoning,
Usage: &ResponsesUsage{
InputTokens: chatResponse.PromptEvalCount,
OutputTokens: chatResponse.EvalCount,
TotalTokens: chatResponse.PromptEvalCount + chatResponse.EvalCount,
// TODO(drifkin): wire through the actual values
InputTokensDetails: ResponsesInputTokensDetails{CachedTokens: 0},
InputTokens: chatResponse.PromptEvalCount,
OutputTokens: chatResponse.EvalCount,
TotalTokens: chatResponse.PromptEvalCount + chatResponse.EvalCount,
InputTokensDetails: ResponsesInputTokensDetails{CachedTokens: chatResponse.PromptEvalCachedCount},
// TODO(drifkin): wire through the actual values
OutputTokensDetails: ResponsesOutputTokensDetails{ReasoningTokens: 0},
},
@@ -1366,7 +1365,7 @@ func (c *ResponsesStreamConverter) processCompletion(r api.ChatResponse) []Respo
"output_tokens": r.EvalCount,
"total_tokens": r.PromptEvalCount + r.EvalCount,
"input_tokens_details": map[string]any{
"cached_tokens": 0,
"cached_tokens": r.PromptEvalCachedCount,
},
"output_tokens_details": map[string]any{
"reasoning_tokens": 0,

42
openai/responses_test.go
View File

@@ -1506,6 +1506,32 @@ func TestToResponse_WithReasoning(t *testing.T) {
}
}
func TestToResponse_UsageIncludesCachedTokens(t *testing.T) {
response := ToResponse("gpt-oss:20b", "resp_123", "msg_456", api.ChatResponse{
CreatedAt: time.Now(),
Message: api.Message{Content: "The answer is 42"},
Done: true,
Metrics: api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 3,
},
}, ResponsesRequest{})
if response.Usage == nil {
t.Fatal("expected usage")
}
if response.Usage.InputTokens != 10 {
t.Errorf("InputTokens = %d, want 10", response.Usage.InputTokens)
}
if response.Usage.InputTokensDetails.CachedTokens != 4 {
t.Errorf("CachedTokens = %d, want 4", response.Usage.InputTokensDetails.CachedTokens)
}
if response.Usage.TotalTokens != 13 {
t.Errorf("TotalTokens = %d, want 13", response.Usage.TotalTokens)
}
}
func TestFromResponsesRequest_Instructions(t *testing.T) {
reqJSON := `{
"model": "gpt-oss:20b",
@@ -1858,6 +1884,11 @@ func TestResponsesStreamConverter_OutputIncludesContent(t *testing.T) {
events := converter.Process(api.ChatResponse{
Message: api.Message{},
Done: true,
Metrics: api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 3,
},
})
// Find the output_item.done event
@@ -1906,6 +1937,11 @@ func TestResponsesStreamConverter_ResponseCompletedIncludesOutput(t *testing.T)
events := converter.Process(api.ChatResponse{
Message: api.Message{},
Done: true,
Metrics: api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: 4,
EvalCount: 3,
},
})
// Find the response.completed event
@@ -1935,6 +1971,12 @@ func TestResponsesStreamConverter_ResponseCompletedIncludesOutput(t *testing.T)
if item["type"] != "message" {
t.Errorf("output[0].type = %q, want %q", item["type"], "message")
}
usage := response["usage"].(map[string]any)
inputTokensDetails := usage["input_tokens_details"].(map[string]any)
if inputTokensDetails["cached_tokens"] != 4 {
t.Errorf("cached_tokens = %v, want 4", inputTokensDetails["cached_tokens"])
}
}
func TestResponsesStreamConverter_ResponseCreatedIncludesOutput(t *testing.T) {

27
server/routes.go
View File

@@ -681,10 +681,11 @@ func (s *Server) GenerateHandler(c *gin.Context) {
Response: cr.Content,
Done: cr.Done,
Metrics: api.Metrics{
PromptEvalCount: cr.PromptEvalCount,
PromptEvalDuration: cr.PromptEvalDuration,
EvalCount: cr.EvalCount,
EvalDuration: cr.EvalDuration,
PromptEvalCount: cr.PromptEvalCount,
PromptEvalCachedCount: cr.PromptEvalCachedCount,
PromptEvalDuration: cr.PromptEvalDuration,
EvalCount: cr.EvalCount,
EvalDuration: cr.EvalDuration,
},
Logprobs: toAPILogprobs(cr.Logprobs),
}
@@ -2814,10 +2815,11 @@ func (s *Server) ChatHandler(c *gin.Context) {
Message: api.Message{Role: "assistant", Content: r.Content},
Done: r.Done,
Metrics: api.Metrics{
PromptEvalCount: r.PromptEvalCount,
PromptEvalDuration: r.PromptEvalDuration,
EvalCount: r.EvalCount,
EvalDuration: r.EvalDuration,
PromptEvalCount: r.PromptEvalCount,
PromptEvalCachedCount: r.PromptEvalCachedCount,
PromptEvalDuration: r.PromptEvalDuration,
EvalCount: r.EvalCount,
EvalDuration: r.EvalDuration,
},
Logprobs: toAPILogprobs(r.Logprobs),
}
@@ -3022,10 +3024,11 @@ func (s *Server) handleNativeChat(c *gin.Context, req api.ChatRequest, m *Model,
Message: r.Message,
Done: r.Done,
Metrics: api.Metrics{
PromptEvalCount: r.PromptEvalCount,
PromptEvalDuration: r.PromptEvalDuration,
EvalCount: r.EvalCount,
EvalDuration: r.EvalDuration,
PromptEvalCount: r.PromptEvalCount,
PromptEvalCachedCount: r.PromptEvalCachedCount,
PromptEvalDuration: r.PromptEvalDuration,
EvalCount: r.EvalCount,
EvalDuration: r.EvalDuration,
},
Logprobs: toAPILogprobs(r.Logprobs),
}

34
server/routes_generate_test.go
View File

@@ -754,12 +754,13 @@ func TestGenerateChat(t *testing.T) {
mock := mockRunner{
CompletionResponse: llm.CompletionResponse{
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: 1,
PromptEvalDuration: 1,
EvalCount: 1,
EvalDuration: 1,
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: 1,
PromptEvalCachedCount: 1,
PromptEvalDuration: 1,
EvalCount: 1,
EvalDuration: 1,
},
}
@@ -972,6 +973,10 @@ func TestGenerateChat(t *testing.T) {
t.Errorf("expected prompt eval count > 0, got 0")
}
if actual.PromptEvalCachedCount != 1 {
t.Errorf("expected prompt eval cached count 1, got %d", actual.PromptEvalCachedCount)
}
if actual.PromptEvalDuration == 0 {
t.Errorf("expected prompt eval duration > 0, got 0")
}
@@ -1471,12 +1476,13 @@ func TestGenerate(t *testing.T) {
mock := mockRunner{
CompletionResponse: llm.CompletionResponse{
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: 1,
PromptEvalDuration: 1,
EvalCount: 1,
EvalDuration: 1,
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: 1,
PromptEvalCachedCount: 1,
PromptEvalDuration: 1,
EvalCount: 1,
EvalDuration: 1,
},
}
@@ -1669,6 +1675,10 @@ func TestGenerate(t *testing.T) {
t.Errorf("expected prompt eval count > 0, got 0")
}
if actual.PromptEvalCachedCount != 1 {
t.Errorf("expected prompt eval cached count 1, got %d", actual.PromptEvalCachedCount)
}
if actual.PromptEvalDuration == 0 {
t.Errorf("expected prompt eval duration > 0, got 0")
}

26
x/mlxrunner/client.go
View File

@@ -107,10 +107,11 @@ type CompletionResponse struct {
Done bool
DoneReason int
PromptEvalCount int
PromptEvalDuration time.Duration
EvalCount int
EvalDuration time.Duration
PromptEvalCount int
PromptEvalCachedCount int
PromptEvalDuration time.Duration
EvalCount int
EvalDuration time.Duration
Logprobs []llm.Logprob
@@ -186,14 +187,15 @@ func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn f
}
cresp := llm.CompletionResponse{
Content: raw.Content,
Done: raw.Done,
DoneReason: llm.DoneReason(raw.DoneReason),
PromptEvalCount: raw.PromptEvalCount,
PromptEvalDuration: raw.PromptEvalDuration,
EvalCount: raw.EvalCount,
EvalDuration: raw.EvalDuration,
Logprobs: raw.Logprobs,
Content: raw.Content,
Done: raw.Done,
DoneReason: llm.DoneReason(raw.DoneReason),
PromptEvalCount: raw.PromptEvalCount,
PromptEvalCachedCount: raw.PromptEvalCachedCount,
PromptEvalDuration: raw.PromptEvalDuration,
EvalCount: raw.EvalCount,
EvalDuration: raw.EvalDuration,
Logprobs: raw.Logprobs,
}
fn(cresp)

7
x/mlxrunner/pipeline.go
View File

@@ -189,7 +189,12 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess
wantTopLogprobs: request.SamplerOpts.TopLogprobs,
}
final := CompletionResponse{Done: true, PromptEvalCount: len(request.Tokens), DoneReason: 1}
final := CompletionResponse{
Done: true,
PromptEvalCount: len(request.Tokens),
PromptEvalCachedCount: len(session.inputs) - len(session.remaining),
DoneReason: 1,
}
final.PromptEvalDuration = promptEval
now := time.Now()