diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 67f778af..925dbdc6 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -111,23 +111,6 @@ impl TranslatingLlmClient { }) } - /// A configured Anthropic backend to forward a direct `count_tokens` call to, - /// paired with its upstream model id (the id the inbound route name is - /// restamped to). `None` when this client has no Anthropic backend; - /// `count_tokens` is Anthropic-only. - fn anthropic_backend(&self) -> Option<(&str, &Backend)> { - self.model_to_config.values().find_map(|config| { - if config.default_backend.is_anthropic() { - return Some((config.model_name.as_str(), &config.default_backend)); - } - config - .other_backends - .as_ref() - .and_then(|backends| backends.iter().find(|backend| backend.is_anthropic())) - .map(|backend| (config.model_name.as_str(), backend)) - }) - } - /// The backend serving `model` over `format` — the default backend when its /// format matches, otherwise a matching entry in `other_backends`; `None` when /// the model is unknown or has no backend for `format`. @@ -144,6 +127,50 @@ impl TranslatingLlmClient { }) } + /// Whether `model` has an Anthropic backend that supports token counting. + pub fn supports_count_tokens(&self, model: &str) -> bool { + self.backend_for(model, WireFormat::AnthropicMessages) + .is_some() + } + + /// Counts input tokens with `model`'s Anthropic backend. + /// + /// Returns an error when the model has no Anthropic backend or the upstream + /// request fails or returns invalid JSON. + pub async fn count_tokens(&self, model: &str, request: Request) -> Result { + let backend = self + .backend_for(model, WireFormat::AnthropicMessages) + .ok_or_else(|| LlmClientError::Configuration { + message: format!("model {model} has no Anthropic backend for count_tokens"), + })?; + let Request { + llm_request, + metadata, + .. + } = request; + let http_response = self + .send_encoded( + backend, + WireFormat::AnthropicMessages, + llm_request, + metadata.as_ref(), + model, + UpstreamEndpoint::CountTokens, + ) + .await?; + let body = match http_response { + EncodedResponse::Buffered { body, .. } => body, + EncodedResponse::Streaming(_) => { + return Err(LlmClientError::InvalidRequest { + message: "count_tokens does not support streaming requests".to_string(), + }); + } + }; + serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { + source: Box::new(error), + }) + } + /// Encode `llm_request` (its model restamped to `model`) for `wire_format`, /// POST it to `url` with the request's forwarded headers plus the backend's /// static headers and auth, and return the successful upstream response. A @@ -153,8 +180,8 @@ impl TranslatingLlmClient { /// overflow via the backend's provider rules. Shared by /// [`call_rewrite_model`](Self::call_rewrite_model) (which POSTs to the /// backend's completion URL and decodes a response) and - /// [`count_tokens`](RoutedLlmClient::count_tokens) (which POSTs to the - /// `count_tokens` URL and returns the raw JSON). + /// [`count_tokens`](Self::count_tokens) (which POSTs to the `count_tokens` + /// URL and returns the raw JSON). async fn send_encoded( &self, backend: &Backend, @@ -488,49 +515,6 @@ impl RoutedLlmClient for TranslatingLlmClient { let model_name = Some(decision.selected_model()); self.call_rewrite_model(ctx, request, model_name).await } - - fn supports_count_tokens(&self) -> bool { - self.anthropic_backend().is_some() - } - - async fn count_tokens(&self, request: Request) -> Result { - // Direct passthrough: forward straight to this client's Anthropic - // backend's count_tokens endpoint. Not routed — no decision. - let (model, backend) = - self.anthropic_backend() - .ok_or_else(|| LlmClientError::Configuration { - message: "count_tokens is anthropic-only; this client has no anthropic backend" - .to_string(), - })?; - // Same encode-and-forward path as the completion call, only to the - // `count_tokens` URL; the response is the raw `{"input_tokens": N}` JSON. - let Request { - llm_request, - metadata, - .. - } = request; - let http_response = self - .send_encoded( - backend, - WireFormat::AnthropicMessages, - llm_request, - metadata.as_ref(), - model, - UpstreamEndpoint::CountTokens, - ) - .await?; - let body = match http_response { - EncodedResponse::Buffered { body, .. } => body, - EncodedResponse::Streaming(_) => { - return Err(LlmClientError::InvalidRequest { - message: "count_tokens does not support streaming requests".to_string(), - }); - } - }; - serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse { - source: Box::new(error), - }) - } } #[derive(Clone, Copy)] diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index d7d2381a..8d1314a5 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -30,7 +30,7 @@ use crate::core::algorithm::{self, Algorithm, Driver, LlmTarget, LlmTargetSet, S use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::processor::{Event, Processor}; use crate::{LibsyError, Result}; -use switchyard_protocol::{Context, Decision, Request, Response, RoutedLlmClient}; +use switchyard_protocol::{Context, Decision, Request, Response}; struct SessionState { state: Arc>, @@ -420,10 +420,6 @@ where &self.name } - fn count_tokens_client(&self) -> Option> { - self.targets.count_tokens_client() - } - async fn create_run_task( self: Arc, ctx: Context, @@ -441,8 +437,8 @@ mod tests { use crate::{SystemPromptProcessor, TargetPrompts}; use switchyard_protocol::{ - LlmClientError, LlmRequest, LlmResponse, Message, Metadata, Role, completion_text, - text_request, text_response, + LlmClientError, LlmRequest, LlmResponse, Message, Metadata, Role, RoutedLlmClient, + completion_text, text_request, text_response, }; #[derive(Debug, thiserror::Error)] diff --git a/crates/libsy/src/algorithms/llm_class.rs b/crates/libsy/src/algorithms/llm_class.rs index 4c627b09..c2cce86e 100644 --- a/crates/libsy/src/algorithms/llm_class.rs +++ b/crates/libsy/src/algorithms/llm_class.rs @@ -26,7 +26,7 @@ use crate::core::classifier::{Classification, Classifier, Score}; use crate::core::state::{State, StateValue}; use crate::{LibsyError, Result}; use switchyard_protocol::{ - AggLlmResponse, Context, LlmClientError, LlmResponse, Request, Response, RoutedLlmClient, + AggLlmResponse, Context, LlmClientError, LlmResponse, Request, Response, }; const PROMPT_TEMPLATE: &str = include_str!("../prompts/capability-classifier/prompt.md"); @@ -936,10 +936,6 @@ impl Algorithm for LlmTaskClassifier { "llm_task_classifier" } - fn count_tokens_client(&self) -> Option> { - self.route.count_tokens_client() - } - async fn create_run_task( self: Arc, ctx: Context, diff --git a/crates/libsy/src/algorithms/passthrough.rs b/crates/libsy/src/algorithms/passthrough.rs index 2a922f28..9643fc4b 100644 --- a/crates/libsy/src/algorithms/passthrough.rs +++ b/crates/libsy/src/algorithms/passthrough.rs @@ -9,7 +9,7 @@ use switchyard_protocol::{Request, Response}; use crate::Result; use crate::core::algorithm::{Algorithm, Driver, LlmTarget}; -use switchyard_protocol::{Context, Decision, RoutedLlmClient}; +use switchyard_protocol::{Context, Decision}; /// Routing algorithm that always calls one configured target. pub struct Passthrough { @@ -47,14 +47,6 @@ impl Algorithm for Passthrough { "passthrough" } - fn count_tokens_client(&self) -> Option> { - self.target - .llm_client - .as_ref() - .filter(|client| client.supports_count_tokens()) - .cloned() - } - async fn create_run_task( self: Arc, ctx: Context, diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index 2eb93f02..3f41500c 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -19,7 +19,7 @@ use crate::algorithms::fall_through::{FallThrough, FallThroughDecision}; use crate::core::algorithm::{Algorithm, Driver, LlmTargetSet}; use crate::core::classifier::{Classification, Classifier, Score}; use crate::{LibsyError, Result}; -use switchyard_protocol::{Context, Request, Response, RoutedLlmClient}; +use switchyard_protocol::{Context, Request, Response}; /// Compatibility name for the decision produced by [`Random`]. pub type RandomDecision = FallThroughDecision; @@ -162,10 +162,6 @@ impl Algorithm for Random { "random" } - fn count_tokens_client(&self) -> Option> { - self.inner.count_tokens_client() - } - async fn create_run_task( self: Arc, ctx: Context, diff --git a/crates/libsy/src/algorithms/stage.rs b/crates/libsy/src/algorithms/stage.rs index a2612d69..c82577c5 100644 --- a/crates/libsy/src/algorithms/stage.rs +++ b/crates/libsy/src/algorithms/stage.rs @@ -29,7 +29,7 @@ use crate::core::algorithm::{Algorithm, Driver, LlmTarget, LlmTargetSet}; use crate::core::classifier::{Classification, Classifier}; use crate::core::state::State; use crate::{LibsyError, Result}; -use switchyard_protocol::{Context, Request, Response, RoutedLlmClient}; +use switchyard_protocol::{Context, Request, Response}; /// Telemetry name for a router this module assembles. const STAGE_ROUTER: &str = "stage_router"; @@ -142,10 +142,6 @@ impl Algorithm for StageRouter { STAGE_ROUTER } - fn count_tokens_client(&self) -> Option> { - self.route.count_tokens_client() - } - async fn create_run_task( self: Arc, ctx: Context, diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 676328bc..d63294ee 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -405,19 +405,6 @@ impl LlmTargetSet { .cloned() .ok_or(LibsyError::AllTargetsExcluded) } - - /// The first target's client that can serve `count_tokens` (an Anthropic - /// upstream), or `None` when no target has one. Used by an algorithm's - /// [`count_tokens_client`](crate::Algorithm::count_tokens_client). - pub fn count_tokens_client(&self) -> Option> { - self.targets.iter().find_map(|target| { - target - .llm_client - .as_ref() - .filter(|client| client.supports_count_tokens()) - .cloned() - }) - } } /// Bounds process-local overflow history. Dropping a live session's entry costs one @@ -590,39 +577,6 @@ pub trait Algorithm: Send + Sync + 'static { Ok(()) } - /// The client [`count_tokens`](Self::count_tokens) forwards to: the first of - /// this algorithm's targets whose client can count tokens (an Anthropic - /// upstream). The default is `None` — an algorithm with no Anthropic target - /// does not support token counting. - /// - /// CAVEAT: this picks the *first* Anthropic target, not a routed one — - /// count_tokens is a direct passthrough, so it does not run the routing - /// cascade. For a route with several Anthropic tiers "first" is arbitrary; - /// choosing which tier count_tokens should reflect is deferred. - fn count_tokens_client(&self) -> Option> { - None - } - - /// Count the tokens `request` would use — a **direct passthrough** to this - /// algorithm's Anthropic target (via - /// [`count_tokens_client`](Self::count_tokens_client)), **not** a routed - /// call. Token counting is a pre-flight estimate with no routing decision, - /// so it deliberately bypasses the classifier cascade (which runs only for - /// completions via [`run`](Self::run)). Returns the upstream's JSON - /// verbatim. Errors when the algorithm has no Anthropic target. - async fn count_tokens(&self, request: Request) -> Result { - let client = self - .count_tokens_client() - .ok_or_else(|| LibsyError::AlgorithmError { - message: "no target supports count_tokens (needs an Anthropic upstream)" - .to_string(), - })?; - client - .count_tokens(request) - .await - .map_err(|source| LibsyError::client_call("count_tokens", source)) - } - /// Process a request to completion, returning a stream of [`Step`]s. /// /// The consumer must fulfill every [`Step::CallLlm`] before the algorithm can diff --git a/crates/protocol/src/client.rs b/crates/protocol/src/client.rs index 233ddb38..05124509 100644 --- a/crates/protocol/src/client.rs +++ b/crates/protocol/src/client.rs @@ -180,24 +180,4 @@ pub trait RoutedLlmClient: Send + Sync { request: Request, decision: Arc, ) -> Result; - - /// Whether this client can serve [`count_tokens`](Self::count_tokens) — i.e. - /// it has an Anthropic upstream. The default is `false`. - fn supports_count_tokens(&self) -> bool { - false - } - - /// Count the tokens `request` would use — a **direct passthrough**, not a - /// routed call. Forwards `request` to this client's Anthropic - /// `/v1/messages/count_tokens` endpoint (model restamped to the upstream - /// target id) and returns the JSON verbatim. Token counting is a pre-flight - /// estimate with no routing decision, so unlike [`call`](Self::call) it - /// takes no [`Decision`]. The default errors; only an Anthropic-backed - /// client overrides it. - async fn count_tokens(&self, request: Request) -> Result { - let _ = request; - Err(LlmClientError::Configuration { - message: "count_tokens is not supported by this client".to_string(), - }) - } } diff --git a/crates/switchyard-server/README.md b/crates/switchyard-server/README.md index edbf5a25..afab8f98 100644 --- a/crates/switchyard-server/README.md +++ b/crates/switchyard-server/README.md @@ -119,7 +119,7 @@ documented in [Stage-Router Routing](../../docs/routing_algorithms/stage_router_ | `POST` | `/v1/chat/completions` | OpenAI Chat Completions | | `POST` | `/v1/messages` | Anthropic Messages | | `POST` | `/v1/responses` | OpenAI Responses | -| `POST` | `/v1/messages/count_tokens` | Anthropic token count | +| `POST` | `/v1/messages/count_tokens` | Token count from a route's Anthropic target | | `GET` | `/v1/models` | Routes served by this deployment | | `GET` | `/v1/stats` | Per-model request, token, and cost totals | | `POST` | `/v1/stats/reset` | Clear accumulated stats | @@ -130,6 +130,9 @@ Requests name a route by its `id`, so `POST /v1/chat/completions` with `"model": routes through the `[routes.general]` entry above. Any of the three request formats can address any route, and the server translates between them. +Token counting selects an Anthropic-format completion target, preferring target names or model IDs +containing `opus`, `sonnet`, then `haiku`. Other ties preserve the route's target order. + ## Metrics `GET /metrics` exposes Prometheus text from the server's process-wide OpenTelemetry provider. diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index 95cc57b6..18ad3f7f 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -21,7 +21,7 @@ use switchyard_llm_client::{ }; use switchyard_protocol::RoutedLlmClient; -use crate::{ModelCapabilities, ServerError, ServerResult, ServerState}; +use crate::{CountTokensTarget, ModelCapabilities, ServerError, ServerResult, ServerState}; const SUPPORTED_SCHEMA_VERSION: u32 = 1; const MAX_CONFIGURED_RETRIES: u32 = 10; @@ -98,12 +98,18 @@ impl ServerConfig { ))); } let algorithm = build_algorithm(route_name, config, &targets)?; - routes.push((config.id().to_string(), algorithm, capabilities)); + let count_tokens_target = self.build_count_tokens_target(config, &clients); + routes.push(( + config.id().to_string(), + algorithm, + capabilities, + count_tokens_target, + )); } ServerState::new_with_capabilities(routes) } - fn build_clients(&self) -> ServerResult>> { + fn build_clients(&self) -> ServerResult>> { let mut models_by_client = self .llm_clients .keys() @@ -132,7 +138,7 @@ impl ServerConfig { let mut clients = BTreeMap::new(); for (name, model_configs) in models_by_client { - let client: Arc = Arc::new( + let client = Arc::new( TranslatingLlmClient::new(&model_configs) .map_err(|error| ServerError::new(error.to_string()))?, ); @@ -143,7 +149,7 @@ impl ServerConfig { fn build_targets( &self, - clients: &BTreeMap>, + clients: &BTreeMap>, ) -> ServerResult> { self.targets .iter() @@ -151,16 +157,53 @@ impl ServerConfig { let client = clients.get(&config.llm_client).ok_or_else(|| { ServerError::new(format!("target {name} has no constructed llm client")) })?; + let client: Arc = client.clone(); Ok(( name.clone(), LlmTarget { semantic_name: config.id.clone(), - llm_client: Some(Arc::clone(client)), + llm_client: Some(client), }, )) }) .collect() } + + fn build_count_tokens_target( + &self, + route_config: &RouteConfig, + clients: &BTreeMap>, + ) -> Option { + route_config + .routing_target_names() + .into_iter() + .enumerate() + .filter_map(|(index, name)| { + let target = self.targets.get(name)?; + let client = clients.get(&target.llm_client)?; + client.supports_count_tokens(&target.id).then_some(( + count_tokens_priority(name, &target.id), + index, + target, + client, + )) + }) + .min_by_key(|(priority, index, _, _)| (*priority, *index)) + .map(|(_, _, target, client)| CountTokensTarget { + model: target.id.clone(), + client: client.clone(), + }) + } +} + +// Prefer known Claude families, then preserve the route's target order. +fn count_tokens_priority(target_name: &str, model_id: &str) -> usize { + let target_name = target_name.to_ascii_lowercase(); + let model_id = model_id.to_ascii_lowercase(); + ["opus", "sonnet", "haiku"] + .iter() + .position(|hint| target_name.contains(hint) || model_id.contains(hint)) + .unwrap_or(3) } #[derive(Debug, Deserialize)] @@ -399,6 +442,44 @@ impl RouteConfig { } } + // Completion targets in algorithm order; judge-only targets are excluded. + fn routing_target_names(&self) -> Vec<&str> { + match self { + Self::Noop { .. } => Vec::new(), + Self::Random { targets, .. } => targets.iter().map(String::as_str).collect(), + Self::Passthrough { target, .. } => vec![target], + Self::LlmClassifier { + mode, + strong_target, + weak_target, + escalation, + targets, + .. + } => match mode.unwrap_or(if escalation.is_some() { + ClassifierMode::Escalation + } else { + ClassifierMode::Capability + }) { + ClassifierMode::Capability => weak_target + .iter() + .chain(strong_target) + .map(String::as_str) + .collect(), + ClassifierMode::Escalation => strong_target + .iter() + .chain(weak_target) + .map(String::as_str) + .collect(), + ClassifierMode::Custom => targets.iter().flatten().map(String::as_str).collect(), + }, + Self::StageRouter { + capable_target, + efficient_target, + .. + } => vec![capable_target, efficient_target], + } + } + fn capabilities(&self) -> ModelCapabilities { use RouteConfig::*; match self { diff --git a/crates/switchyard-server/src/lib.rs b/crates/switchyard-server/src/lib.rs index edcdac37..8e39620e 100644 --- a/crates/switchyard-server/src/lib.rs +++ b/crates/switchyard-server/src/lib.rs @@ -36,6 +36,7 @@ use libsy::{Algorithm, LibsyError, RunObservation, RunObserver}; use parking_lot::Mutex; use serde::Deserialize; use serde_json::{Value, json}; +use switchyard_llm_client::TranslatingLlmClient; use switchyard_protocol::{Context, Decision, LlmClientError, Metadata, Request, Usage}; use tokio::net::{TcpListener, TcpSocket}; use tokio::task; @@ -97,12 +98,24 @@ struct ModelCapabilities { tool_calling: Option, } -/// A registered route: the libsy algorithm that serves it and the capabilities -/// advertised for it on `GET /v1/models`. One entry owns both so the routing -/// runtime and the model listing can never drift apart. +/// A registered algorithm route and its server-owned endpoint metadata. struct RouteEntry { algorithm: Arc, capabilities: ModelCapabilities, + count_tokens_target: Option, +} + +/// Exact upstream model used by the server's Anthropic token-count endpoint. +#[derive(Clone)] +struct CountTokensTarget { + model: String, + client: Arc, +} + +impl CountTokensTarget { + async fn count_tokens(&self, request: Request) -> Result { + self.client.count_tokens(&self.model, request).await + } } /// Shared server state used by all endpoint handlers. @@ -157,15 +170,22 @@ impl ServerState { Self::new_with_capabilities( routes .into_iter() - .map(|(model, algorithm)| (model, algorithm, ModelCapabilities::default())), + .map(|(model, algorithm)| (model, algorithm, ModelCapabilities::default(), None)), ) } fn new_with_capabilities( - routes: impl IntoIterator, ModelCapabilities)>, + routes: impl IntoIterator< + Item = ( + String, + Arc, + ModelCapabilities, + Option, + ), + >, ) -> ServerResult { let mut entries = BTreeMap::new(); - for (model, algorithm, capabilities) in routes { + for (model, algorithm, capabilities, count_tokens_target) in routes { let model = model.trim(); if model.is_empty() { return Err(ServerError::new("route model must not be empty")); @@ -173,6 +193,7 @@ impl ServerState { let entry = RouteEntry { algorithm, capabilities, + count_tokens_target, }; if entries.insert(model.to_string(), entry).is_some() { return Err(ServerError::new(format!("duplicate route model {model}"))); @@ -202,10 +223,8 @@ impl ServerState { self.routes.keys().map(String::as_str) } - fn algorithm_for_model(&self, model: &str) -> Option> { - self.routes - .get(model) - .map(|entry| Arc::clone(&entry.algorithm)) + fn route_for_model(&self, model: &str) -> Option<&RouteEntry> { + self.routes.get(model) } } @@ -463,11 +482,7 @@ async fn openai_responses( handle_endpoint(state, started, headers, body, WireFormat::OpenAiResponses).await } -/// Anthropic token counting. Resolves the route named by `model`, then does a -/// **direct passthrough** via [`Algorithm::count_tokens`] to that route's -/// Anthropic target — it does *not* run the routing cascade (count_tokens is a -/// pre-flight estimate with no routing decision). Unknown route → 404; a route -/// with no Anthropic target → 400. +/// Anthropic token counting against the route's explicitly configured target. async fn anthropic_count_tokens( State(state): State, headers: HeaderMap, @@ -477,7 +492,7 @@ async fn anthropic_count_tokens( Ok(body) => body, Err(message) => return anthropic_error_response(invalid_body_error(message)), }; - let (algorithm, request) = match resolve_route( + let (route, request) = match resolve_route( &state, metadata_from_headers(headers), body, @@ -486,27 +501,23 @@ async fn anthropic_count_tokens( Ok(resolved) => resolved, Err(response) => return anthropic_error_response(response), }; - anthropic_error_response(match algorithm.count_tokens(request).await { + let Some(target) = route.count_tokens_target.as_ref() else { + return anthropic_error_response(error_response( + StatusCode::BAD_REQUEST, + "route has no Anthropic target for token counting", + "invalid_request_error", + "count_tokens_unsupported", + )); + }; + anthropic_error_response(match target.count_tokens(request).await { Ok(payload) => (StatusCode::OK, Json(payload)).into_response(), Err(error) => count_tokens_error(error), }) } -/// Map a [`count_tokens`](Algorithm::count_tokens) failure to an HTTP response: -/// the route has no Anthropic target → 400, an upstream HTTP error → its own -/// status, anything else → 502. -fn count_tokens_error(error: LibsyError) -> Response { - // The one count_tokens-specific case is "no Anthropic target in the route"; - // every upstream/client failure gets the same mapping completions use. - match &error { - LibsyError::AlgorithmError { message } => error_response( - StatusCode::BAD_REQUEST, - message.clone(), - "invalid_request_error", - "count_tokens_unsupported", - ), - _ => algorithm_error(error), - } +/// Maps a token-count client failure with the same policy as a routed client call. +fn count_tokens_error(error: LlmClientError) -> Response { + client_error(&error) } async fn handle_endpoint( @@ -585,7 +596,7 @@ fn llm_json_body( /// Decode `body`, resolve the route named by its `model`, and build the /// [`Request`]. Shared by the completion and `count_tokens` handlers. Returns -/// the resolved algorithm and the built request — or an error [`Response`] +/// the resolved route and the built request — or an error [`Response`] /// (invalid body, empty `model` → 400, unknown route → 404). // Both callers immediately return the `Err(Response)` as the HTTP response, so // the large error type is intentional, not propagated up a call stack. @@ -595,7 +606,7 @@ fn resolve_route( metadata: Metadata, body: Value, wire_format: WireFormat, -) -> std::result::Result<(Arc, Request), Response> { +) -> std::result::Result<(&RouteEntry, Request), Response> { let llm_request = decode_request(wire_format, &body) .map_err(|error| invalid_body_error(error.to_string()))?; let requested_model = llm_request @@ -610,7 +621,7 @@ fn resolve_route( "invalid_request_error", ) })?; - let algorithm = state.algorithm_for_model(&requested_model).ok_or_else(|| { + let route = state.route_for_model(&requested_model).ok_or_else(|| { error_response( StatusCode::NOT_FOUND, format!("No route registered for model {requested_model}"), @@ -623,7 +634,7 @@ fn resolve_route( raw_request: Some(body), metadata: Some(metadata), }; - Ok((algorithm, request)) + Ok((route, request)) } async fn handle_llm_request( @@ -635,10 +646,11 @@ async fn handle_llm_request( routing_log_context: Option, ) -> Response { let cache_probe = state.track_cache_eligibility.then(|| prefix_probe(&body)); - let (algorithm, request) = match resolve_route(&state, metadata, body, wire_format) { + let (route, request) = match resolve_route(&state, metadata, body, wire_format) { Ok(resolved) => resolved, Err(response) => return response, }; + let algorithm = Arc::clone(&route.algorithm); let observer = stats_observer(state.stats.clone()); let (trace, response) = match algorithm .run_observed(Context::default(), request, Some(observer)) @@ -784,7 +796,11 @@ fn algorithm_error(error: LibsyError) -> Response { let LibsyError::ClientCall { source, .. } = &error else { return server_error(error.to_string()); }; - match source { + client_error(source) +} + +fn client_error(error: &LlmClientError) -> Response { + match error { LlmClientError::InvalidRequest { message } | LlmClientError::RequestTranslation(message) => error_response( StatusCode::BAD_REQUEST, diff --git a/crates/switchyard-server/tests/server.rs b/crates/switchyard-server/tests/server.rs index 42aac166..3b7b1100 100644 --- a/crates/switchyard-server/tests/server.rs +++ b/crates/switchyard-server/tests/server.rs @@ -980,10 +980,14 @@ base_url = "{base_url}" id = "real/opus" llm_client = "claude" +[targets.other] +id = "real/sonnet" +llm_client = "claude" + [routes.random] id = "switchyard/random" type = "random" -targets = ["strong"] +targets = ["other", "strong"] "#, base_url = upstream.base_url ))?; @@ -1044,15 +1048,14 @@ targets = ["weak"] ) .await?; assert_eq!(response.status, StatusCode::BAD_REQUEST); - // The route's picked target is OpenAI, so count_tokens (Anthropic-only) is - // unsupported for it. + // This route has no Anthropic-format target. assert_eq!( response.json()?, json!({ "type": "error", "error": { "type": "invalid_request_error", - "message": "no target supports count_tokens (needs an Anthropic upstream)" + "message": "route has no Anthropic target for token counting" } }) );