mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-09-11 22:53:00 +00:00
feat(llm): add input token counting (#359)
## Summary Adds an optional `fabro-llm` API for counting model-visible input tokens without creating a completion, with provider-native counting where available and deterministic local estimates when exact counting is unavailable or intentionally avoided. ## Details - Adds `Client::count_input_tokens` plus `InputTokenCountPreference` modes for provider-preferred, provider-required, and estimate-only behavior. - Implements provider count endpoints for Anthropic, Gemini, and OpenAI while filtering request bodies to count-supported fields. - Adds strict fallback semantics so local estimates do not hide bad credentials, invalid requests, unsupported models, context-length/content-filter failures, or other deterministic provider errors. - Adds a deterministic local estimator with explicit warning codes for local estimates, media heuristics, opaque provider context, and provider options. - Documents privacy implications: provider-native counting sends the provider-serialized model-visible request to the upstream count endpoint, while `EstimateOnly` keeps counting local. ## Verification - `cargo nextest run -p fabro-llm` - 385 passed, 10 skipped - `cargo +nightly-2026-04-14 fmt --check --all` - `cargo +nightly-2026-04-14 clippy -p fabro-llm --all-targets -- -D warnings` - `cargo build --workspace` --- [](https://github.com/EveryInc/compound-engineering-plugin) 🤖 Generated with GPT-5 (context unknown, thinking not disclosed) via [Codex](https://openai.com/codex)
This commit is contained in:
Bryan Helmkamp
GitHub
parent
e67756c10c
commit
8f36772af3
10 changed files with 1251 additions and 14 deletions
|
|
@ -504,6 +504,7 @@ mod tests {
|
|||
use fabro_llm::types::{
|
||||
ContentPart, FinishReason, Message as LlmMessage, Response, Role, TokenCounts, ToolCall,
|
||||
};
|
||||
use tokio::fs;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::*;
|
||||
|
|
@ -901,7 +902,9 @@ mod tests {
|
|||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("src/lib.rs");
|
||||
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
|
||||
std::fs::write(&path, "fn hello() {\n println!(\"old\");\n}\n").unwrap();
|
||||
fs::write(&path, "fn hello() {\n println!(\"old\");\n}\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let env = LocalSandbox::new(dir.path().to_path_buf());
|
||||
let patch = "\
|
||||
*** Begin Patch
|
||||
|
|
@ -919,7 +922,7 @@ mod tests {
|
|||
"Success. Updated the following files:\nM src/lib.rs\n"
|
||||
);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&path).unwrap(),
|
||||
fs::read_to_string(&path).await.unwrap(),
|
||||
"fn hello() {\n println!(\"new\");\n}\n"
|
||||
);
|
||||
}
|
||||
|
|
@ -1039,7 +1042,7 @@ mod tests {
|
|||
async fn pure_addition_update_hunk_uses_raw_local_file_text() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("insert_only.txt");
|
||||
std::fs::write(&path, "alpha\nomega\n").unwrap();
|
||||
fs::write(&path, "alpha\nomega\n").await.unwrap();
|
||||
let env = LocalSandbox::new(dir.path().to_path_buf());
|
||||
let patch = "\
|
||||
*** Begin Patch
|
||||
|
|
@ -1056,7 +1059,7 @@ mod tests {
|
|||
"Success. Updated the following files:\nM insert_only.txt\n"
|
||||
);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&path).unwrap(),
|
||||
fs::read_to_string(&path).await.unwrap(),
|
||||
"alpha\nomega\ninserted\n"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) |
|
||||
|
||||
|
|
|
|||
|
|
@ -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<InputTokenCount, Error> {
|
||||
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<T: ToString>(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<Result<Option<InputTokenCount>, Error>>,
|
||||
count_calls: Arc<AtomicUsize>,
|
||||
reject_named: bool,
|
||||
}
|
||||
|
||||
impl CountingProvider {
|
||||
fn new(result: Result<Option<InputTokenCount>, 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<AtomicUsize> {
|
||||
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<Response, Error> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
async fn stream(&self, _request: &Request) -> Result<StreamEventStream, Error> {
|
||||
unimplemented!()
|
||||
}
|
||||
|
||||
fn supports_tool_choice(&self, mode: &str) -> bool {
|
||||
!(self.reject_named && mode == "named")
|
||||
}
|
||||
|
||||
async fn count_input_tokens(
|
||||
&self,
|
||||
_request: &Request,
|
||||
) -> Result<Option<InputTokenCount>, 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<AtomicUsize>) {
|
||||
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<ApiCredential>,
|
||||
}
|
||||
|
|
@ -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![]);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<StreamEventStream, Error>;
|
||||
|
||||
/// 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<Option<InputTokenCount>, Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Release resources. Called by `Client::close()`.
|
||||
async fn close(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
|
|
|
|||
|
|
@ -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<ApiMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
system: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tools: Option<Vec<ApiToolDef>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tool_choice: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thinking: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl From<ApiRequest> 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<i64>,
|
||||
}
|
||||
|
||||
#[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<Option<InputTokenCount>, 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<Response, Error> {
|
||||
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(),
|
||||
|
|
|
|||
|
|
@ -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<i64>,
|
||||
}
|
||||
|
||||
#[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<Option<InputTokenCount>, 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<Response, Error> {
|
||||
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();
|
||||
|
|
|
|||
|
|
@ -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<ApiUsage>,
|
||||
}
|
||||
|
||||
#[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<ContentPart>, bool) {
|
||||
let mut parts = Vec::new();
|
||||
|
|
@ -1211,6 +1246,51 @@ impl ProviderAdapter for Adapter {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn count_input_tokens(
|
||||
&self,
|
||||
request: &Request,
|
||||
) -> Result<Option<InputTokenCount>, 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<Response, Error> {
|
||||
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();
|
||||
|
|
|
|||
382
lib/crates/fabro-llm/src/token_count.rs
Normal file
382
lib/crates/fabro-llm/src/token_count.rs
Normal file
|
|
@ -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<Warning>,
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn estimate_input_tokens(request: &Request, provider: impl Into<String>) -> 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::<usize>();
|
||||
}
|
||||
|
||||
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<Warning>,
|
||||
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<Message>) -> 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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -206,6 +206,12 @@ impl DockerSandbox {
|
|||
}
|
||||
|
||||
async fn download_file_bytes(&self, remote_path: &str) -> crate::Result<Vec<u8>> {
|
||||
#[expect(
|
||||
clippy::disallowed_types,
|
||||
reason = "tar entries are synchronous in-memory readers; bytes are collected before any await"
|
||||
)]
|
||||
use std::io::Read as _;
|
||||
|
||||
let container_id = self.container_id()?;
|
||||
let container_path = self.resolve_container_path(remote_path);
|
||||
let opts = DownloadFromContainerOptions {
|
||||
|
|
@ -225,12 +231,6 @@ impl DockerSandbox {
|
|||
archive_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::disallowed_types,
|
||||
reason = "tar entries are synchronous in-memory readers; bytes are collected before any await"
|
||||
)]
|
||||
use std::io::Read as _;
|
||||
|
||||
let mut archive = tar::Archive::new(Cursor::new(archive_bytes));
|
||||
let entries = archive.entries().map_err(|e| {
|
||||
crate::Error::context(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue