Expose detailed completion token usage

This commit is contained in:
Bryan Helmkamp 2026-07-24 07:36:56 -04:00
parent df3c834581
commit 142862f342
No known key found for this signature in database
7 changed files with 179 additions and 13 deletions

View file

@ -8507,15 +8507,38 @@ components:
description: Provider-specific options.
CompletionUsage:
description: >
Five disjoint token buckets for one completion. `input_tokens` excludes
cache reads and writes, while `output_tokens` excludes reasoning tokens
when the provider reports them separately.
type: object
required: [input_tokens, output_tokens]
required:
- input_tokens
- output_tokens
- reasoning_tokens
- cache_read_tokens
- cache_write_tokens
properties:
input_tokens:
type: integer
format: int64
description: Number of uncached input tokens consumed.
output_tokens:
type: integer
format: int64
description: Number of non-reasoning output tokens generated.
reasoning_tokens:
type: integer
format: int64
description: Number of separately reported reasoning tokens.
cache_read_tokens:
type: integer
format: int64
description: Number of input tokens served from a provider cache.
cache_write_tokens:
type: integer
format: int64
description: Number of input tokens written to a provider cache.
CompletionResponse:
type: object

View file

@ -4,10 +4,10 @@ use std::sync::Arc;
use fabro_model::{Catalog, ModelSelectionError};
use super::super::{
ApiError, AppState, CompletionResponse, CompletionToolChoiceMode, CompletionUsage,
CreateCompletionRequest, FinishReason, GenerateParams, IntoResponse, Json, LlmMessage,
LlmRequest, ProviderId, RequiredUser, Response, Router, State, StatusCode, ToolChoice,
ToolDefinition, Ulid, error, generate_object, info, post, warn,
ApiError, AppState, CompletionResponse, CompletionToolChoiceMode, CreateCompletionRequest,
FinishReason, GenerateParams, IntoResponse, Json, LlmMessage, LlmRequest, ProviderId,
RequiredUser, Response, Router, State, StatusCode, ToolChoice, ToolDefinition, Ulid, error,
generate_object, info, post, warn,
};
use super::llm_sse;
@ -169,10 +169,7 @@ async fn create_completion(
provider: selected_provider,
message: response.message,
stop_reason,
usage: CompletionUsage {
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
},
usage: response.usage,
output,
cost_usd: response.cost_usd,
cost_source: response.cost_source,
@ -192,10 +189,7 @@ async fn create_completion(
provider: ProviderId::new(response.provider),
message: response.message,
stop_reason,
usage: CompletionUsage {
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
},
usage: response.usage,
output: None,
cost_usd: response.cost_usd,
cost_source: response.cost_source,

View file

@ -15272,6 +15272,73 @@ async fn create_completion_unknown_provider_returns_clear_error() {
);
}
#[tokio::test]
async fn create_completion_returns_disjoint_usage_buckets() {
let upstream = MockServer::start();
let completion = upstream.mock(|when, then| {
when.method(POST).path("/chat/completions");
then.status(200)
.header("content-type", "application/json")
.json_body(json!({
"id": "chatcmpl-usage",
"model": "kimi-k3",
"choices": [{
"message": {"role": "assistant", "content": "OK"},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 200,
"completion_tokens": 30,
"total_tokens": 230,
"prompt_tokens_details": {
"cached_tokens": 50,
"cache_write_tokens": 100
},
"completion_tokens_details": {
"reasoning_tokens": 20
}
}
}));
});
let state = TestAppStateBuilder::new()
.provider_base_url("kimi", upstream.base_url())
.vault_entries([(EnvVars::KIMI_API_KEY, "test-kimi-api-key")])
.build();
let app = crate::test_support::build_test_router(state);
let req = Request::builder()
.method("POST")
.uri(api("/completions"))
.header("content-type", "application/json")
.body(Body::from(
json!({
"provider": "kimi",
"model": "kimi-k3",
"stream": false,
"messages": [{
"role": "user",
"content": [{"kind": "text", "data": "hi"}]
}]
})
.to_string(),
))
.unwrap();
let response = app.oneshot(req).await.unwrap();
let body = response_json!(response, StatusCode::OK).await;
assert_eq!(
body["usage"],
json!({
"input_tokens": 50,
"output_tokens": 10,
"reasoning_tokens": 20,
"cache_read_tokens": 50,
"cache_write_tokens": 100
})
);
completion.assert();
}
#[tokio::test]
async fn create_completion_default_model_uses_app_state_catalog() {
let upstream = MockServer::start();

View file

@ -461,6 +461,7 @@ fn main() {
"fabro_types::PendingInterviewRecord",
&[],
),
("CompletionUsage", "fabro_model::TokenCounts", &[]),
("BilledTokenCounts", "fabro_types::BilledTokenCounts", &[]),
("BillingModelRef", "fabro_model::ModelRef", &[]),
("BillingSpeed", "fabro_model::Speed", &[]),

View file

@ -22,6 +22,7 @@ pub mod types {
pub use fabro_model::{
CostSource, Model, ModelCosts, ModelFeatures, ModelLimits, ModelRef as BillingModelRef,
ModelTestMode, Provider, ReasoningEffort, ReasoningEffortFeature, Speed as BillingSpeed,
TokenCounts as CompletionUsage,
};
pub use fabro_types::run_event::AgentSessionActivatedProps;
pub use fabro_types::settings::run::McpHttpProtocol;

View file

@ -0,0 +1,59 @@
use std::any::{TypeId, type_name};
use fabro_api::types::CompletionUsage as ApiCompletionUsage;
use fabro_model::TokenCounts;
use serde_json::json;
#[test]
fn completion_usage_reuses_canonical_type() {
assert_same_type::<ApiCompletionUsage, TokenCounts>();
}
#[test]
fn completion_usage_json_matches_openapi_shape() {
let usage = TokenCounts {
input_tokens: 10,
output_tokens: 20,
reasoning_tokens: 3,
cache_read_tokens: 4,
cache_write_tokens: 5,
};
let json = serde_json::to_value(&usage).unwrap();
assert_eq!(json["input_tokens"], 10);
assert_eq!(json["output_tokens"], 20);
assert_eq!(json["reasoning_tokens"], 3);
assert_eq!(json["cache_read_tokens"], 4);
assert_eq!(json["cache_write_tokens"], 5);
let round_trip: ApiCompletionUsage = serde_json::from_value(json).unwrap();
assert_eq!(round_trip, usage);
}
#[test]
fn completion_usage_keeps_zero_counts_present() {
let json = serde_json::to_value(TokenCounts::default()).unwrap();
assert_eq!(
json,
json!({
"input_tokens": 0,
"output_tokens": 0,
"reasoning_tokens": 0,
"cache_read_tokens": 0,
"cache_write_tokens": 0
})
);
let round_trip: ApiCompletionUsage = serde_json::from_value(json).unwrap();
assert_eq!(round_trip, TokenCounts::default());
}
fn assert_same_type<T: 'static, U: 'static>() {
assert_eq!(
TypeId::of::<T>(),
TypeId::of::<U>(),
"{} should be the same type as {}",
type_name::<T>(),
type_name::<U>()
);
}

View file

@ -14,7 +14,28 @@
/**
* Five disjoint token buckets for one completion. `input_tokens` excludes cache reads and writes, while `output_tokens` excludes reasoning tokens when the provider reports them separately.
*/
export interface CompletionUsage {
/**
* Number of uncached input tokens consumed.
*/
'input_tokens': number;
/**
* Number of non-reasoning output tokens generated.
*/
'output_tokens': number;
/**
* Number of separately reported reasoning tokens.
*/
'reasoning_tokens': number;
/**
* Number of input tokens served from a provider cache.
*/
'cache_read_tokens': number;
/**
* Number of input tokens written to a provider cache.
*/
'cache_write_tokens': number;
}