Skip to content

llm-router: let caller decide model provider #124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aswink wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

llm-router: let caller decide model provider #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 5 additions & 0 deletions crates/braintrust-llm-router/examples/custom_auth.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,11 +132,16 @@ async fn main() -> Result<()> {

println!(" Sending authenticated request to GPT-4...");
let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.unwrap();
let provider = aliases.first().unwrap();
let bytes = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await?;
Expand Down
11 changes: 10 additions & 1 deletion crates/braintrust-llm-router/examples/multi_provider.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,11 +79,16 @@ async fn main() -> Result<()> {
});

let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.unwrap();
let provider = aliases.first().unwrap();
match router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand All @@ -99,7 +104,6 @@ async fn main() -> Result<()> {
}
}

// Test Anthropic if available
if anthropic_key.is_some() {
println!("📍 Sending request to Claude...");
let model = "claude-3-5-haiku-20241022";
Expand All @@ -110,11 +114,16 @@ async fn main() -> Result<()> {
});

let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.unwrap();
let provider = aliases.first().unwrap();
match router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand Down
9 changes: 6 additions & 3 deletions crates/braintrust-llm-router/examples/simple.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@

use anyhow::Result;
use braintrust_llm_router::{
serde_json::json, ClientHeaders, ModelCatalog, OpenAIConfig, OpenAIProvider, ProviderFormat,
Router,
serde_json::json, ClientHeaders, Error, ModelCatalog, OpenAIConfig, OpenAIProvider,
ProviderFormat, Router,
};
use bytes::Bytes;
use serde_json::Value;
Expand Down Expand Up @@ -55,13 +55,16 @@ async fn main() -> Result<()> {
"max_tokens": 500
});

// Convert payload to bytes and send request
let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)?;
let provider = aliases.first().ok_or(Error::NoProvider(ProviderFormat::ChatCompletions))?;
let bytes = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await?;
Expand Down
10 changes: 10 additions & 0 deletions crates/braintrust-llm-router/examples/streaming.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,16 @@ async fn main() -> Result<()> {
});

let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.unwrap();
let provider = aliases.first().unwrap();
let mut stream = router
.complete_stream(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await?;
Expand Down Expand Up @@ -138,11 +143,16 @@ async fn main() -> Result<()> {
"stream": true
});
let body = Bytes::from(serde_json::to_vec(&payload)?);
let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.unwrap();
let provider = aliases.first().unwrap();
let stream = router
.complete_stream(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await?;
Expand Down
60 changes: 43 additions & 17 deletions crates/braintrust-llm-router/src/router.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,9 +106,9 @@ type ResolvedRoute<'a> = (
pub struct Router {
catalog: Arc<ModelCatalog>,
resolver: ModelResolver,
providers: HashMap<String, Arc<dyn Provider>>, // alias -> provider
formats: HashMap<ProviderFormat, String>, // format -> alias
auth_configs: HashMap<String, AuthConfig>, // alias -> auth
providers: HashMap<String, Arc<dyn Provider>>, // alias -> provider
format_providers: HashMap<ProviderFormat, Vec<String>>, // format -> ordered aliases
auth_configs: HashMap<String, AuthConfig>, // alias -> auth
retry_policy: RetryPolicy,
}

Expand All @@ -121,6 +121,27 @@ impl Router {
Arc::clone(&self.catalog)
}

/// Return the format and provider aliases for a model.
///
/// Callers use this to select a provider, then pass the chosen alias to
/// `complete` and `complete_stream`.
pub fn provider_aliases_for_model(
&self,
model: &str,
output_format: ProviderFormat,
) -> Result<(ProviderFormat, Vec<String>)> {
let (_, format, default) = self.resolver.resolve(model)?;
if format != output_format {
return Err(Error::NoProvider(output_format));
}
let aliases = self
.format_providers
.get(&format)
.cloned()
.unwrap_or_else(|| vec![default]);
Ok((format, aliases))
}

/// Execute a completion request with the given body bytes.
///
/// # Arguments
Expand All @@ -145,10 +166,11 @@ impl Router {
body: Bytes,
model: &str,
output_format: ProviderFormat,
provider_alias: &str,
client_headers: &ClientHeaders,
) -> Result<Bytes> {
let (provider, auth, spec, format, strategy) =
self.resolve_provider(model, output_format)?;
self.resolve_provider(model, output_format, provider_alias)?;
let payload = match lingua::transform_request(body.clone(), format, Some(&spec.model)) {
Ok(TransformResult::PassThrough(bytes)) => bytes,
Ok(TransformResult::Transformed { bytes, .. }) => bytes,
Expand Down Expand Up @@ -202,9 +224,11 @@ impl Router {
body: Bytes,
model: &str,
output_format: ProviderFormat,
provider_alias: &str,
client_headers: &ClientHeaders,
) -> Result<ResponseStream> {
let (provider, auth, spec, format, _) = self.resolve_provider(model, output_format)?;
let (provider, auth, spec, format, _) =
self.resolve_provider(model, output_format, provider_alias)?;
let payload = match lingua::transform_request(body.clone(), format, Some(&spec.model)) {
Ok(TransformResult::PassThrough(bytes)) => bytes,
Ok(TransformResult::Transformed { bytes, .. }) => bytes,
Expand All @@ -219,18 +243,14 @@ impl Router {
Ok(transform_stream(raw_stream, output_format))
}

pub fn provider_alias(&self, model: &str) -> Result<String> {
let (_, format, alias) = self.resolver.resolve(model)?;
Ok(self.formats.get(&format).cloned().unwrap_or(alias))
}

fn resolve_provider(
&self,
model: &str,
output_format: ProviderFormat,
provider_alias: &str,
) -> Result<ResolvedRoute<'_>> {
let (spec, catalog_format, alias) = self.resolver.resolve(model)?;
let alias = self.formats.get(&catalog_format).cloned().unwrap_or(alias);
let (spec, catalog_format, _default_alias) = self.resolver.resolve(model)?;
let alias = provider_alias.to_string();
let provider = self
.providers
.get(&alias)
Expand Down Expand Up @@ -318,7 +338,7 @@ impl Router {
pub struct RouterBuilder {
catalog: Option<Arc<ModelCatalog>>,
providers: HashMap<String, Arc<dyn Provider>>,
formats: HashMap<ProviderFormat, String>,
format_providers: HashMap<ProviderFormat, Vec<String>>,
auth_configs: HashMap<String, AuthConfig>,
retry_policy: RetryPolicy,
}
Expand All @@ -334,7 +354,7 @@ impl RouterBuilder {
Self {
catalog: None,
providers: HashMap::new(),
formats: HashMap::new(),
format_providers: HashMap::new(),
auth_configs: HashMap::new(),
retry_policy: RetryPolicy::default(),
}
Expand Down Expand Up @@ -362,7 +382,10 @@ impl RouterBuilder {
{
let alias = alias.into();
for format in provider.provider_formats() {
self.formats.insert(format, alias.clone());
self.format_providers
.entry(format)
.or_default()
.push(alias.clone());
}
self.providers.insert(alias, Arc::new(provider));
self
Expand All @@ -376,7 +399,10 @@ impl RouterBuilder {
) -> Self {
let alias = alias.into();
for format in provider.provider_formats() {
self.formats.insert(format, alias.clone());
self.format_providers
.entry(format)
.or_default()
.push(alias.clone());
}
self.providers.insert(alias, provider);
self
Expand Down Expand Up @@ -409,7 +435,7 @@ impl RouterBuilder {
catalog,
resolver,
providers: self.providers,
formats: self.formats,
format_providers: self.format_providers,
auth_configs: self.auth_configs,
retry_policy: self.retry_policy,
})
Expand Down
22 changes: 21 additions & 1 deletion crates/braintrust-llm-router/tests/router.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,11 +132,16 @@ async fn router_routes_to_stub_provider() {
]
}));

let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.expect("aliases");
let provider = aliases.first().expect("has provider");
let bytes = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand Down Expand Up @@ -187,11 +192,16 @@ async fn router_requires_auth_for_provider() {
"messages": [{"role": "user", "content": "Ping"}]
}));

let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.expect("aliases");
let provider = aliases.first().expect("has provider");
let err = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand Down Expand Up @@ -233,11 +243,16 @@ async fn router_reports_missing_provider() {
"messages": [{"role": "user", "content": "Ping"}]
}));

let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.expect("aliases");
let provider = aliases.first().expect("has default");
let err = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand All @@ -264,7 +279,6 @@ async fn router_propagates_validation_errors() {
.build()
.expect("router builds");

// Empty model should fail validation
let body = to_body(json!({
"model": "",
"messages": []
Expand All @@ -274,6 +288,7 @@ async fn router_propagates_validation_errors() {
body,
"",
ProviderFormat::ChatCompletions,
"stub",
&ClientHeaders::default(),
)
.await
Expand Down Expand Up @@ -384,11 +399,16 @@ async fn router_retries_and_propagates_terminal_error() {
"messages": [{"role": "user", "content": "Ping"}]
}));

let (_, aliases) = router
.provider_aliases_for_model(model, ProviderFormat::ChatCompletions)
.expect("aliases");
let provider = aliases.first().expect("has provider");
let err = router
.complete(
body,
model,
ProviderFormat::ChatCompletions,
provider,
&ClientHeaders::default(),
)
.await
Expand Down
Loading