diff --git a/docs/filters/token_count.md b/docs/filters/token_count.md index 21db517b..fe2b0228 100644 --- a/docs/filters/token_count.md +++ b/docs/filters/token_count.md @@ -7,17 +7,17 @@ Extracts token usage from AI inference responses and writes unified counts to [` ## Configuration Notes -Supports both streaming (SSE) and non-streaming (JSON) responses across all five providers (OpenAI, Anthropic, Google, Bedrock, Azure). +Supports both streaming (SSE) and non-streaming (JSON) responses across five providers (OpenAI, Anthropic, Google, Bedrock Converse, Azure), plus a header-only extraction path for Bedrock `InvokeModel`. ## Configuration | Field | Type | Required | Description | |-------|------|---------|-------------| -| `provider` | TokenUsageProvider | yes | AI provider whose response format to parse. | +| `provider` | `openai` \| `anthropic` \| `google` \| `bedrock` \| `bedrock_invoke_model` \| `azure` | yes | AI provider whose response format to parse. | ## Example ```yaml filter: token_count -provider: openai +provider: openai # openai | anthropic | google | bedrock | bedrock_invoke_model | azure ``` diff --git a/examples/configs/token-usage-headers.yaml b/examples/configs/token-usage-headers.yaml index 8a5ec3c2..aa47ee90 100644 --- a/examples/configs/token-usage-headers.yaml +++ b/examples/configs/token-usage-headers.yaml @@ -12,9 +12,12 @@ # filter metadata and injects them as response headers. When no token # data is present the filter is a no-op. # -# NOTE: A token-counting filter (or any filter that calls -# set_token_usage()) must precede token_usage_headers in the pipeline -# for headers to appear. Without it the filter is always a no-op. +# NOTE: response hooks run in reverse declared order (the last filter's +# on_response fires first), so a token-counting filter (or any filter +# that calls set_token_usage()) must be declared AFTER token_usage_headers +# in the pipeline for headers to appear. Without it the filter is always +# a no-op. See token-counting.yaml for a filter that populates +# token.input/token.output/token.total. listeners: - name: default diff --git a/filters/src/token_count/mod.rs b/filters/src/token_count/mod.rs index 677e1962..c490ff32 100644 --- a/filters/src/token_count/mod.rs +++ b/filters/src/token_count/mod.rs @@ -8,6 +8,10 @@ //! for downstream consumers. The filter is transparent: response bodies //! and status codes pass through unchanged. //! +//! Bedrock `InvokeModel` is the one exception: it reports token counts as +//! HTTP response headers rather than in the response body, so it is read +//! directly in `on_response` with no body buffering at all. +//! //! [`filter_metadata`]: HttpFilterContext::filter_metadata #[cfg(test)] @@ -77,6 +81,12 @@ const META_SSE_PREV_CR: &str = "token_count.sse_prev_cr"; /// Metadata key for SSE scanner scratch byte count. const META_SSE_SCRATCH: &str = "token_count.sse_scratch_bytes"; +/// Bedrock `InvokeModel` response header carrying the input token count. +const HEADER_BEDROCK_INPUT: &str = "x-amzn-bedrock-input-token-count"; + +/// Bedrock `InvokeModel` response header carrying the output token count. +const HEADER_BEDROCK_OUTPUT: &str = "x-amzn-bedrock-output-token-count"; + // ----------------------------------------------------------------------------- // Configuration // ----------------------------------------------------------------------------- @@ -86,7 +96,59 @@ const META_SSE_SCRATCH: &str = "token_count.sse_scratch_bytes"; #[serde(deny_unknown_fields)] struct TokenCountConfig { /// AI provider whose response format to parse. - provider: TokenUsageProvider, + provider: ProviderKind, +} + +/// AI provider selecting the token extraction strategy. +/// +/// Distinct from [`TokenUsageProvider`] because Bedrock `InvokeModel` has no +/// body format to parse — its counts live in response headers — so it needs +/// a variant that [`TokenUsageProvider`] intentionally has no equivalent for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +enum ProviderKind { + /// OpenAI Chat Completions API. + #[serde(rename = "openai")] + OpenAi, + /// Anthropic Claude API. + Anthropic, + /// Google Gemini API. + Google, + /// AWS Bedrock Converse API (JSON body / SSE, like other providers). + Bedrock, + /// AWS Bedrock `InvokeModel` API (token counts in response headers only). + BedrockInvokeModel, + /// Azure OpenAI (same JSON schema as OpenAI). + Azure, +} + +impl ProviderKind { + /// Maps to the shared library provider used for body-based extraction. + /// + /// Returns `None` for `BedrockInvokeModel`, which has no body format + /// and is instead handled directly via response headers. + fn to_library_provider(self) -> Option { + match self { + Self::OpenAi => Some(TokenUsageProvider::OpenAi), + Self::Anthropic => Some(TokenUsageProvider::Anthropic), + Self::Google => Some(TokenUsageProvider::Google), + Self::Bedrock => Some(TokenUsageProvider::Bedrock), + Self::Azure => Some(TokenUsageProvider::Azure), + Self::BedrockInvokeModel => None, + } + } +} + +impl From for ProviderKind { + fn from(provider: TokenUsageProvider) -> Self { + match provider { + TokenUsageProvider::OpenAi => Self::OpenAi, + TokenUsageProvider::Anthropic => Self::Anthropic, + TokenUsageProvider::Google => Self::Google, + TokenUsageProvider::Bedrock => Self::Bedrock, + TokenUsageProvider::Azure => Self::Azure, + } + } } // ----------------------------------------------------------------------------- @@ -96,20 +158,21 @@ struct TokenCountConfig { /// Extracts token usage from AI inference responses and writes unified /// counts to [`filter_metadata`]. /// -/// Supports both streaming (SSE) and non-streaming (JSON) responses -/// across all five providers (OpenAI, Anthropic, Google, Bedrock, Azure). +/// Supports both streaming (SSE) and non-streaming (JSON) responses across +/// five providers (OpenAI, Anthropic, Google, Bedrock Converse, Azure), plus +/// a header-only extraction path for Bedrock `InvokeModel`. /// /// # YAML /// /// ```yaml /// filter: token_count -/// provider: openai +/// provider: openai # openai | anthropic | google | bedrock | bedrock_invoke_model | azure /// ``` /// /// [`filter_metadata`]: HttpFilterContext::filter_metadata pub struct TokenCountFilter { /// Which provider's response format to parse. - provider: TokenUsageProvider, + provider: ProviderKind, } impl TokenCountFilter { @@ -132,7 +195,11 @@ impl HttpFilter for TokenCountFilter { } fn response_body_access(&self) -> BodyAccess { - BodyAccess::ReadOnly + if self.provider == ProviderKind::BedrockInvokeModel { + BodyAccess::None + } else { + BodyAccess::ReadOnly + } } fn response_body_mode(&self) -> BodyMode { @@ -143,6 +210,11 @@ impl HttpFilter for TokenCountFilter { Ok(FilterAction::Continue) } + /// Skips extraction entirely for non-success statuses, for every + /// provider. For `bedrock_invoke_model`, reads token counts directly + /// from response headers since that provider has no body format to + /// parse. For all other providers, detects the content-type to select + /// the body extraction strategy used by `on_response_body`. async fn on_response(&self, ctx: &mut HttpFilterContext<'_>) -> Result { let is_success = ctx.response_header.as_ref().is_some_and(|r| r.status.is_success()); @@ -151,6 +223,11 @@ impl HttpFilter for TokenCountFilter { return Ok(FilterAction::Continue); } + if self.provider == ProviderKind::BedrockInvokeModel { + extract_bedrock_headers(ctx); + return Ok(FilterAction::Continue); + } + let content_type = ctx .response_header .as_ref() @@ -178,11 +255,15 @@ impl HttpFilter for TokenCountFilter { body: &mut Option, end_of_stream: bool, ) -> Result { + let Some(provider) = self.provider.to_library_provider() else { + return Ok(FilterAction::Continue); + }; + 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, self.provider), - Some("json") => handle_json_body(ctx, body, end_of_stream, self.provider), + Some("sse") => handle_sse_body(ctx, body, end_of_stream, provider), + Some("json") => handle_json_body(ctx, body, end_of_stream, provider), _ => {}, } @@ -190,6 +271,40 @@ impl HttpFilter for TokenCountFilter { } } +// ----------------------------------------------------------------------------- +// Bedrock InvokeModel (Header-Only) Path +// ----------------------------------------------------------------------------- + +/// Reads Bedrock `InvokeModel` token counts from response headers; no-op if +/// either header is absent or unparseable. Unlike every other provider, no +/// body access is required for this path. +fn extract_bedrock_headers(ctx: &mut HttpFilterContext<'_>) { + let input = ctx + .response_header + .as_ref() + .and_then(|r| r.headers.get(HEADER_BEDROCK_INPUT)) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()); + + let output = ctx + .response_header + .as_ref() + .and_then(|r| r.headers.get(HEADER_BEDROCK_OUTPUT)) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()); + + let (Some(input), Some(output)) = (input, output) else { + trace!("Bedrock InvokeModel token headers not present or unparseable"); + return; + }; + + set_token_usage(ctx, input, output, None); + debug!( + input, + output, "extracted Bedrock InvokeModel token counts from response headers" + ); +} + // ----------------------------------------------------------------------------- // JSON (Non-Streaming) Path // ----------------------------------------------------------------------------- diff --git a/filters/src/token_count/tests.rs b/filters/src/token_count/tests.rs index ef255c4d..c22acbd5 100644 --- a/filters/src/token_count/tests.rs +++ b/filters/src/token_count/tests.rs @@ -23,7 +23,14 @@ fn from_config_with_valid_provider() { #[test] fn from_config_all_providers() { - for provider in ["openai", "anthropic", "google", "bedrock", "azure"] { + for provider in [ + "openai", + "anthropic", + "google", + "bedrock", + "bedrock_invoke_model", + "azure", + ] { let yaml = format!("provider: {provider}"); let config: serde_yaml::Value = serde_yaml::from_str(&yaml).unwrap(); let result = TokenCountFilter::from_config(&config); @@ -593,6 +600,170 @@ fn decode_hex_rejects_invalid_chars() { assert!(decode_hex("zz").is_none(), "invalid hex chars should return None"); } +// ----------------------------------------------------------------------------- +// Bedrock InvokeModel: Header-Only Path +// ----------------------------------------------------------------------------- + +#[tokio::test] +async fn bedrock_invoke_model_extracts_headers_on_response() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = crate::test_utils::make_response(); + resp.headers.insert(HEADER_BEDROCK_INPUT, "25".parse().unwrap()); + resp.headers.insert(HEADER_BEDROCK_OUTPUT, "50".parse().unwrap()); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert_eq!( + ctx.get_metadata("token.input"), + Some("25"), + "Bedrock InvokeModel input tokens" + ); + assert_eq!( + ctx.get_metadata("token.output"), + Some("50"), + "Bedrock InvokeModel output tokens" + ); + assert_eq!( + ctx.get_metadata("token.total"), + Some("75"), + "Bedrock InvokeModel total tokens (computed)" + ); +} + +#[tokio::test] +async fn bedrock_invoke_model_headers_absent_is_noop() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = crate::test_utils::make_response(); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.get_metadata("token.input").is_none(), + "missing headers should not set input" + ); + assert!( + ctx.get_metadata("token.output").is_none(), + "missing headers should not set output" + ); +} + +#[tokio::test] +async fn bedrock_invoke_model_partial_headers_is_noop() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = crate::test_utils::make_response(); + resp.headers.insert(HEADER_BEDROCK_INPUT, "25".parse().unwrap()); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.get_metadata("token.input").is_none(), + "only-input header should not set tokens" + ); +} + +#[tokio::test] +async fn bedrock_invoke_model_non_numeric_headers_is_noop() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = crate::test_utils::make_response(); + resp.headers.insert(HEADER_BEDROCK_INPUT, "abc".parse().unwrap()); + resp.headers.insert(HEADER_BEDROCK_OUTPUT, "50".parse().unwrap()); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.get_metadata("token.input").is_none(), + "non-numeric input header should not set tokens" + ); + assert!( + ctx.get_metadata("token.output").is_none(), + "non-numeric input header should also suppress the otherwise-valid output header" + ); +} + +#[tokio::test] +async fn bedrock_invoke_model_skips_non_success_status() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = + make_response_with_status_and_content_type(http::StatusCode::INTERNAL_SERVER_ERROR, "application/json"); + resp.headers.insert(HEADER_BEDROCK_INPUT, "25".parse().unwrap()); + resp.headers.insert(HEADER_BEDROCK_OUTPUT, "50".parse().unwrap()); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.get_metadata("token.input").is_none(), + "non-success status should skip extraction even with valid token headers present" + ); + assert!( + ctx.get_metadata("token.output").is_none(), + "non-success status should skip extraction even with valid token headers present" + ); +} + +#[tokio::test] +async fn bedrock_invoke_model_does_not_set_content_type_mode() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = make_response_with_content_type("application/json"); + resp.headers.insert(HEADER_BEDROCK_INPUT, "25".parse().unwrap()); + resp.headers.insert(HEADER_BEDROCK_OUTPUT, "50".parse().unwrap()); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.get_metadata(META_MODE).is_none(), + "header-only path should never set the body extraction mode" + ); +} + +#[test] +fn bedrock_invoke_model_response_body_access_is_none() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + assert_eq!(filter.response_body_access(), BodyAccess::None); +} + +#[test] +fn other_providers_response_body_access_is_read_only() { + let filter = make_filter(TokenUsageProvider::OpenAi); + assert_eq!(filter.response_body_access(), BodyAccess::ReadOnly); +} + +#[test] +fn on_response_body_noop_for_bedrock_invoke_model() { + let filter = make_filter_kind(ProviderKind::BedrockInvokeModel); + let req = crate::test_utils::make_request(http::Method::POST, "/model/amazon.titan/invoke"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut body: Option = None; + let action = filter.on_response_body(&mut ctx, &mut body, true).unwrap(); + + assert!(matches!(action, FilterAction::Continue)); + assert!(ctx.get_metadata("token.input").is_none()); +} + // ----------------------------------------------------------------------------- // Test Utilities // ----------------------------------------------------------------------------- @@ -600,6 +771,12 @@ fn decode_hex_rejects_invalid_chars() { use std::fmt::Write as _; fn make_filter(provider: TokenUsageProvider) -> TokenCountFilter { + TokenCountFilter { + provider: provider.into(), + } +} + +fn make_filter_kind(provider: ProviderKind) -> TokenCountFilter { TokenCountFilter { provider } } diff --git a/tests/integration/tests/suite/examples/token_count.rs b/tests/integration/tests/suite/examples/token_count.rs index 2f2099c4..f7c810ba 100644 --- a/tests/integration/tests/suite/examples/token_count.rs +++ b/tests/integration/tests/suite/examples/token_count.rs @@ -4,13 +4,21 @@ //! Tests for the token_count filter example configuration. //! //! The filter writes to `filter_metadata` which is not observable from -//! an HTTP response, so we only verify the proxy starts and proxies -//! traffic correctly. Token extraction correctness is covered by unit -//! tests in `praxis-ai-filters`. +//! an HTTP response, so most cases only verify the proxy starts and +//! proxies traffic correctly. Token extraction correctness is covered +//! by unit tests in `praxis-ai-filters`. +//! +//! The Bedrock `InvokeModel` header path is the exception: it is +//! observable end-to-end by chaining `token_count` with +//! `token_usage_headers`, which surfaces the extracted counts as +//! response headers. use std::collections::HashMap; -use praxis_test_utils::{free_port, http_send, parse_status, start_backend_with_shutdown, start_proxy}; +use praxis_core::config::Config; +use praxis_test_utils::{ + Backend, free_port, http_send, parse_header, parse_status, start_backend_with_shutdown, start_proxy, +}; // ----------------------------------------------------------------------------- // Tests @@ -34,3 +42,75 @@ fn token_count_proxies_response() { ); assert_eq!(parse_status(&raw), 200, "proxy should return 200"); } + +#[test] +fn token_count_bedrock_invoke_model_extracts_header_counts() { + let backend_port_guard = Backend::fixed("ok") + .header("x-amzn-bedrock-input-token-count", "25") + .header("x-amzn-bedrock-output-token-count", "50") + .start_with_shutdown(); + let backend_port = backend_port_guard.port(); + let proxy_port = free_port(); + let config = Config::from_yaml(&make_bedrock_invoke_model_yaml(proxy_port, backend_port)).unwrap(); + let proxy = start_proxy(&config); + + let raw = http_send( + proxy.addr(), + "GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n", + ); + + assert_eq!(parse_status(&raw), 200, "proxy should return 200"); + assert_eq!( + parse_header(&raw, "praxis-token-input"), + Some("25".to_owned()), + "input token count should be extracted from Bedrock InvokeModel headers" + ); + assert_eq!( + parse_header(&raw, "praxis-token-output"), + Some("50".to_owned()), + "output token count should be extracted from Bedrock InvokeModel headers" + ); + assert_eq!( + parse_header(&raw, "praxis-token-total"), + Some("75".to_owned()), + "total token count should be computed from input + output" + ); +} + +// ----------------------------------------------------------------------------- +// Test Utilities +// ----------------------------------------------------------------------------- + +/// Builds a config chaining `token_count` (Bedrock `InvokeModel`) with +/// `token_usage_headers`, so the extracted counts become observable as +/// response headers. +/// +/// Response hooks run in *reverse* declared order (last filter's +/// `on_response` fires first), so `token_usage_headers` is declared +/// before `token_count` here to ensure `token_count` populates +/// `filter_metadata` before `token_usage_headers` reads it. +fn make_bedrock_invoke_model_yaml(proxy_port: u16, backend_port: u16) -> String { + format!( + r#" +listeners: + - name: default + address: "127.0.0.1:{proxy_port}" + filter_chains: [main] +filter_chains: + - name: main + filters: + - filter: router + routes: + - path_prefix: "/" + cluster: backend + - filter: token_usage_headers + - filter: token_count + provider: bedrock_invoke_model + - filter: load_balancer + clusters: + - name: backend + endpoints: + - "127.0.0.1:{backend_port}" +"# + ) +}