Skip to content

Commit d1eec00

Browse files
committed
Avoid duplicate entries in router builder
Currently the router builder silently drops duplicate entries, solve this by returning an error if the builder sees that. Also explicitly set default format -> provider mappings, and also error if there are duplicate entries in the builder there. Also clean up some override client stuff.
1 parent a6a6904 commit d1eec00

18 files changed

Lines changed: 269 additions & 120 deletions

File tree

crates/braintrust-llm-router/README.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,11 @@ let stream = router.complete_stream(body, "gpt-4", ProviderFormat::ChatCompletio
123123
### Authentication
124124

125125
```rust
126-
use braintrust_llm_router::AuthConfig;
126+
use braintrust_llm_router::{api_key_auth, AuthConfig};
127127

128128
// OpenAI-style (Authorization: Bearer <key>)
129+
api_key_auth("sk-...")
130+
// Or equivalently:
129131
AuthConfig::ApiKey {
130132
key: "sk-...".into(),
131133
header: Some("authorization".into()),

crates/braintrust-llm-router/examples/custom_auth.rs

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
1616
use anyhow::Result;
1717
use braintrust_llm_router::{
18-
serde_json::json, AuthConfig, ClientHeaders, OpenAIConfig, OpenAIProvider, ProviderFormat,
19-
Router,
18+
api_key_auth, serde_json::json, AuthConfig, ClientHeaders, OpenAIConfig, OpenAIProvider,
19+
ProviderFormat, Router,
2020
};
2121
use bytes::Bytes;
2222
use serde_json::Value;
@@ -119,8 +119,12 @@ async fn main() -> Result<()> {
119119

120120
let router = Router::builder()
121121
.with_catalog(catalog)
122-
.add_provider("openai", openai_provider)
123-
.add_api_key("openai", api_key) // Convenience method
122+
.add_provider(
123+
"openai",
124+
openai_provider,
125+
api_key_auth(&api_key),
126+
vec![ProviderFormat::ChatCompletions],
127+
)
124128
.build()?;
125129

126130
let model = "gpt-4";
@@ -132,7 +136,7 @@ async fn main() -> Result<()> {
132136

133137
println!(" Sending authenticated request to GPT-4...");
134138
let body = Bytes::from(serde_json::to_vec(&payload)?);
135-
let bytes = router
139+
let bytes: Bytes = router
136140
.complete(
137141
body,
138142
model,

crates/braintrust-llm-router/examples/multi_provider.rs

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010
1111
use anyhow::Result;
1212
use braintrust_llm_router::{
13-
serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, ClientHeaders, ModelCatalog,
14-
OpenAIConfig, OpenAIProvider, ProviderFormat, Router,
13+
api_key_auth, serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, ClientHeaders,
14+
ModelCatalog, OpenAIConfig, OpenAIProvider, ProviderFormat, Router,
1515
};
1616
use bytes::Bytes;
1717
use serde_json::Value;
@@ -40,25 +40,28 @@ async fn main() -> Result<()> {
4040
// Add OpenAI provider if key is available
4141
if let Some(key) = &openai_key {
4242
let openai_provider = OpenAIProvider::new(OpenAIConfig::default())?;
43-
builder = builder
44-
.add_provider("openai", openai_provider)
45-
.add_api_key("openai", key.clone());
43+
builder = builder.add_provider(
44+
"openai",
45+
openai_provider,
46+
api_key_auth(key),
47+
vec![ProviderFormat::ChatCompletions],
48+
);
4649
println!("✅ OpenAI provider configured");
4750
}
4851

4952
// Add Anthropic provider if key is available
5053
if let Some(key) = &anthropic_key {
5154
let anthropic_provider = AnthropicProvider::new(AnthropicConfig::default())?;
52-
builder = builder
53-
.add_provider("anthropic", anthropic_provider)
54-
.add_auth(
55-
"anthropic",
56-
AuthConfig::ApiKey {
57-
key: key.clone(),
58-
header: Some("x-api-key".into()),
59-
prefix: None,
60-
},
61-
);
55+
builder = builder.add_provider(
56+
"anthropic",
57+
anthropic_provider,
58+
AuthConfig::ApiKey {
59+
key: key.clone(),
60+
header: Some("x-api-key".into()),
61+
prefix: None,
62+
},
63+
vec![ProviderFormat::ChatCompletions],
64+
);
6265
println!("✅ Anthropic provider configured");
6366
}
6467

crates/braintrust-llm-router/examples/simple.rs

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
88
use anyhow::Result;
99
use braintrust_llm_router::{
10-
serde_json::json, ClientHeaders, ModelCatalog, OpenAIConfig, OpenAIProvider, ProviderFormat,
11-
Router,
10+
api_key_auth, serde_json::json, ClientHeaders, ModelCatalog, OpenAIConfig, OpenAIProvider,
11+
ProviderFormat, Router,
1212
};
1313
use bytes::Bytes;
1414
use serde_json::Value;
@@ -37,8 +37,12 @@ async fn main() -> Result<()> {
3737

3838
let router = Router::builder()
3939
.with_catalog(catalog)
40-
.add_provider("openai", openai_provider)
41-
.add_api_key("openai", api_key)
40+
.add_provider(
41+
"openai",
42+
openai_provider,
43+
api_key_auth(&api_key),
44+
vec![ProviderFormat::ChatCompletions],
45+
)
4246
.build()?;
4347

4448
// Simple chat completion
@@ -57,7 +61,7 @@ async fn main() -> Result<()> {
5761

5862
// Convert payload to bytes and send request
5963
let body = Bytes::from(serde_json::to_vec(&payload)?);
60-
let bytes = router
64+
let bytes: Bytes = router
6165
.complete(
6266
body,
6367
model,
@@ -91,7 +95,7 @@ async fn main() -> Result<()> {
9195
// Print model information
9296
println!("\n🔍 Model Information:");
9397
if let Some(spec) = router.catalog().get("gpt-4") {
94-
let spec = spec.as_ref();
98+
let spec: &braintrust_llm_router::ModelSpec = spec.as_ref();
9599
println!(
96100
" Display Name: {}",
97101
spec.display_name.as_deref().unwrap_or("GPT-4")

crates/braintrust-llm-router/examples/streaming.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,15 @@ async fn main() -> Result<()> {
3030

3131
let router = Router::builder()
3232
.with_catalog(catalog)
33-
.add_provider("anthropic", anthropic_provider)
34-
.add_auth(
33+
.add_provider(
3534
"anthropic",
35+
anthropic_provider,
3636
AuthConfig::ApiKey {
3737
key: anthropic_api_key,
3838
header: Some("x-api-key".into()),
3939
prefix: None,
4040
},
41+
vec![ProviderFormat::ChatCompletions],
4142
)
4243
.build()?;
4344

@@ -59,7 +60,7 @@ async fn main() -> Result<()> {
5960
});
6061

6162
let body = Bytes::from(serde_json::to_vec(&payload)?);
62-
let mut stream = router
63+
let mut stream: ResponseStream = router
6364
.complete_stream(
6465
body,
6566
model,

crates/braintrust-llm-router/src/auth/mod.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,15 @@ pub enum AuthConfig {
5050
},
5151
}
5252

53+
/// Convenience function to create an API key authentication configuration.
54+
pub fn api_key_auth(api_key: &str) -> AuthConfig {
55+
AuthConfig::ApiKey {
56+
key: api_key.to_string(),
57+
header: Some("authorization".into()),
58+
prefix: Some("Bearer".into()),
59+
}
60+
}
61+
5362
impl AuthConfig {
5463
pub fn auth_type(&self) -> AuthType {
5564
match self {

crates/braintrust-llm-router/src/client.rs

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,22 +40,16 @@ pub fn build_client(settings: &ClientSettings) -> Result<Client> {
4040
}
4141

4242
pub fn build_middleware_client(settings: &ClientSettings) -> Result<ClientWithMiddleware> {
43+
if let Some(existing) = OVERRIDE_CLIENT.read().clone() {
44+
return Ok(existing);
45+
}
4346
let client = build_client(settings)?;
4447
Ok(reqwest_middleware::ClientBuilder::new(client).build())
4548
}
4649

4750
static OVERRIDE_CLIENT: Lazy<RwLock<Option<ClientWithMiddleware>>> =
4851
Lazy::new(|| RwLock::new(None));
4952

50-
pub fn override_client() -> Result<ClientWithMiddleware> {
51-
if let Some(existing) = OVERRIDE_CLIENT.read().clone() {
52-
return Ok(existing);
53-
}
54-
let client = build_middleware_client(&ClientSettings::default())?;
55-
*OVERRIDE_CLIENT.write() = Some(client.clone());
56-
Ok(client)
57-
}
58-
5953
pub fn set_override_client(client: ClientWithMiddleware) {
6054
*OVERRIDE_CLIENT.write() = Some(client);
6155
}

crates/braintrust-llm-router/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ mod streaming;
1414
pub use lingua::serde_json;
1515

1616
pub use auth::{
17+
api_key_auth,
1718
azure::{AzureEntraCredentials, AzureEntraTokenManager},
1819
databricks::{DatabricksCredentials, DatabricksTokenManager},
1920
google::{GoogleServiceAccountConfig, GoogleTokenManager, ServiceAccountKey},

crates/braintrust-llm-router/src/providers/anthropic.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use reqwest_middleware::ClientWithMiddleware;
88

99
use crate::auth::AuthConfig;
1010
use crate::catalog::ModelSpec;
11-
use crate::client::{build_middleware_client, override_client, ClientSettings};
11+
use crate::client::{build_middleware_client, ClientSettings};
1212
use crate::error::{Error, Result, UpstreamHttpError};
1313
use crate::providers::ClientHeaders;
1414
use crate::streaming::{single_bytes_stream, sse_stream, RawResponseStream};
@@ -48,7 +48,7 @@ impl AnthropicProvider {
4848
if let Some(timeout) = config.timeout {
4949
settings.request_timeout = timeout;
5050
}
51-
let client = override_client().or_else(|_| build_middleware_client(&settings))?;
51+
let client = build_middleware_client(&settings)?;
5252
Ok(Self { client, config })
5353
}
5454

crates/braintrust-llm-router/src/providers/azure.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use reqwest_middleware::ClientWithMiddleware;
99

1010
use crate::auth::AuthConfig;
1111
use crate::catalog::ModelSpec;
12-
use crate::client::{build_middleware_client, override_client, ClientSettings};
12+
use crate::client::{build_middleware_client, ClientSettings};
1313
use crate::error::{Error, Result, UpstreamHttpError};
1414
use crate::providers::ClientHeaders;
1515
use crate::streaming::{single_bytes_stream, sse_stream, RawResponseStream};
@@ -49,7 +49,7 @@ impl AzureProvider {
4949
if let Some(timeout) = config.timeout {
5050
settings.request_timeout = timeout;
5151
}
52-
let client = override_client().or_else(|_| build_middleware_client(&settings))?;
52+
let client = build_middleware_client(&settings)?;
5353
Ok(Self { client, config })
5454
}
5555

0 commit comments

Comments
 (0)