Skip to content

feat(server): count hybrid upstream generation with the model's tokenizer #620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
krisztian-gajdar merged 1 commit into from
Oct 8, 2026
Merged

feat(server): count hybrid upstream generation with the model's tokenizer #620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions packages/sie_gateway/src/handlers/proxy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17335,6 +17335,7 @@ mod tests {
text: "ok".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -21196,6 +21197,7 @@ mod tests {
text: "Hello world!".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -26157,6 +26159,7 @@ mod tests {
text: String::new(),
finish_reason: "stop".to_string(),
usage: Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -26215,6 +26218,7 @@ mod tests {
text: String::new(),
finish_reason: "tool_calls".to_string(),
usage: Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -26285,6 +26289,7 @@ mod tests {
text: String::new(),
finish_reason: "stop".to_string(),
usage: Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -26670,6 +26675,7 @@ mod tests {
text: "a continuation".to_string(),
finish_reason: "length".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -27296,6 +27302,7 @@ mod tests {
text: "a joke".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -27331,6 +27338,7 @@ mod tests {
#[test]
fn test_responses_usage_reports_cached_input_tokens() {
let usage = crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: Some(crate::queue::streaming::PromptTokensDetails {
Expand Down Expand Up @@ -27380,6 +27388,7 @@ mod tests {
text: "Hi there!".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -27473,6 +27482,7 @@ mod tests {
text: "Hi".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -27513,6 +27523,7 @@ mod tests {
text: "Hi".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -27548,6 +27559,7 @@ mod tests {
text: String::new(),
finish_reason: "tool_calls".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -28205,6 +28217,7 @@ mod tests {
text: "Hi".to_string(),
finish_reason: "stop".to_string(),
usage: Some(crate::queue::streaming::UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down
11 changes: 11 additions & 0 deletions packages/sie_gateway/src/handlers/sse.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1994,6 +1994,7 @@ mod tests {
// whose usage is the count-so-far.
finish_reason: "cancelled".to_string(),
usage: Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -2304,6 +2305,7 @@ mod tests {
let mut terminal = _terminal_chunk("error", None);
terminal.seq = 42;
terminal.usage = Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand All @@ -2330,6 +2332,7 @@ mod tests {
let mut terminal = _terminal_chunk("error", None);
terminal.seq = 42;
terminal.usage = Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3058,6 +3061,7 @@ mod tests {
let chunk = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3478,6 +3482,7 @@ mod tests {
collector.apply(_terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3516,6 +3521,7 @@ mod tests {
let terminal = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3639,6 +3645,7 @@ mod tests {
let terminal = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3687,6 +3694,7 @@ mod tests {
let terminal = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3755,6 +3763,7 @@ mod tests {
let terminal = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3792,6 +3801,7 @@ mod tests {
let terminal = _terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -3927,6 +3937,7 @@ mod tests {
collector.apply(_terminal_chunk(
"stop",
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down
89 changes: 89 additions & 0 deletions packages/sie_gateway/src/queue/streaming.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,26 @@ pub struct UsageBlock {
/// engine does not report prefix-cache hits (and from older workers).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub prompt_tokens_details: Option<PromptTokensDetails>,
/// The upstream's own counts for a generation it served, when the counts
/// above are the worker's count with the model's tokenizer. Never
/// serialized, so no response surface shows it.
///
/// `dead_code`-allowed because the `sie-gateway` binary compiles this
/// module tree independently of the library (see the note in `lib.rs`),
/// and only metering consumers of the stream outcome read it.
#[allow(dead_code)]
#[serde(default, skip_serializing)]
pub upstream_usage: Option<UpstreamTokenUsage>,
}

/// Token counts an upstream reported for a generation it served.
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct UpstreamTokenUsage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
/// Clamped to `prompt_tokens` when decoded.
#[serde(default)]
pub cached_tokens: Option<u32>,
}

impl UsageBlock {
Expand Down Expand Up @@ -240,6 +260,8 @@ struct WorkerUsageBlock {
gpu_second: Option<u64>,
#[serde(default)]
prompt_tokens_details: Option<PromptTokensDetails>,
#[serde(default)]
upstream_usage: Option<UpstreamTokenUsage>,
}

impl From<WorkerUsageBlock> for UsageBlock {
Expand All @@ -249,13 +271,20 @@ impl From<WorkerUsageBlock> for UsageBlock {
.map(|details| PromptTokensDetails {
cached_tokens: details.cached_tokens.min(raw.prompt_tokens),
});
let upstream_usage = raw.upstream_usage.map(|upstream| UpstreamTokenUsage {
cached_tokens: upstream
.cached_tokens
.map(|cached| cached.min(upstream.prompt_tokens)),
..upstream
});
Self {
prompt_tokens: raw.prompt_tokens,
completion_tokens: raw.completion_tokens,
total_tokens: raw.total_tokens,
images: raw.images,
gpu_second: raw.gpu_second,
prompt_tokens_details,
upstream_usage,
}
}
}
Expand Down Expand Up @@ -1263,6 +1292,7 @@ mod tests {
finish_reason: if done { Some("stop".to_string()) } else { None },
usage: if done {
Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down Expand Up @@ -1400,6 +1430,64 @@ mod tests {
}
}

#[test]
fn test_upstream_usage_is_decoded_clamped_and_never_serialized() {
let decode = |upstream_usage: Option<serde_json::Value>| {
let mut usage = serde_json::json!({
"prompt_tokens": 2,
"completion_tokens": 5,
"total_tokens": 7,
});
if let Some(value) = upstream_usage {
usage["upstream_usage"] = value;
}
let bytes = rmp_serde::to_vec_named(&usage).expect("encode usage block");
rmp_serde::from_slice::<UsageBlock>(&bytes)
};

let counted = decode(Some(serde_json::json!({
"prompt_tokens": 37,
"completion_tokens": 9,
"cached_tokens": 40,
})))
.expect("upstream usage decodes");
assert_eq!(
counted.upstream_usage,
Some(UpstreamTokenUsage {
prompt_tokens: 37,
completion_tokens: 9,
cached_tokens: Some(37),
})
);
assert_eq!(counted.prompt_tokens, 2);
let serialized = serde_json::to_value(&counted).expect("serialize usage");
assert!(serialized.get("upstream_usage").is_none());
assert!(!serialized.to_string().contains("37"));

let uncached = decode(Some(serde_json::json!({
"prompt_tokens": 11,
"completion_tokens": 4,
})))
.expect("upstream usage without a cache count decodes");
assert_eq!(
uncached
.upstream_usage
.and_then(|usage| usage.cached_tokens),
None
);

let legacy = decode(None).expect("usage without an upstream figure decodes");
assert_eq!(legacy.upstream_usage, None);

for invalid in [
serde_json::json!({"prompt_tokens": -1, "completion_tokens": 1}),
serde_json::json!({"prompt_tokens": 1}),
serde_json::json!("37"),
] {
assert!(decode(Some(invalid)).is_err());
}
}

#[test]
fn test_worker_error_public_contract_fails_closed_for_parameters() {
let public_message = "PUBLIC_MESSAGE";
Expand Down Expand Up @@ -1750,6 +1838,7 @@ mod tests {
collector.last_output_at = Some(first + std::time::Duration::from_millis(400));
collector.output_event_count = 2;
collector.final_meta.as_mut().expect("terminal").usage = Some(UsageBlock {
upstream_usage: None,
gpu_second: None,
images: None,
prompt_tokens_details: None,
Expand Down
14 changes: 14 additions & 0 deletions packages/sie_server/REMOTE_BACKENDS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,20 @@ upstream's `usage` without `input_tokens`. Generation fails closed: a
generation, chat or completion response without exact final usage is an
error, never a success.

Queued generation for a model with local weights, served by an `openai`
upstream, is counted as if it had run locally. The worker renders the model's
chat template and counts the prompt with the model's tokenizer. It counts the
completion from every text the upstream returned, private reasoning included,
with the same tokenizer. The terminal chunk reports that count without cached
tokens, and carries the upstream's own counts in `usage.upstream_usage`. The
gateway decodes that field for consumers of its stream outcome and never
returns it to a caller. A request whose messages the local template cannot
render, or whose counted prompt exceeds the context length, fails before it is
sent. When the local tokenizer does not load, or has no chat template for a
chat request, the upstream's counts are reported instead. An `sie` upstream
serves the same model and counts with its tokenizer, so its counts are reported
as they are. Single-node serving reports the upstream's counts.

## OpenAI-compatible embeddings and rerank

Define a `kind: openai` upstream with a base URL that includes the provider's
Expand Down
Loading
Loading