diff --git a/apps/fabro-web/app/routes/run-detail.tsx b/apps/fabro-web/app/routes/run-detail.tsx index e2fec8160..91e912ea9 100644 --- a/apps/fabro-web/app/routes/run-detail.tsx +++ b/apps/fabro-web/app/routes/run-detail.tsx @@ -680,7 +680,7 @@ export default function RunDetail({ params }: { params: { id: string } }) { diff --git a/lib/crates/fabro-llm/README.md b/lib/crates/fabro-llm/README.md index 2e251b97c..aa992daaa 100644 --- a/lib/crates/fabro-llm/README.md +++ b/lib/crates/fabro-llm/README.md @@ -208,6 +208,39 @@ let anthropic_models = list_models(Some("anthropic")); let best_reasoner = get_latest_model("anthropic", Some("reasoning")); ``` +### Input token counting + +Use `count_input_tokens` when you need the current model-visible context size +without creating a completion: + +```rust +use fabro_llm::{InputTokenCountPreference, Client}; + +let count = client + .count_input_tokens(&request, InputTokenCountPreference::PreferProvider) + .await?; +``` + +`InputTokenCountPreference` controls precision and data exposure: + +- `PreferProvider` sends the provider-serialized request to the upstream + token-count endpoint when supported, then falls back to a local estimate only + for unsupported adapters, network/timeout failures, rate limits, and provider + server errors. +- `RequireProvider` sends the provider-serialized request and returns either a + provider count or an error. It never returns a local estimate. +- `EstimateOnly` validates and resolves the provider locally, does not call the + adapter count endpoint, and returns a deterministic local estimate. + +Provider-native counting sends model-visible request content to the provider's +token-count endpoint. That can include messages, system/developer instructions, +tools, schemas, structured content, and media metadata/content after provider +serialization. Use `EstimateOnly` when that extra upstream exposure is not +acceptable. + +`InputTokenCount` is for input/context sizing. It is not billing usage and does +not include output, reasoning-output, cache-read, or cache-write token buckets. + ## Key types | Type | Description | @@ -222,7 +255,8 @@ let best_reasoner = get_latest_model("anthropic", Some("reasoning")); | `GenerateResult` | Result containing response, tool results, total usage, and step history | | `ToolDefinition` | Tool name, description, and JSON Schema parameters | | `ToolChoice` | Auto, None, Required, or Named tool selection | -| `Usage` | Token counts including input, output, reasoning, and cache tokens | +| `InputTokenCount` | Input/context token count from a provider count API or local estimate | +| `TokenCounts` | Billing-oriented token counts including input, output, reasoning, and cache tokens | | `RetryPolicy` | Configurable retry with exponential backoff, jitter, and max delay | | `Model` | Metadata about a model (context window, capabilities, costs) | diff --git a/lib/crates/fabro-llm/src/client.rs b/lib/crates/fabro-llm/src/client.rs index eed70a1f6..7ad3153d4 100644 --- a/lib/crates/fabro-llm/src/client.rs +++ b/lib/crates/fabro-llm/src/client.rs @@ -6,10 +6,13 @@ use fabro_model::{Catalog, ProviderId}; use tracing::debug; use crate::adapter_registry::{AdapterConfig, factory_for}; -use crate::error::Error; +use crate::error::{Error, ProviderErrorKind}; use crate::middleware::{Middleware, NextFn, NextStreamFn}; use crate::provider::{ProviderAdapter, StreamEventStream}; -use crate::types::{Request, Response, Speed}; +use crate::token_count::{ + InputTokenCount, InputTokenCountMethod, InputTokenCountPreference, estimate_input_tokens, +}; +use crate::types::{Request, Response, Speed, Warning}; /// The core client that routes requests to provider adapters (Section 2.2, 3). #[derive(Clone)] @@ -389,6 +392,59 @@ impl Client { chain(request.clone()).await } + /// Count the model-visible input/context tokens for a request without + /// creating a completion. + /// + /// # Errors + /// + /// Returns request validation/provider resolution errors, and returns + /// provider count errors when the selected preference requires provider + /// semantics or when the error is not fallback-eligible. + pub async fn count_input_tokens( + &self, + request: &Request, + preference: InputTokenCountPreference, + ) -> Result { + self.validate_request_controls(request)?; + let provider = self.resolve_provider(request)?; + provider.validate_request(request)?; + + if preference == InputTokenCountPreference::EstimateOnly { + return Ok(estimate_input_tokens(request, provider.name())); + } + + match provider.count_input_tokens(request).await { + Ok(Some(count)) => Ok(count), + Ok(None) if preference == InputTokenCountPreference::PreferProvider => { + Ok(fallback_estimate( + request, + provider.name(), + "provider_token_count_unsupported", + "provider does not support input token counting; returned local estimate", + )) + } + Ok(None) => Err(Error::Configuration { + message: format!( + "provider '{}' does not support input token counting", + provider.name() + ), + source: None, + }), + Err(error) + if preference == InputTokenCountPreference::PreferProvider + && token_count_fallback_eligible(&error) => + { + Ok(fallback_estimate( + request, + provider.name(), + "provider_token_count_failed", + "provider input token counting failed; returned local estimate", + )) + } + Err(error) => Err(error), + } + } + /// Close all provider adapters. /// /// # Errors @@ -428,6 +484,39 @@ impl Client { } } +fn token_count_fallback_eligible(error: &Error) -> bool { + matches!( + error, + Error::Network { .. } + | Error::RequestTimeout { .. } + | Error::Provider { + kind: ProviderErrorKind::RateLimit | ProviderErrorKind::Server, + .. + } + ) +} + +fn fallback_estimate( + request: &Request, + provider: &str, + code: &'static str, + message: &'static str, +) -> InputTokenCount { + let mut count = estimate_input_tokens(request, provider); + if count.method == InputTokenCountMethod::LocalEstimate + && !count + .warnings + .iter() + .any(|warning| warning.code.as_deref() == Some(code)) + { + count.warnings.push(Warning { + message: message.to_string(), + code: Some(code.to_string()), + }); + } + count +} + fn format_control_values(values: &[T]) -> String { if values.is_empty() { "none".to_string() @@ -450,6 +539,8 @@ fn format_additional_speeds(values: &[Speed]) -> String { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use async_trait::async_trait; use fabro_auth::{ApiKeyHeader, CredentialSource, ResolvedCredentials}; use fabro_model::ProviderId; @@ -457,6 +548,7 @@ mod tests { use futures::stream; use super::*; + use crate::error::ProviderErrorDetail; use crate::types::*; /// A mock provider for testing. @@ -542,6 +634,99 @@ mod tests { } } + struct CountingProvider { + provider_name: String, + count_result: std::sync::Mutex, Error>>, + count_calls: Arc, + reject_named: bool, + } + + impl CountingProvider { + fn new(result: Result, Error>) -> Self { + Self { + provider_name: "counter".to_string(), + count_result: std::sync::Mutex::new(result), + count_calls: Arc::new(AtomicUsize::new(0)), + reject_named: false, + } + } + + fn with_name(mut self, name: &str) -> Self { + self.provider_name = name.to_string(); + self + } + + fn count_calls(&self) -> Arc { + Arc::clone(&self.count_calls) + } + + fn rejecting_named(mut self) -> Self { + self.reject_named = true; + self + } + } + + #[async_trait::async_trait] + impl ProviderAdapter for CountingProvider { + fn name(&self) -> &str { + &self.provider_name + } + + async fn complete(&self, _request: &Request) -> Result { + unimplemented!() + } + + async fn stream(&self, _request: &Request) -> Result { + unimplemented!() + } + + fn supports_tool_choice(&self, mode: &str) -> bool { + !(self.reject_named && mode == "named") + } + + async fn count_input_tokens( + &self, + _request: &Request, + ) -> Result, Error> { + self.count_calls.fetch_add(1, Ordering::SeqCst); + self.count_result.lock().unwrap().clone() + } + } + + fn provider_count(tokens: i64) -> InputTokenCount { + InputTokenCount { + input_tokens: tokens, + method: InputTokenCountMethod::ProviderApi, + provider: "counter".to_string(), + model: "mock-model".to_string(), + warnings: vec![], + } + } + + fn warning_codes(count: &InputTokenCount) -> Vec<&str> { + count + .warnings + .iter() + .filter_map(|warning| warning.code.as_deref()) + .collect() + } + + fn provider_error(kind: ProviderErrorKind) -> Error { + Error::Provider { + kind, + detail: Box::new(ProviderErrorDetail::new("provider failed", "counter")), + } + } + + async fn client_with_counting_provider( + provider: CountingProvider, + ) -> (Client, Arc) { + let calls = provider.count_calls(); + let mut client = Client::new(HashMap::new(), None, vec![]); + client.register_provider(Arc::new(provider)).await.unwrap(); + (client, calls) + } + struct StubSource { credentials: Vec, } @@ -583,6 +768,167 @@ mod tests { assert_eq!(response.provider, "test"); } + #[tokio::test] + async fn count_input_tokens_returns_provider_result() { + let (client, calls) = + client_with_counting_provider(CountingProvider::new(Ok(Some(provider_count(42))))) + .await; + + let count = client + .count_input_tokens(&test_request(), InputTokenCountPreference::PreferProvider) + .await + .unwrap(); + + assert_eq!(count.input_tokens, 42); + assert_eq!(count.method, InputTokenCountMethod::ProviderApi); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn count_input_tokens_prefer_provider_falls_back_for_unsupported_adapter() { + let mut client = Client::new(HashMap::new(), None, vec![]); + client + .register_provider(Arc::new(MockProvider::new("test", ""))) + .await + .unwrap(); + + let count = client + .count_input_tokens(&test_request(), InputTokenCountPreference::PreferProvider) + .await + .unwrap(); + + assert_eq!(count.method, InputTokenCountMethod::LocalEstimate); + assert!(warning_codes(&count).contains(&"provider_token_count_unsupported")); + } + + #[tokio::test] + async fn count_input_tokens_require_provider_errors_for_unsupported_adapter() { + let mut client = Client::new(HashMap::new(), None, vec![]); + client + .register_provider(Arc::new(MockProvider::new("test", ""))) + .await + .unwrap(); + + let error = client + .count_input_tokens(&test_request(), InputTokenCountPreference::RequireProvider) + .await + .unwrap_err(); + + assert!(matches!(error, Error::Configuration { .. })); + } + + #[tokio::test] + async fn count_input_tokens_prefer_provider_falls_back_for_eligible_errors() { + let errors = vec![ + Error::Network { + message: "network down".to_string(), + source: None, + }, + Error::RequestTimeout { + message: "timed out".to_string(), + source: None, + }, + provider_error(ProviderErrorKind::RateLimit), + provider_error(ProviderErrorKind::Server), + ]; + + for error in errors { + let (client, _) = + client_with_counting_provider(CountingProvider::new(Err(error))).await; + let count = client + .count_input_tokens(&test_request(), InputTokenCountPreference::PreferProvider) + .await + .unwrap(); + + assert_eq!(count.method, InputTokenCountMethod::LocalEstimate); + assert!(warning_codes(&count).contains(&"provider_token_count_failed")); + } + } + + #[tokio::test] + async fn count_input_tokens_prefer_provider_returns_non_fallback_errors() { + let errors = vec![ + provider_error(ProviderErrorKind::InvalidRequest), + provider_error(ProviderErrorKind::Authentication), + provider_error(ProviderErrorKind::AccessDenied), + provider_error(ProviderErrorKind::NotFound), + provider_error(ProviderErrorKind::ContextLength), + provider_error(ProviderErrorKind::ContentFilter), + provider_error(ProviderErrorKind::QuotaExceeded), + Error::Configuration { + message: "bad config".to_string(), + source: None, + }, + Error::UnsupportedToolChoice { + message: "bad tool choice".to_string(), + }, + ]; + + for error in errors { + let (client, _) = + client_with_counting_provider(CountingProvider::new(Err(error))).await; + let err = client + .count_input_tokens(&test_request(), InputTokenCountPreference::PreferProvider) + .await + .unwrap_err(); + + assert!(!token_count_fallback_eligible(&err)); + } + } + + #[tokio::test] + async fn count_input_tokens_require_provider_returns_fallback_eligible_errors() { + let (client, _) = client_with_counting_provider(CountingProvider::new(Err( + provider_error(ProviderErrorKind::RateLimit), + ))) + .await; + + let err = client + .count_input_tokens(&test_request(), InputTokenCountPreference::RequireProvider) + .await + .unwrap_err(); + + assert!(matches!(err, Error::Provider { + kind: ProviderErrorKind::RateLimit, + .. + })); + } + + #[tokio::test] + async fn count_input_tokens_estimate_only_does_not_call_adapter() { + let provider = CountingProvider::new(Ok(Some(provider_count(99)))); + let calls = provider.count_calls(); + let (client, _) = client_with_counting_provider(provider).await; + + let count = client + .count_input_tokens(&test_request(), InputTokenCountPreference::EstimateOnly) + .await + .unwrap(); + + assert_eq!(count.method, InputTokenCountMethod::LocalEstimate); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn count_input_tokens_validation_errors_still_return_err() { + let (client, calls) = client_with_counting_provider( + CountingProvider::new(Ok(Some(provider_count(1)))) + .with_name("restricted") + .rejecting_named(), + ) + .await; + let mut request = test_request(); + request.tool_choice = Some(ToolChoice::named("search")); + + let err = client + .count_input_tokens(&request, InputTokenCountPreference::PreferProvider) + .await + .unwrap_err(); + + assert!(matches!(err, Error::UnsupportedToolChoice { .. })); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn complete_routes_to_named_provider() { let mut client = Client::new(HashMap::new(), None, vec![]); diff --git a/lib/crates/fabro-llm/src/lib.rs b/lib/crates/fabro-llm/src/lib.rs index 5e756b989..63066b622 100644 --- a/lib/crates/fabro-llm/src/lib.rs +++ b/lib/crates/fabro-llm/src/lib.rs @@ -7,8 +7,12 @@ pub mod model_test; pub mod provider; pub mod providers; pub mod retry; +pub mod token_count; pub mod tools; pub mod types; pub use error::{Error, ProviderErrorDetail, ProviderErrorKind, Result}; pub use fabro_model::{ModelHandle, ProviderId}; +pub use token_count::{ + InputTokenCount, InputTokenCountMethod, InputTokenCountPreference, estimate_input_tokens, +}; diff --git a/lib/crates/fabro-llm/src/provider.rs b/lib/crates/fabro-llm/src/provider.rs index 076c3ecf3..b49a70cdd 100644 --- a/lib/crates/fabro-llm/src/provider.rs +++ b/lib/crates/fabro-llm/src/provider.rs @@ -4,6 +4,7 @@ pub use fabro_model::{ModelHandle, ProviderId}; use futures::Stream; use crate::error::Error; +use crate::token_count::InputTokenCount; use crate::types::{Request, Response, Speed, StreamEvent, ToolChoice}; // --------------------------------------------------------------------------- @@ -25,6 +26,15 @@ pub trait ProviderAdapter: Send + Sync { /// Send a request and return an async stream of events (Section 4.2). async fn stream(&self, request: &Request) -> Result; + /// Count model-visible input/context tokens without creating a completion, + /// when the provider exposes a count endpoint. + async fn count_input_tokens( + &self, + _request: &Request, + ) -> Result, Error> { + Ok(None) + } + /// Release resources. Called by `Client::close()`. async fn close(&self) -> Result<(), Error> { Ok(()) diff --git a/lib/crates/fabro-llm/src/providers/anthropic.rs b/lib/crates/fabro-llm/src/providers/anthropic.rs index c7979cb38..75bf0ea55 100644 --- a/lib/crates/fabro-llm/src/providers/anthropic.rs +++ b/lib/crates/fabro-llm/src/providers/anthropic.rs @@ -11,6 +11,7 @@ use crate::providers::common::{ self as common, extract_system_prompt, parse_error_body, parse_rate_limit_headers, parse_retry_after, send_and_read_response, }; +use crate::token_count::{InputTokenCount, InputTokenCountMethod}; use crate::types::{ AdapterTimeout, ContentPart, FinishReason, Message, RateLimitInfo, ReasoningEffort, Request, Response, ResponseFormatType, Role, Speed, StreamEvent, ThinkingData, TokenCounts, ToolCall, @@ -77,6 +78,10 @@ impl Adapter { format!("{}/messages", self.http.base_url) } + fn count_tokens_url(&self) -> String { + format!("{}/messages/count_tokens", self.http.base_url) + } + /// Collect a streaming response into a single [`Response`]. /// /// Used by non-Anthropic providers (e.g. Kimi) that require `stream=true`. @@ -137,6 +142,33 @@ struct ApiRequest { stream: bool, } +#[derive(serde::Serialize)] +struct CountTokensRequest { + model: String, + messages: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + system: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tools: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + tool_choice: Option, + #[serde(skip_serializing_if = "Option::is_none")] + thinking: Option, +} + +impl From for CountTokensRequest { + fn from(request: ApiRequest) -> Self { + Self { + model: request.model, + messages: request.messages, + system: request.system, + tools: request.tools, + tool_choice: request.tool_choice, + thinking: request.thinking, + } + } +} + /// Anthropic messages use structured content blocks, not plain strings. #[derive(serde::Serialize)] struct ApiMessage { @@ -194,6 +226,11 @@ struct ApiUsage { cache_creation_input_tokens: Option, } +#[derive(serde::Deserialize)] +struct CountTokensResponse { + input_tokens: i64, +} + fn token_counts_from_api_usage(usage: &ApiUsage) -> TokenCounts { // Anthropic does not expose a separate billed thinking/reasoning token // count. Thinking tokens are billed as part of `output_tokens`. When @@ -1329,6 +1366,67 @@ impl ProviderAdapter for Adapter { &self.provider_name } + async fn count_input_tokens( + &self, + request: &Request, + ) -> Result, Error> { + if self.provider_name != "anthropic" { + return Ok(None); + } + + self.validate_request(request)?; + let (api_request, _req_builder) = build_api_request(self, request, false).await; + let count_request = CountTokensRequest::from(api_request); + + let model_info = common::catalog_model(self.catalog.as_deref(), &request.model); + let supports_prompt_cache = model_info.is_some_and(|m| m.features.prompt_cache); + let auto_cache = + supports_prompt_cache && is_auto_cache_enabled(request.provider_options.as_ref()); + let is_fast = request.speed == Some(Speed::Fast); + let include_1m_context = model_info.is_some_and(|m| m.context_window() >= 1_000_000); + + let url = self.count_tokens_url(); + let mut req = self.http.client.post(&url); + for (key, value) in &self.http.default_headers { + req = req.header(key, value); + } + if let Some(api_key) = &self.http.api_key { + req = req.header("x-api-key", api_key); + } + req = req.header("anthropic-version", "2023-06-01"); + if let Some(beta_str) = build_beta_header( + request.provider_options.as_ref(), + auto_cache, + is_fast, + include_1m_context, + ) { + req = req.header("anthropic-beta", beta_str); + } + + let mut req = req.json(&count_request); + if let Some(t) = self.http.request_timeout { + req = req.timeout(t); + } + + let (body, _headers) = send_and_read_response(req, &self.provider_name, "type").await?; + let response: CountTokensResponse = + serde_json::from_str(&body).map_err(|e| Error::Configuration { + message: format!( + "failed to parse {} token count response: {e}", + self.provider_name + ), + source: None, + })?; + + Ok(Some(InputTokenCount { + input_tokens: response.input_tokens, + method: InputTokenCountMethod::ProviderApi, + provider: self.provider_name.clone(), + model: request.model.clone(), + warnings: vec![], + })) + } + async fn complete(&self, request: &Request) -> Result { self.validate_request(request)?; @@ -1486,6 +1584,7 @@ impl ProviderAdapter for Adapter { #[cfg(test)] mod tests { use fabro_model::catalog::LlmCatalogSettings; + use httpmock::prelude::*; use super::*; use crate::error::ProviderErrorKind; @@ -1946,6 +2045,90 @@ mod tests { ); } + #[tokio::test] + async fn count_input_tokens_posts_count_request_and_parses_response() { + let server = MockServer::start(); + let mock = server.mock(|when, then| { + when.method(POST) + .path("/messages/count_tokens") + .header("x-api-key", "test-key") + .header("anthropic-version", "2023-06-01"); + then.status(200) + .header("content-type", "application/json") + .json_body(serde_json::json!({"input_tokens": 123})); + }); + let adapter = Adapter::new("test-key").with_base_url(server.base_url()); + let request = Request { + messages: vec![Message::system("Be concise"), Message::user("Hello")], + tools: Some(vec![ToolDefinition::function( + "search", + "Search files", + serde_json::json!({"type": "object"}), + )]), + ..make_base_request() + }; + + let count = adapter + .count_input_tokens(&request) + .await + .unwrap() + .expect("anthropic should count tokens"); + + mock.assert(); + assert_eq!(count.input_tokens, 123); + assert_eq!(count.method, InputTokenCountMethod::ProviderApi); + } + + #[tokio::test] + async fn count_request_omits_generation_only_fields_for_reasoning_effort() { + let adapter = Adapter::new("test-key").with_catalog(catalog_with_anthropic_model( + r#" +reasoning_effort = "levels" +"#, + )); + let request = Request { + model: "test-claude".to_string(), + reasoning_effort: Some(ReasoningEffort::High), + temperature: Some(0.2), + top_p: Some(0.9), + metadata: Some(std::collections::HashMap::from([( + "trace".to_string(), + "abc".to_string(), + )])), + ..make_base_request() + }; + + let (api_request, _req_builder) = build_api_request(&adapter, &request, false).await; + assert!(api_request.output_config.is_some()); + let body = serde_json::to_value(CountTokensRequest::from(api_request)).unwrap(); + + assert!(body.get("output_config").is_none()); + assert!(body.get("max_tokens").is_none()); + assert!(body.get("temperature").is_none()); + assert!(body.get("top_p").is_none()); + assert!(body.get("metadata").is_none()); + assert!(body.get("stream").is_none()); + } + + #[tokio::test] + async fn count_request_includes_explicit_thinking_when_translated_request_has_it() { + let adapter = Adapter::new("test-key"); + let request = Request { + provider_options: Some(serde_json::json!({ + "anthropic": { + "thinking": {"type": "enabled", "budget_tokens": 1024} + } + })), + ..make_base_request() + }; + + let (api_request, _req_builder) = build_api_request(&adapter, &request, false).await; + let body = serde_json::to_value(CountTokensRequest::from(api_request)).unwrap(); + + assert_eq!(body["thinking"]["type"], "enabled"); + assert_eq!(body["thinking"]["budget_tokens"], 1024); + } + fn make_base_request() -> Request { Request { model: "claude-sonnet-4-20250514".to_string(), diff --git a/lib/crates/fabro-llm/src/providers/gemini.rs b/lib/crates/fabro-llm/src/providers/gemini.rs index b9c26990f..af4fc4860 100644 --- a/lib/crates/fabro-llm/src/providers/gemini.rs +++ b/lib/crates/fabro-llm/src/providers/gemini.rs @@ -16,6 +16,7 @@ use crate::providers::common::{ self as common, extract_system_prompt, parse_error_body, parse_rate_limit_headers, parse_retry_after, }; +use crate::token_count::{InputTokenCount, InputTokenCountMethod}; use crate::types::{ AdapterTimeout, ContentPart, FinishReason, Message, RateLimitInfo, Request, Response, ResponseFormat, ResponseFormatType, Role, StreamEvent, ThinkingData, TokenCounts, ToolCall, @@ -174,6 +175,12 @@ struct UsageMetadata { tool_use_prompt_token_count: Option, } +#[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct CountTokensResponse { + total_tokens: i64, +} + /// Map Gemini's finish reason, inferring `ToolCalls` from content when needed. fn map_finish_reason(reason: Option<&str>, has_function_calls: bool) -> FinishReason { if has_function_calls { @@ -934,6 +941,42 @@ impl ProviderAdapter for Adapter { Ok(()) } + async fn count_input_tokens( + &self, + request: &Request, + ) -> Result, Error> { + self.validate_request(request)?; + let api_body = build_api_request(request).await; + let api_model = common::api_model_id(self.catalog.as_deref(), &request.model); + let url = format!("{}/models/{}:countTokens", self.http.base_url, api_model); + + let mut req = self.http.client.post(&url); + if let Some(api_key) = &self.http.api_key { + req = req.header("x-goog-api-key", api_key); + } + for (key, value) in &self.http.default_headers { + req = req.header(key, value); + } + let mut req = req.json(&serde_json::json!({ "generateContentRequest": api_body })); + if let Some(t) = self.http.request_timeout { + req = req.timeout(t); + } + let (body, _headers) = send_gemini_response(req).await?; + let response: CountTokensResponse = + serde_json::from_str(&body).map_err(|e| Error::Configuration { + message: format!("failed to parse Gemini token count: {e}"), + source: None, + })?; + + Ok(Some(InputTokenCount { + input_tokens: response.total_tokens, + method: InputTokenCountMethod::ProviderApi, + provider: self.provider_name.clone(), + model: request.model.clone(), + warnings: vec![], + })) + } + async fn complete(&self, request: &Request) -> Result { self.validate_request(request)?; let api_body = build_api_request(request).await; @@ -1033,6 +1076,8 @@ impl ProviderAdapter for Adapter { #[cfg(test)] mod tests { + use httpmock::prelude::*; + use super::*; use crate::types::{AudioData, DocumentData}; @@ -1083,6 +1128,51 @@ mod tests { assert!(body.get("cachedContent").is_none()); } + #[tokio::test] + async fn count_input_tokens_posts_generate_content_request_and_parses_response() { + let server = MockServer::start(); + let mock = server.mock(|when, then| { + when.method(POST) + .path("/models/gemini-2.0-flash:countTokens") + .header("x-goog-api-key", "test-key"); + then.status(200) + .header("content-type", "application/json") + .json_body(serde_json::json!({"totalTokens": 456})); + }); + let adapter = Adapter::new("test-key").with_base_url(server.base_url()); + + let count = adapter + .count_input_tokens(&minimal_request()) + .await + .unwrap() + .expect("gemini should count tokens"); + + mock.assert(); + assert_eq!(count.input_tokens, 456); + assert_eq!(count.method, InputTokenCountMethod::ProviderApi); + } + + #[tokio::test] + async fn count_tokens_body_uses_only_generate_content_request_top_level() { + let mut request = minimal_request(); + request.tools = Some(vec![ToolDefinition::function( + "search", + "Search files", + serde_json::json!({"type": "object"}), + )]); + let api_body = build_api_request(&request).await; + let count_body = serde_json::json!({ "generateContentRequest": api_body }); + + assert!(count_body.get("generateContentRequest").is_some()); + assert!(count_body.get("contents").is_none()); + assert!( + count_body["generateContentRequest"] + .get("contents") + .is_some() + ); + assert!(count_body["generateContentRequest"].get("tools").is_some()); + } + #[tokio::test] async fn provider_options_gemini_safety_settings_merged() { let mut request = minimal_request(); diff --git a/lib/crates/fabro-llm/src/providers/openai.rs b/lib/crates/fabro-llm/src/providers/openai.rs index ea19d151b..b569893b8 100644 --- a/lib/crates/fabro-llm/src/providers/openai.rs +++ b/lib/crates/fabro-llm/src/providers/openai.rs @@ -14,6 +14,7 @@ use crate::providers::common::{ self as common, parse_error_body, parse_rate_limit_headers, parse_retry_after, send_and_read_response, }; +use crate::token_count::{InputTokenCount, InputTokenCountMethod}; use crate::types::{ AdapterTimeout, ContentPart, FinishReason, Message, RateLimitInfo, Request, Response, ResponseFormat, ResponseFormatType, Role, StreamEvent, TokenCounts, ToolCall, ToolChoice, @@ -190,6 +191,12 @@ struct ApiResponse { usage: Option, } +#[derive(serde::Deserialize)] +struct InputTokensResponse { + input_tokens: i64, + object: String, +} + #[derive(serde::Deserialize)] struct ApiUsage { input_tokens: i64, @@ -630,6 +637,34 @@ async fn build_request_body_with_catalog( body } +fn filter_input_tokens_request_body(body: &serde_json::Value) -> serde_json::Value { + const ALLOWED_FIELDS: &[&str] = &[ + "conversation", + "input", + "instructions", + "model", + "parallel_tool_calls", + "previous_response_id", + "reasoning", + "text", + "tool_choice", + "tools", + "truncation", + ]; + + let Some(source) = body.as_object() else { + return serde_json::json!({}); + }; + + let mut filtered = serde_json::Map::new(); + for field in ALLOWED_FIELDS { + if let Some(value) = source.get(*field) { + filtered.insert((*field).to_string(), value.clone()); + } + } + serde_json::Value::Object(filtered) +} + /// Parse output items from the Responses API into content parts. fn parse_output(output: &[serde_json::Value]) -> (Vec, bool) { let mut parts = Vec::new(); @@ -1211,6 +1246,51 @@ impl ProviderAdapter for Adapter { Ok(()) } + async fn count_input_tokens( + &self, + request: &Request, + ) -> Result, Error> { + self.validate_request(request)?; + let request_body = build_request_body_with_catalog( + request, + false, + self.codex_mode, + self.catalog.as_deref(), + ) + .await; + let request_body = filter_input_tokens_request_body(&request_body); + let url = format!("{}/responses/input_tokens", self.http.base_url); + + let mut req = self.build_request(&url).json(&request_body); + if let Some(t) = self.http.request_timeout { + req = req.timeout(t); + } + let (body, _headers) = send_and_read_response(req, "openai", "type").await?; + let response: InputTokensResponse = + serde_json::from_str(&body).map_err(|e| Error::Configuration { + message: format!("failed to parse OpenAI input token response: {e}"), + source: None, + })?; + + if response.object != "response.input_tokens" { + return Err(Error::Configuration { + message: format!( + "failed to parse OpenAI input token response: unexpected object '{}'", + response.object + ), + source: None, + }); + } + + Ok(Some(InputTokenCount { + input_tokens: response.input_tokens, + method: InputTokenCountMethod::ProviderApi, + provider: self.provider_name.clone(), + model: request.model.clone(), + warnings: vec![], + })) + } + async fn complete(&self, request: &Request) -> Result { self.validate_request(request)?; @@ -1337,7 +1417,7 @@ mod tests { use super::*; use crate::error::ProviderErrorKind; use crate::providers::common::LineReader; - use crate::types::{AudioData, DocumentData, ToolResult}; + use crate::types::{AudioData, DocumentData, ReasoningEffort, ToolResult}; fn minimal_request() -> Request { Request { @@ -1433,6 +1513,111 @@ mod tests { assert!(body.get("stream").is_none()); } + #[tokio::test] + async fn filter_input_tokens_request_body_keeps_only_count_fields() { + let mut metadata = HashMap::new(); + metadata.insert("trace".to_string(), "abc".to_string()); + + let mut request = minimal_request(); + request.tools = Some(vec![ToolDefinition::function( + "search", + "Search files", + serde_json::json!({"type": "object"}), + )]); + request.reasoning_effort = Some(ReasoningEffort::Low); + request.response_format = Some(ResponseFormat { + kind: ResponseFormatType::JsonSchema, + json_schema: Some(serde_json::json!({"type": "object"})), + strict: true, + }); + request.temperature = Some(0.2); + request.top_p = Some(0.9); + request.max_tokens = Some(32); + request.stop_sequences = Some(vec!["END".to_string()]); + request.metadata = Some(metadata); + + let body = build_request_body(&request, true, false).await; + let filtered = filter_input_tokens_request_body(&body); + + assert_eq!( + filtered, + serde_json::json!({ + "input": [{"type": "message", "content": [{"text": "Hello", "type": "input_text"}], "role": "user"}], + "model": "gpt-4o", + "reasoning": {"effort": "low"}, + "text": {"format": {"name": "response", "schema": {"type": "object"}, "strict": true, "type": "json_schema"}}, + "tools": [{"description": "Search files", "name": "search", "parameters": {"type": "object"}, "type": "function"}] + }) + ); + assert!(filtered.get("store").is_none()); + assert!(filtered.get("include").is_none()); + assert!(filtered.get("stream").is_none()); + assert!(filtered.get("max_output_tokens").is_none()); + assert!(filtered.get("metadata").is_none()); + assert!(filtered.get("temperature").is_none()); + assert!(filtered.get("top_p").is_none()); + assert!(filtered.get("stop").is_none()); + } + + #[tokio::test] + async fn filter_input_tokens_request_body_preserves_codex_serialization() { + let body = build_request_body(&minimal_request(), false, true).await; + let filtered = filter_input_tokens_request_body(&body); + + assert_eq!(filtered["instructions"], ""); + assert!(filtered.get("input").is_some()); + assert!(filtered.get("model").is_some()); + assert!(filtered.get("max_output_tokens").is_none()); + assert!(filtered.get("include").is_none()); + } + + #[tokio::test] + async fn count_input_tokens_posts_count_request_and_parses_response() { + let server = MockServer::start(); + let mock = server.mock(|when, then| { + when.method(POST).path("/responses/input_tokens"); + then.status(200) + .header("content-type", "application/json") + .json_body(serde_json::json!({ + "object": "response.input_tokens", + "input_tokens": 789 + })); + }); + let adapter = Adapter::new("sk-test").with_base_url(server.base_url()); + + let count = adapter + .count_input_tokens(&minimal_request()) + .await + .unwrap() + .expect("openai should count tokens"); + + mock.assert(); + assert_eq!(count.input_tokens, 789); + assert_eq!(count.method, InputTokenCountMethod::ProviderApi); + } + + #[tokio::test] + async fn count_input_tokens_rejects_wrong_response_object() { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(POST).path("/responses/input_tokens"); + then.status(200) + .header("content-type", "application/json") + .json_body(serde_json::json!({ + "object": "other", + "input_tokens": 789 + })); + }); + let adapter = Adapter::new("sk-test").with_base_url(server.base_url()); + + let err = adapter + .count_input_tokens(&minimal_request()) + .await + .unwrap_err(); + + assert!(matches!(err, Error::Configuration { .. })); + } + #[tokio::test] async fn build_request_body_includes_encrypted_reasoning_for_stateless_requests() { let request = minimal_request(); diff --git a/lib/crates/fabro-llm/src/token_count.rs b/lib/crates/fabro-llm/src/token_count.rs new file mode 100644 index 000000000..46fb12f2c --- /dev/null +++ b/lib/crates/fabro-llm/src/token_count.rs @@ -0,0 +1,382 @@ +use std::collections::HashSet; + +use serde::{Deserialize, Serialize}; + +use crate::types::{ + AudioData, ContentPart, DocumentData, ImageData, Message, Request, Role, ToolDefinition, + ToolResult, Warning, +}; + +const LOCAL_ESTIMATE_WARNING: &str = "local_token_estimate"; +const MEDIA_ESTIMATE_WARNING: &str = "media_token_estimate"; +const OPAQUE_CONTEXT_ESTIMATE_WARNING: &str = "opaque_context_estimate"; +const PROVIDER_OPTIONS_ESTIMATE_WARNING: &str = "provider_options_estimate"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputTokenCountPreference { + PreferProvider, + RequireProvider, + EstimateOnly, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum InputTokenCountMethod { + ProviderApi, + LocalEstimate, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct InputTokenCount { + pub input_tokens: i64, + pub method: InputTokenCountMethod, + pub provider: String, + pub model: String, + #[serde(default)] + pub warnings: Vec, +} + +#[must_use] +pub fn estimate_input_tokens(request: &Request, provider: impl Into) -> InputTokenCount { + let mut estimator = Estimator::default(); + let mut tokens = 0usize; + + for message in &request.messages { + tokens += 4; + tokens += estimate_text_tokens(message.role_name()); + if let Some(name) = &message.name { + tokens += estimate_text_tokens(name); + } + if let Some(tool_call_id) = &message.tool_call_id { + tokens += estimate_text_tokens(tool_call_id); + } + for part in &message.content { + tokens += 1 + estimator.estimate_content_part(part); + } + } + + if let Some(tools) = &request.tools { + tokens += tools.iter().map(estimate_tool).sum::(); + } + + if let Some(tool_choice) = &request.tool_choice { + if let Ok(value) = serde_json::to_value(tool_choice) { + tokens += estimate_json_tokens(&value); + } + } + + if let Some(response_format) = &request.response_format { + if let Ok(value) = serde_json::to_value(response_format) { + tokens += estimate_json_tokens(&value); + } + } + + if let Some(reasoning_effort) = request.reasoning_effort { + tokens += estimate_text_tokens(reasoning_effort.to_string().as_str()); + } + + if let Some(provider_options) = &request.provider_options { + tokens += estimate_json_tokens(provider_options); + estimator.warn( + PROVIDER_OPTIONS_ESTIMATE_WARNING, + "provider options estimated from JSON", + ); + } + + estimator.warn( + LOCAL_ESTIMATE_WARNING, + "input token count is a local estimate", + ); + + InputTokenCount { + input_tokens: i64::try_from(tokens).unwrap_or(i64::MAX), + method: InputTokenCountMethod::LocalEstimate, + provider: provider.into(), + model: request.model.clone(), + warnings: estimator.warnings, + } +} + +#[must_use] +pub fn estimate_text_tokens(text: &str) -> usize { + text.chars().count().div_ceil(4) +} + +#[must_use] +pub fn estimate_json_tokens(value: &serde_json::Value) -> usize { + serde_json::to_string(value).map_or(0, |json| json.len().div_ceil(4)) +} + +#[derive(Default)] +struct Estimator { + warnings: Vec, + seen_codes: HashSet<&'static str>, +} + +impl Estimator { + fn estimate_content_part(&mut self, part: &ContentPart) -> usize { + match part { + ContentPart::Text(text) => estimate_text_tokens(text), + ContentPart::Image(image) => self.estimate_image(image), + ContentPart::Audio(audio) => self.estimate_audio(audio), + ContentPart::Document(document) => self.estimate_document(document), + ContentPart::ToolCall(tool_call) => estimate_json_tokens(&serde_json::json!(tool_call)), + ContentPart::ToolResult(result) => self.estimate_tool_result(result), + ContentPart::Thinking(thinking) => { + estimate_text_tokens(&thinking.text) + + thinking + .signature + .as_deref() + .map_or(0, estimate_text_tokens) + + usize::from(thinking.redacted) + } + ContentPart::Other { kind, data } => { + self.warn( + OPAQUE_CONTEXT_ESTIMATE_WARNING, + "opaque provider context estimated from JSON", + ); + estimate_text_tokens(kind) + estimate_json_tokens(data) + } + } + } + + fn estimate_tool_result(&mut self, result: &ToolResult) -> usize { + let mut tokens = + estimate_text_tokens(&result.tool_call_id) + estimate_json_tokens(&result.content); + if let Some(image_data) = &result.image_data { + tokens += estimate_embedded_bytes(image_data.len()); + self.warn( + MEDIA_ESTIMATE_WARNING, + "media content estimated heuristically", + ); + } + if let Some(media_type) = &result.image_media_type { + tokens += estimate_text_tokens(media_type); + } + tokens + usize::from(result.is_error) + } + + fn estimate_image(&mut self, image: &ImageData) -> usize { + let mut tokens = + self.estimate_media_common(image.url.as_deref(), image.media_type.as_deref()); + if let Some(detail) = &image.detail { + tokens += estimate_text_tokens(detail); + } + tokens + + image + .data + .as_ref() + .map_or(2000, |data| estimate_embedded_bytes(data.len()).max(2000)) + } + + fn estimate_audio(&mut self, audio: &AudioData) -> usize { + let tokens = self.estimate_media_common(audio.url.as_deref(), audio.media_type.as_deref()); + tokens + + audio + .data + .as_ref() + .map_or(2000, |data| estimate_embedded_bytes(data.len())) + } + + fn estimate_document(&mut self, document: &DocumentData) -> usize { + let mut tokens = + self.estimate_media_common(document.url.as_deref(), document.media_type.as_deref()); + if let Some(file_name) = &document.file_name { + tokens += estimate_text_tokens(file_name); + } + tokens + + document + .data + .as_ref() + .map_or(2000, |data| estimate_embedded_bytes(data.len())) + } + + fn estimate_media_common(&mut self, url: Option<&str>, media_type: Option<&str>) -> usize { + self.warn( + MEDIA_ESTIMATE_WARNING, + "media content estimated heuristically", + ); + url.map_or(0, estimate_text_tokens) + media_type.map_or(0, estimate_text_tokens) + } + + fn warn(&mut self, code: &'static str, message: &'static str) { + if self.seen_codes.insert(code) { + self.warnings.push(Warning { + message: message.to_string(), + code: Some(code.to_string()), + }); + } + } +} + +fn estimate_tool(tool: &ToolDefinition) -> usize { + 8 + estimate_text_tokens(&tool.name) + + estimate_text_tokens(&tool.description) + + estimate_json_tokens(&tool.parameters) +} + +fn estimate_embedded_bytes(byte_len: usize) -> usize { + byte_len.div_ceil(4) +} + +trait RoleName { + fn role_name(&self) -> &'static str; +} + +impl RoleName for Message { + fn role_name(&self) -> &'static str { + match self.role { + Role::System => "system", + Role::User => "user", + Role::Assistant => "assistant", + Role::Tool => "tool", + Role::Developer => "developer", + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::types::{ + DocumentData, ImageData, Request, ResponseFormat, ResponseFormatType, ToolDefinition, + }; + + fn request(messages: Vec) -> Request { + Request { + model: "model-a".to_string(), + messages, + provider: Some("test".to_string()), + tools: None, + tool_choice: None, + response_format: None, + temperature: None, + top_p: None, + max_tokens: None, + stop_sequences: None, + reasoning_effort: None, + speed: None, + metadata: None, + provider_options: None, + } + } + + fn warning_codes(count: &InputTokenCount) -> Vec<&str> { + count + .warnings + .iter() + .filter_map(|warning| warning.code.as_deref()) + .collect() + } + + #[test] + fn text_only_request_returns_positive_local_estimate() { + let count = estimate_input_tokens(&request(vec![Message::user("hello world")]), "test"); + + assert!(count.input_tokens > 0); + assert_eq!(count.method, InputTokenCountMethod::LocalEstimate); + assert!(warning_codes(&count).contains(&LOCAL_ESTIMATE_WARNING)); + } + + #[test] + fn adding_tool_increases_estimate() { + let mut with_tool = request(vec![Message::user("hello")]); + let without_tool = estimate_input_tokens(&with_tool, "test"); + + with_tool.tools = Some(vec![ToolDefinition::function( + "search", + "Search files", + json!({"type": "object", "properties": {"query": {"type": "string"}}}), + )]); + let with_tool = estimate_input_tokens(&with_tool, "test"); + + assert!(with_tool.input_tokens > without_tool.input_tokens); + } + + #[test] + fn adding_response_format_increases_estimate() { + let mut with_schema = request(vec![Message::user("hello")]); + let without_schema = estimate_input_tokens(&with_schema, "test"); + + with_schema.response_format = Some(ResponseFormat { + kind: ResponseFormatType::JsonSchema, + json_schema: Some( + json!({"type": "object", "properties": {"answer": {"type": "string"}}}), + ), + strict: true, + }); + let with_schema = estimate_input_tokens(&with_schema, "test"); + + assert!(with_schema.input_tokens > without_schema.input_tokens); + } + + #[test] + fn media_content_gets_media_warning_and_sized_estimate() { + let count = estimate_input_tokens( + &request(vec![Message { + role: Role::User, + content: vec![ + ContentPart::Image(ImageData { + url: Some("https://example.test/image.png".to_string()), + data: None, + media_type: Some("image/png".to_string()), + detail: Some("high".to_string()), + }), + ContentPart::Document(DocumentData { + url: None, + data: Some(vec![0; 4096]), + media_type: Some("application/pdf".to_string()), + file_name: Some("doc.pdf".to_string()), + }), + ], + name: None, + tool_call_id: None, + }]), + "test", + ); + + assert!(count.input_tokens >= 3000); + assert!(warning_codes(&count).contains(&MEDIA_ESTIMATE_WARNING)); + } + + #[test] + fn provider_options_produce_provider_options_warning() { + let mut req = request(vec![Message::user("hello")]); + req.provider_options = Some(json!({"gemini": {"cached_content": "cachedContents/1"}})); + + let count = estimate_input_tokens(&req, "test"); + + assert!(warning_codes(&count).contains(&PROVIDER_OPTIONS_ESTIMATE_WARNING)); + } + + #[test] + fn opaque_content_produces_opaque_warning() { + let count = estimate_input_tokens( + &request(vec![Message { + role: Role::Assistant, + content: vec![ContentPart::Other { + kind: "openai_reasoning".to_string(), + data: json!({"id": "rs_123", "summary": []}), + }], + name: None, + tool_call_id: None, + }]), + "test", + ); + + assert!(warning_codes(&count).contains(&OPAQUE_CONTEXT_ESTIMATE_WARNING)); + } + + #[test] + fn estimate_is_deterministic() { + let req = request(vec![Message::user("repeatable")]); + + assert_eq!( + estimate_input_tokens(&req, "test"), + estimate_input_tokens(&req, "test") + ); + } +}