diff --git a/apis/src/token_usage/mod.rs b/apis/src/token_usage/mod.rs index 5c4c30a5..70964c90 100644 --- a/apis/src/token_usage/mod.rs +++ b/apis/src/token_usage/mod.rs @@ -9,8 +9,6 @@ mod providers; mod streaming; -pub use streaming::extract_streaming_tokens; - #[cfg(test)] #[expect(clippy::allow_attributes, reason = "blanket test suppressions")] #[allow( @@ -79,7 +77,7 @@ impl TokenUsage { } /// AI provider identifier for response format selection. -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq)] +#[derive(Debug, Deserialize, Eq, PartialEq)] #[serde(rename_all = "lowercase")] pub enum TokenUsageProvider { /// `OpenAI` API (`usage.prompt_tokens`, `usage.completion_tokens`). @@ -119,28 +117,49 @@ pub fn set_token_usage(ctx: &mut HttpFilterContext<'_>, input: u64, output: u64, ctx.set_metadata("token.total", total.to_string()); } -/// Extracts token usage from a provider's JSON response body. -/// -/// Returns `None` if the response doesn't contain usage information -/// (e.g., error responses, malformed JSON, or missing fields). -/// -/// # Example -/// -/// ``` -/// use praxis_ai_apis::token_usage::{TokenUsageProvider, extract_token_usage}; -/// -/// let openai_response = -/// br#"{"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}}"#; -/// let usage = extract_token_usage(TokenUsageProvider::OpenAi, openai_response).unwrap(); -/// assert_eq!(usage.input_tokens(), 10); -/// assert_eq!(usage.output_tokens(), 20); -/// assert_eq!(usage.total_tokens(), 30); -/// ``` -pub fn extract_token_usage(provider: TokenUsageProvider, body: &[u8]) -> Option { - match provider { - TokenUsageProvider::OpenAi | TokenUsageProvider::Azure => parse_openai(body), - TokenUsageProvider::Anthropic => parse_anthropic(body), - TokenUsageProvider::Google => parse_google(body), - TokenUsageProvider::Bedrock => parse_bedrock(body), +impl TokenUsageProvider { + /// Extracts token usage from a provider's JSON response body. + /// + /// Returns `None` if the response doesn't contain usage information + /// (e.g., error responses, malformed JSON, or missing fields). + /// + /// # Example + /// + /// ``` + /// use praxis_ai_apis::token_usage::TokenUsageProvider; + /// + /// let openai_response = + /// br#"{"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}}"#; + /// let usage = TokenUsageProvider::OpenAi + /// .extract_token_usage(openai_response) + /// .unwrap(); + /// assert_eq!(usage.input_tokens(), 10); + /// assert_eq!(usage.output_tokens(), 20); + /// assert_eq!(usage.total_tokens(), 30); + /// ``` + pub fn extract_token_usage(&self, body: &[u8]) -> Option { + match *self { + TokenUsageProvider::OpenAi | TokenUsageProvider::Azure => parse_openai(body), + TokenUsageProvider::Anthropic => parse_anthropic(body), + TokenUsageProvider::Google => parse_google(body), + TokenUsageProvider::Bedrock => parse_bedrock(body), + } + } + + /// Extracts partial token counts from a single SSE event payload. + /// + /// Returns `(input, output)` where at least one field will be `Some` + /// when the event contains token data. Returns `(None, None)` when + /// the event contains no token usage data. + /// + /// Used as a fallback when [`Self::extract_token_usage`] returns `None` for + /// providers that distribute token counts across multiple events + /// (Anthropic, Bedrock). + pub fn extract_streaming_tokens(&self, event_data: &[u8]) -> (Option, Option) { + match *self { + TokenUsageProvider::Anthropic => streaming::parse_anthropic_event(event_data), + TokenUsageProvider::Bedrock => streaming::parse_bedrock_event(event_data), + TokenUsageProvider::OpenAi | TokenUsageProvider::Azure | TokenUsageProvider::Google => (None, None), + } } } diff --git a/apis/src/token_usage/streaming.rs b/apis/src/token_usage/streaming.rs index 5ec6c407..97f038c1 100644 --- a/apis/src/token_usage/streaming.rs +++ b/apis/src/token_usage/streaming.rs @@ -8,31 +8,6 @@ use serde::Deserialize; -use super::TokenUsageProvider; - -// ----------------------------------------------------------------------------- -// Public API -// ----------------------------------------------------------------------------- - -/// Extracts partial token counts from a single SSE event payload. -/// -/// Returns `(input, output)` where at least one field will be `Some` -/// when the event contains token data. Returns `(None, None)` when -/// the event contains no token usage data. -/// -/// Used as a fallback when [`extract_token_usage`] returns `None` for -/// providers that distribute token counts across multiple events -/// (Anthropic, Bedrock). -/// -/// [`extract_token_usage`]: super::extract_token_usage -pub fn extract_streaming_tokens(provider: TokenUsageProvider, event_data: &[u8]) -> (Option, Option) { - match provider { - TokenUsageProvider::Anthropic => parse_anthropic_event(event_data), - TokenUsageProvider::Bedrock => parse_bedrock_event(event_data), - TokenUsageProvider::OpenAi | TokenUsageProvider::Azure | TokenUsageProvider::Google => (None, None), - } -} - // ----------------------------------------------------------------------------- // Anthropic Streaming // ----------------------------------------------------------------------------- @@ -87,7 +62,7 @@ struct AnthropicDeltaUsage { } /// Parses Anthropic streaming events for partial token counts. -fn parse_anthropic_event(data: &[u8]) -> (Option, Option) { +pub(super) fn parse_anthropic_event(data: &[u8]) -> (Option, Option) { if let Ok(start) = serde_json::from_slice::(data) && start.event_type.as_deref() == Some("message_start") && let Some(message) = start.message @@ -142,7 +117,7 @@ struct BedrockStreamUsage { } /// Parses Bedrock `ConverseStream` metadata events for token counts. -fn parse_bedrock_event(data: &[u8]) -> (Option, Option) { +pub(super) fn parse_bedrock_event(data: &[u8]) -> (Option, Option) { let Some(meta) = serde_json::from_slice::(data).ok() else { return (None, None); }; diff --git a/apis/src/token_usage/tests.rs b/apis/src/token_usage/tests.rs index f1c627a7..c73dfaf9 100644 --- a/apis/src/token_usage/tests.rs +++ b/apis/src/token_usage/tests.rs @@ -3,7 +3,7 @@ //! Unit tests for token usage extraction. -use super::{TokenUsageProvider, extract_token_usage, set_token_usage, streaming::extract_streaming_tokens}; +use super::{TokenUsageProvider, set_token_usage}; // ----------------------------------------------------------------------------- // set_token_usage Tests @@ -88,7 +88,7 @@ fn openai_full_response() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::OpenAi, json).unwrap(); + let usage = TokenUsageProvider::OpenAi.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 15, "input_tokens should be 15"); assert_eq!(usage.output_tokens(), 42, "output_tokens should be 42"); @@ -99,7 +99,7 @@ fn openai_full_response() { fn azure_same_as_openai() { let json = br#"{"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}}"#; - let usage = extract_token_usage(TokenUsageProvider::Azure, json).unwrap(); + let usage = TokenUsageProvider::Azure.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 5, "Azure should parse like OpenAI"); assert_eq!(usage.output_tokens(), 10, "Azure should parse like OpenAI"); @@ -117,7 +117,7 @@ fn anthropic_full_response() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::Anthropic, json).unwrap(); + let usage = TokenUsageProvider::Anthropic.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 15, "input_tokens should be 15"); assert_eq!(usage.output_tokens(), 42, "output_tokens should be 42"); @@ -140,7 +140,7 @@ fn anthropic_with_prompt_caching() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::Anthropic, json).unwrap(); + let usage = TokenUsageProvider::Anthropic.extract_token_usage(json).unwrap(); assert_eq!( usage.input_tokens(), @@ -162,7 +162,7 @@ fn google_full_response() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::Google, json).unwrap(); + let usage = TokenUsageProvider::Google.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 15, "promptTokenCount should map to input_tokens"); assert_eq!( @@ -177,7 +177,7 @@ fn google_full_response() { fn google_without_total() { let json = br#"{"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 20}}"#; - let usage = extract_token_usage(TokenUsageProvider::Google, json).unwrap(); + let usage = TokenUsageProvider::Google.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 10, "promptTokenCount should map to input_tokens"); assert_eq!( @@ -204,7 +204,7 @@ fn bedrock_claude_invoke_model_response() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::Bedrock, json).unwrap(); + let usage = TokenUsageProvider::Bedrock.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 15, "input_tokens should map to input_tokens"); assert_eq!(usage.output_tokens(), 42, "output_tokens should map to output_tokens"); @@ -222,7 +222,7 @@ fn bedrock_converse_response() { } }"#; - let usage = extract_token_usage(TokenUsageProvider::Bedrock, json).unwrap(); + let usage = TokenUsageProvider::Bedrock.extract_token_usage(json).unwrap(); assert_eq!(usage.input_tokens(), 15, "inputTokens should map to input_tokens"); assert_eq!(usage.output_tokens(), 42, "outputTokens should map to output_tokens"); @@ -237,7 +237,7 @@ fn bedrock_converse_response() { fn missing_usage_returns_none() { let json = br#"{"id": "123", "choices": []}"#; - let result = extract_token_usage(TokenUsageProvider::OpenAi, json); + let result = TokenUsageProvider::OpenAi.extract_token_usage(json); assert!(result.is_none(), "missing usage field should return None"); } @@ -246,7 +246,7 @@ fn missing_usage_returns_none() { fn google_missing_usage_returns_none() { let json = br#"{"candidates": []}"#; - let result = extract_token_usage(TokenUsageProvider::Google, json); + let result = TokenUsageProvider::Google.extract_token_usage(json); assert!(result.is_none(), "missing usageMetadata field should return None"); } @@ -255,14 +255,14 @@ fn google_missing_usage_returns_none() { fn invalid_json_returns_none() { let json = b"not valid json"; - let result = extract_token_usage(TokenUsageProvider::OpenAi, json); + let result = TokenUsageProvider::OpenAi.extract_token_usage(json); assert!(result.is_none(), "invalid JSON should return None"); } #[test] fn empty_body_returns_none() { - let result = extract_token_usage(TokenUsageProvider::OpenAi, b""); + let result = TokenUsageProvider::OpenAi.extract_token_usage(b""); assert!(result.is_none(), "empty body should return None"); } @@ -271,7 +271,7 @@ fn empty_body_returns_none() { fn error_response_returns_none() { let json = br#"{"error": {"message": "Invalid API key", "type": "invalid_request_error"}}"#; - let result = extract_token_usage(TokenUsageProvider::OpenAi, json); + let result = TokenUsageProvider::OpenAi.extract_token_usage(json); assert!(result.is_none(), "error response should return None"); } @@ -284,7 +284,7 @@ fn error_response_returns_none() { fn anthropic_message_start_yields_input_tokens() { let event = br#"{"type":"message_start","message":{"id":"msg_01","type":"message","role":"assistant","content":[],"usage":{"input_tokens":25}}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!(result, (Some(25), None), "message_start should yield input tokens only"); } @@ -293,7 +293,7 @@ fn anthropic_message_start_yields_input_tokens() { fn anthropic_message_start_with_caching() { let event = br#"{"type":"message_start","message":{"usage":{"input_tokens":10,"cache_creation_input_tokens":100,"cache_read_input_tokens":500}}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!( result, @@ -306,7 +306,7 @@ fn anthropic_message_start_with_caching() { fn anthropic_message_delta_yields_output_tokens() { let event = br#"{"type":"message_delta","usage":{"output_tokens":42}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!( result, @@ -319,7 +319,7 @@ fn anthropic_message_delta_yields_output_tokens() { fn anthropic_content_block_delta_returns_none() { let event = br#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!(result, (None, None), "content_block_delta has no token data"); } @@ -328,7 +328,7 @@ fn anthropic_content_block_delta_returns_none() { fn openai_streaming_returns_none() { let event = br#"{"id":"chatcmpl-abc","choices":[{"delta":{"content":"Hi"}}]}"#; - let result = extract_streaming_tokens(TokenUsageProvider::OpenAi, event); + let result = TokenUsageProvider::OpenAi.extract_streaming_tokens(event); assert_eq!( result, @@ -341,7 +341,7 @@ fn openai_streaming_returns_none() { fn azure_streaming_returns_none() { let event = br#"{"id":"chatcmpl-abc","choices":[{"delta":{"content":"Hi"}}]}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Azure, event); + let result = TokenUsageProvider::Azure.extract_streaming_tokens(event); assert_eq!( result, @@ -354,7 +354,7 @@ fn azure_streaming_returns_none() { fn google_streaming_returns_none() { let event = br#"{"candidates":[{"content":{"parts":[{"text":"Hi"}]}}]}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Google, event); + let result = TokenUsageProvider::Google.extract_streaming_tokens(event); assert_eq!( result, @@ -367,7 +367,7 @@ fn google_streaming_returns_none() { fn bedrock_converse_stream_metadata_yields_both() { let event = br#"{"metadata":{"usage":{"inputTokens":15,"outputTokens":42}}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Bedrock, event); + let result = TokenUsageProvider::Bedrock.extract_streaming_tokens(event); assert_eq!( result, @@ -380,14 +380,14 @@ fn bedrock_converse_stream_metadata_yields_both() { fn bedrock_content_block_returns_none() { let event = br#"{"contentBlockDelta":{"delta":{"text":"Hi"},"contentBlockIndex":0}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Bedrock, event); + let result = TokenUsageProvider::Bedrock.extract_streaming_tokens(event); assert_eq!(result, (None, None), "Bedrock content delta has no token metadata"); } #[test] fn streaming_invalid_json_returns_none() { - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, b"not json"); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(b"not json"); assert_eq!( result, @@ -398,7 +398,7 @@ fn streaming_invalid_json_returns_none() { #[test] fn streaming_empty_returns_none() { - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, b""); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(b""); assert_eq!( result, @@ -444,7 +444,7 @@ fn provider_rejects_unknown() { fn anthropic_missing_type_field_returns_none() { let event = br#"{"message":{"usage":{"input_tokens":25}}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!( result, @@ -457,7 +457,7 @@ fn anthropic_missing_type_field_returns_none() { fn anthropic_wrong_type_with_usage_returns_none() { let event = br#"{"type":"message_stop","message":{"usage":{"input_tokens":25}}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!(result, (None, None), "wrong type should not match message_start"); } @@ -466,7 +466,7 @@ fn anthropic_wrong_type_with_usage_returns_none() { fn anthropic_message_start_with_null_message() { let event = br#"{"type":"message_start","message":null}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!(result, (None, None), "null message should not yield tokens"); } @@ -475,7 +475,7 @@ fn anthropic_message_start_with_null_message() { fn anthropic_message_start_with_null_usage() { let event = br#"{"type":"message_start","message":{"usage":null}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Anthropic, event); + let result = TokenUsageProvider::Anthropic.extract_streaming_tokens(event); assert_eq!(result, (None, None), "null usage should not yield tokens"); } @@ -484,7 +484,7 @@ fn anthropic_message_start_with_null_usage() { fn bedrock_null_metadata() { let event = br#"{"metadata":null}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Bedrock, event); + let result = TokenUsageProvider::Bedrock.extract_streaming_tokens(event); assert_eq!(result, (None, None), "null metadata should not yield tokens"); } @@ -493,7 +493,7 @@ fn bedrock_null_metadata() { fn bedrock_null_usage_inside_metadata() { let event = br#"{"metadata":{"usage":null}}"#; - let result = extract_streaming_tokens(TokenUsageProvider::Bedrock, event); + let result = TokenUsageProvider::Bedrock.extract_streaming_tokens(event); assert_eq!( result, diff --git a/filters/src/token_count/mod.rs b/filters/src/token_count/mod.rs index c490ff32..2ad7223f 100644 --- a/filters/src/token_count/mod.rs +++ b/filters/src/token_count/mod.rs @@ -29,7 +29,7 @@ use std::fmt::Write as _; use async_trait::async_trait; use bytes::Bytes; -use praxis_ai_apis::token_usage::{TokenUsageProvider, extract_streaming_tokens, extract_token_usage, set_token_usage}; +use praxis_ai_apis::token_usage::{TokenUsageProvider, set_token_usage}; use praxis_filter::{ BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, parse_filter_config, }; @@ -262,8 +262,8 @@ impl HttpFilter for TokenCountFilter { let mode = ctx.get_metadata(META_MODE).map(str::to_owned); match mode.as_deref() { - Some("sse") => handle_sse_body(ctx, body, end_of_stream, provider), - Some("json") => handle_json_body(ctx, body, end_of_stream, provider), + Some("sse") => handle_sse_body(ctx, body, end_of_stream, &provider), + Some("json") => handle_json_body(ctx, body, end_of_stream, &provider), _ => {}, } @@ -314,7 +314,7 @@ fn handle_json_body( ctx: &mut HttpFilterContext<'_>, body: &Option, end_of_stream: bool, - provider: TokenUsageProvider, + provider: &TokenUsageProvider, ) { if let Some(chunk) = body.as_ref() && !accumulate_response_hex(ctx, chunk, DEFAULT_MAX_BODY_BYTES) @@ -327,7 +327,7 @@ fn handle_json_body( let bytes = ctx.filter_metadata.get(META_BUF_HEX).and_then(|hex| decode_hex(hex)); if let Some(data) = bytes - && let Some(usage) = extract_token_usage(provider, &data) + && let Some(usage) = provider.extract_token_usage(&data) { set_token_usage( ctx, @@ -356,7 +356,7 @@ fn handle_sse_body( ctx: &mut HttpFilterContext<'_>, body: &Option, end_of_stream: bool, - provider: TokenUsageProvider, + provider: &TokenUsageProvider, ) { if let Some(chunk) = body.as_ref() { let mut state = load_sse_scan_state(ctx); @@ -387,7 +387,7 @@ fn handle_sse_body( } /// Try to extract token usage from a single SSE data payload. -fn process_sse_payload(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: TokenUsageProvider) { +fn process_sse_payload(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: &TokenUsageProvider) { if payload == b"[DONE]" { return; } @@ -400,8 +400,8 @@ fn process_sse_payload(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider } /// Try complete usage extraction (OpenAI, Google, Azure final events). -fn try_complete_usage(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: TokenUsageProvider) -> bool { - let Some(usage) = extract_token_usage(provider, payload) else { +fn try_complete_usage(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: &TokenUsageProvider) -> bool { + let Some(usage) = provider.extract_token_usage(payload) else { return false; }; @@ -418,8 +418,8 @@ fn try_complete_usage(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: } /// Try partial extraction (Anthropic, Bedrock streaming). -fn try_partial_usage(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: TokenUsageProvider) { - let (input, output) = extract_streaming_tokens(provider, payload); +fn try_partial_usage(ctx: &mut HttpFilterContext<'_>, payload: &[u8], provider: &TokenUsageProvider) { + let (input, output) = provider.extract_streaming_tokens(payload); if let Some(inp) = input { merge_accumulated_count(ctx, META_INPUT, inp);