Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

28 changes: 0 additions & 28 deletions crates/libsy-llm-client/src/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -391,34 +391,6 @@ mod tests {
);
}

#[test]
fn count_tokens_url_joins_every_base_url_shape() {
assert_eq!(
Backend::Anthropic(config("https://host")).count_tokens_url(),
"https://host/v1/messages/count_tokens"
);
assert_eq!(
Backend::Anthropic(config("https://host/v1")).count_tokens_url(),
"https://host/v1/messages/count_tokens"
);
assert_eq!(
Backend::Anthropic(config("https://host/v1/messages")).count_tokens_url(),
"https://host/v1/messages/count_tokens"
);
// Trailing slash is trimmed before the join.
assert_eq!(
Backend::Anthropic(config("https://host/v1/")).count_tokens_url(),
"https://host/v1/messages/count_tokens"
);
}

#[test]
fn only_anthropic_backend_is_anthropic() {
assert!(Backend::Anthropic(config("x")).is_anthropic());
assert!(!Backend::OpenAiChat(config("x")).is_anthropic());
assert!(!Backend::OpenAiResponses(config("x")).is_anthropic());
}

#[test]
fn wire_format_matches_variant() {
assert_eq!(
Expand Down
79 changes: 54 additions & 25 deletions crates/libsy-llm-client/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,34 @@ impl ModelConfig {
}
}

/// A model-bearing provider operation outside the normal completion endpoint.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AuxiliaryOperation {
/// Anthropic Messages input-token counting.
AnthropicCountTokens,
/// OpenAI Responses input-token counting.
ResponsesInputTokens,
/// OpenAI Responses compaction.
ResponsesCompact,
}

impl AuxiliaryOperation {
const fn wire_format(self) -> WireFormat {
match self {
Self::AnthropicCountTokens => WireFormat::AnthropicMessages,
Self::ResponsesInputTokens | Self::ResponsesCompact => WireFormat::OpenAiResponses,
}
}

fn url(self, backend: &Backend) -> String {
match self {
Self::AnthropicCountTokens => backend.count_tokens_url(),
Self::ResponsesInputTokens => format!("{}/input_tokens", backend.url()),
Self::ResponsesCompact => format!("{}/compact", backend.url()),
}
}
}

/// A client that dispatches neutral-IR requests to per-model HTTP backends.
///
/// Construct it with a list of [`ModelConfig`]s — one per model, each naming a
Expand Down Expand Up @@ -146,22 +174,27 @@ impl TranslatingLlmClient {
})
}

/// Whether `model` has an Anthropic backend that supports token counting.
pub fn supports_count_tokens(&self, model: &ModelId) -> bool {
self.backend_for(model, WireFormat::AnthropicMessages)
.is_some()
/// Whether `model` has a backend for `operation`.
pub fn supports_auxiliary(&self, model: &ModelId, operation: AuxiliaryOperation) -> bool {
self.backend_for(model, operation.wire_format()).is_some()
}

/// Counts input tokens with `model`'s Anthropic backend.
/// Calls a model-bearing auxiliary provider operation.
///
/// Returns an error when the model has no Anthropic backend or the upstream
/// Returns an error when the model has no compatible backend or the upstream
/// request fails or returns invalid JSON.
pub async fn count_tokens(&self, model: &ModelId, request: Request) -> Result<Value> {
let backend = self
.backend_for(model, WireFormat::AnthropicMessages)
.ok_or_else(|| LlmClientError::Configuration {
message: format!("model {model} has no Anthropic backend for count_tokens"),
})?;
pub async fn call_auxiliary(
&self,
model: &ModelId,
request: Request,
operation: AuxiliaryOperation,
) -> Result<Value> {
let wire_format = operation.wire_format();
let backend =
self.backend_for(model, wire_format)
.ok_or_else(|| LlmClientError::Configuration {
message: format!("model {model} has no backend for {operation:?}"),
})?;
let Request {
mut llm_request,
metadata,
Expand All @@ -171,20 +204,17 @@ impl TranslatingLlmClient {
let http_response = self
.send_encoded(
backend,
WireFormat::AnthropicMessages,
wire_format,
llm_request,
metadata.as_ref(),
model,
UpstreamEndpoint::CountTokens,
UpstreamEndpoint::Auxiliary(operation),
)
.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(),
});
}
let EncodedResponse::Buffered { body, .. } = http_response else {
return Err(LlmClientError::InvalidRequest {
message: "auxiliary endpoints do not support streaming".to_string(),
});
};
serde_json::from_slice(&body).map_err(|error| LlmClientError::InvalidResponse {
source: Box::new(error),
Expand All @@ -200,8 +230,7 @@ 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`](Self::count_tokens) (which POSTs to the `count_tokens`
/// URL and returns the raw JSON).
/// the model-bearing auxiliary operations, which return raw JSON.
async fn send_encoded(
&self,
backend: &Backend,
Expand Down Expand Up @@ -562,14 +591,14 @@ impl RoutedLlmClient for TranslatingLlmClient {
#[derive(Clone, Copy)]
enum UpstreamEndpoint {
Completion,
CountTokens,
Auxiliary(AuxiliaryOperation),
}

impl UpstreamEndpoint {
fn url(self, backend: &Backend) -> String {
match self {
UpstreamEndpoint::Completion => backend.url(),
UpstreamEndpoint::CountTokens => backend.count_tokens_url(),
UpstreamEndpoint::Auxiliary(operation) => operation.url(backend),
}
}

Expand Down
2 changes: 1 addition & 1 deletion crates/libsy-llm-client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ pub mod raw;
pub mod run;

pub use backend::{Backend, DEFAULT_MAX_RETRIES, HttpBackendConfig};
pub use client::{ModelConfig, TranslatingLlmClient};
pub use client::{AuxiliaryOperation, ModelConfig, TranslatingLlmClient};
pub use error::{LlmClientError, Result};
pub use observation::{LlmCallObservation, RunObservation, RunObserver};
pub use raw::RawResponse;
Expand Down
74 changes: 57 additions & 17 deletions crates/switchyard-runner/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@ use serde::de::DeserializeOwned;
use serde::{Deserialize, Deserializer};
use serde_json::Value;
use switchyard_llm_client::{
Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig,
AuxiliaryOperation, Backend, ClientRouter, DEFAULT_MAX_RETRIES, HttpBackendConfig, ModelConfig,
TranslatingLlmClient,
};
use switchyard_protocol::{ModelId, RoutedLlmClient, WireFormat};

use crate::{
AlgorithmSpec, CallerAuthKind, CountTokensTarget, DecisionTarget, ModelCapabilities, Route,
AlgorithmSpec, AuxiliaryTarget, CallerAuthKind, DecisionTarget, ModelCapabilities, Route,
Runner, RunnerError,
};

Expand Down Expand Up @@ -54,6 +54,7 @@ pub(crate) fn runner_from_toml(source: &str) -> RunnerResult<Runner> {
#[serde(deny_unknown_fields)]
pub(crate) struct DeploymentConfig {
schema_version: u32,
fallback_client: Option<String>,
#[serde(default)]
llm_clients: BTreeMap<String, LlmClientConfig>,
targets: BTreeMap<String, TargetConfig>,
Expand Down Expand Up @@ -165,6 +166,7 @@ impl DeploymentConfig {

let clients = self.build_clients()?;
let targets = self.build_targets();
let fallback_base_url = self.fallback_base_url()?;
let mut routes = Vec::with_capacity(self.routes.len());
for (route_name, config) in &self.routes {
validate_value("route name", route_name)?;
Expand All @@ -188,7 +190,10 @@ impl DeploymentConfig {
.map_err(|error| RunnerError::configuration_source(error.to_string(), error))?;
let (route_clients, caller_auth) =
self.build_route_clients(route_name, config, &clients)?;
let count_tokens_target = self.build_count_tokens_target(config, &clients);
let anthropic_auxiliary_target =
self.build_anthropic_auxiliary_target(config, &clients);
let responses_auxiliary_target =
self.build_responses_auxiliary_target(config, &clients);
let decision_targets = config
.routing_target_names()
.into_iter()
Expand All @@ -199,12 +204,14 @@ impl DeploymentConfig {
route_clients,
caller_auth,
capabilities,
count_tokens_target,
anthropic_auxiliary_target,
responses_auxiliary_target,
decision_targets,
);
routes.push((config.id.clone(), route));
}
Ok(Runner::new(routes))
let runner = Runner::new(routes).with_fallback_url(fallback_base_url);
Ok(runner)
}

fn build_clients(&self) -> RunnerResult<BTreeMap<String, Arc<TranslatingLlmClient>>> {
Expand Down Expand Up @@ -291,27 +298,60 @@ impl DeploymentConfig {
Ok((ClientRouter::new(by_model), caller_auth))
}

fn build_count_tokens_target(
fn fallback_base_url(&self) -> RunnerResult<Option<String>> {
let Some(name) = &self.fallback_client else {
return Ok(None);
};
let config = self.llm_clients.get(name).ok_or_else(|| {
RunnerError::configuration(format!(
"fallback_client references unknown llm client {name}"
))
})?;
Ok(Some(config.base_url.as_str().to_string()))
}

fn build_anthropic_auxiliary_target(
&self,
route: &RouteConfig,
clients: &BTreeMap<String, Arc<TranslatingLlmClient>>,
) -> Option<CountTokensTarget> {
) -> Option<AuxiliaryTarget> {
route
.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,
))
let target = self.build_auxiliary_target(
name,
clients,
AuxiliaryOperation::AnthropicCountTokens,
)?;
Some((count_tokens_priority(name, &target.model), index, target))
})
.min_by_key(|(priority, index, _, _)| (*priority, *index))
.map(|(_, _, target, client)| CountTokensTarget {
.min_by_key(|(priority, index, _)| (*priority, *index))
.map(|(_, _, target)| target)
}

fn build_responses_auxiliary_target(
&self,
route: &RouteConfig,
clients: &BTreeMap<String, Arc<TranslatingLlmClient>>,
) -> Option<AuxiliaryTarget> {
route.routing_target_names().into_iter().find_map(|name| {
self.build_auxiliary_target(name, clients, AuxiliaryOperation::ResponsesInputTokens)
})
}

fn build_auxiliary_target(
&self,
name: &str,
clients: &BTreeMap<String, Arc<TranslatingLlmClient>>,
operation: AuxiliaryOperation,
) -> Option<AuxiliaryTarget> {
let target = self.targets.get(name)?;
let client = clients.get(&target.llm_client)?;
client
.supports_auxiliary(&target.id, operation)
.then(|| AuxiliaryTarget {
model: target.id.clone(),
client: client.clone(),
})
Expand Down
4 changes: 2 additions & 2 deletions crates/switchyard-runner/src/failure.rs
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ impl RunnerError {
),
Self::UnknownRouteModel(_)
| Self::IncompatibleCallerFormat(_)
| Self::CountTokensUnsupported => summary(
| Self::AuxiliaryUnsupported => summary(
RouteErrorKind::InvalidRequest,
RouteErrorPhase::BeforeResponse,
None,
Expand Down Expand Up @@ -329,7 +329,7 @@ mod tests {
RouteErrorKind::Configuration
));

let unsupported = RunnerError::CountTokensUnsupported;
let unsupported = RunnerError::AuxiliaryUnsupported;
assert!(matches!(
unsupported.execution_error_summary().kind,
RouteErrorKind::InvalidRequest
Expand Down
2 changes: 1 addition & 1 deletion crates/switchyard-runner/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,6 @@ pub use algorithm::{
};
pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary};
pub use route::{
CallerAuthKind, CountTokensTarget, ModelCapabilities, Route, RunOutput, RunnerError,
AuxiliaryTarget, CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError,
};
pub use runner::{DecisionDescription, DecisionTarget, ModelInfo, Runner};
Loading
Loading