Skip to content

forward client headers for providers #62

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

Open
knjiang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

forward client headers for providers #62

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
12 changes: 10 additions & 2 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 @@ -15,7 +15,8 @@

use anyhow::Result;
use braintrust_llm_router::{
serde_json::json, AuthConfig, OpenAIConfig, OpenAIProvider, ProviderFormat, Router,
serde_json::json, AuthConfig, ClientHeaders, OpenAIConfig, OpenAIProvider, ProviderFormat,
Router,
};
use bytes::Bytes;
use serde_json::Value;
Expand Down Expand Up @@ -127,7 +128,14 @@ async fn main() -> Result<()> {

println!(" Sending authenticated request to GPT-4...");
let body = Bytes::from(serde_json::to_vec(&payload)?);
let bytes = router.complete(body, model, ProviderFormat::OpenAI).await?;
let bytes = router
.complete(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await?;
let response: Value = serde_json::from_slice(&bytes)?;

if let Some(text) = extract_assistant_text(&response) {
Expand Down
24 changes: 20 additions & 4 deletions 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 @@ -10,8 +10,8 @@

use anyhow::Result;
use braintrust_llm_router::{
serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, OpenAIConfig, OpenAIProvider,
ProviderFormat, Router,
serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, ClientHeaders, OpenAIConfig,
OpenAIProvider, ProviderFormat, Router,
};
use bytes::Bytes;
use serde_json::Value;
Expand Down Expand Up @@ -77,7 +77,15 @@ async fn main() -> Result<()> {
});

let body = Bytes::from(serde_json::to_vec(&payload)?);
match router.complete(body, model, ProviderFormat::OpenAI).await {
match router
.complete(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await
{
Ok(bytes) => {
if let Ok(response) = serde_json::from_slice::<Value>(&bytes) {
if let Some(text) = extract_assistant_text(&response) {
Expand All @@ -100,7 +108,15 @@ async fn main() -> Result<()> {
});

let body = Bytes::from(serde_json::to_vec(&payload)?);
match router.complete(body, model, ProviderFormat::OpenAI).await {
match router
.complete(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await
{
Ok(bytes) => {
if let Ok(response) = serde_json::from_slice::<Value>(&bytes) {
if let Some(text) = extract_assistant_text(&response) {
Expand Down
11 changes: 9 additions & 2 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,7 +7,7 @@

use anyhow::Result;
use braintrust_llm_router::{
serde_json::json, OpenAIConfig, OpenAIProvider, ProviderFormat, Router,
serde_json::json, ClientHeaders, OpenAIConfig, OpenAIProvider, ProviderFormat, Router,
};
use bytes::Bytes;
use serde_json::Value;
Expand Down Expand Up @@ -53,7 +53,14 @@ async fn main() -> Result<()> {

// Convert payload to bytes and send request
let body = Bytes::from(serde_json::to_vec(&payload)?);
let bytes = router.complete(body, model, ProviderFormat::OpenAI).await?;
let bytes = router
.complete(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await?;
let response: Value = serde_json::from_slice(&bytes)?;

println!("📝 Response:\n");
Expand Down
18 changes: 14 additions & 4 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 @@ -7,8 +7,8 @@

use anyhow::Result;
use braintrust_llm_router::{
serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, ProviderFormat,
ResponseStream, Router,
serde_json::json, AnthropicConfig, AnthropicProvider, AuthConfig, ClientHeaders,
ProviderFormat, ResponseStream, Router,
};
use bytes::Bytes;
use futures::StreamExt;
Expand Down Expand Up @@ -54,7 +54,12 @@ async fn main() -> Result<()> {

let body = Bytes::from(serde_json::to_vec(&payload)?);
let mut stream = router
.complete_stream(body, model, ProviderFormat::OpenAI)
.complete_stream(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await?;

// Process the stream
Expand Down Expand Up @@ -128,7 +133,12 @@ async fn main() -> Result<()> {
});
let body = Bytes::from(serde_json::to_vec(&payload)?);
let stream = router
.complete_stream(body, model, ProviderFormat::OpenAI)
.complete_stream(
body,
model,
ProviderFormat::OpenAI,
&ClientHeaders::default(),
)
.await?;
streams.push((model.to_string(), stream));
}
Expand Down
6 changes: 3 additions & 3 deletions crates/braintrust-llm-router/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ pub use lingua::ProviderFormat;
pub use lingua::{FinishReason, UniversalStreamChoice, UniversalStreamChunk};
pub use providers::{
is_openai_compatible, openai_compatible_endpoint, AnthropicConfig, AnthropicProvider,
AzureConfig, AzureProvider, BedrockConfig, BedrockProvider, GoogleConfig, GoogleProvider,
MistralConfig, MistralProvider, OpenAICompatibleEndpoint, OpenAIConfig, OpenAIProvider,
OpenAIResponsesProvider, Provider, VertexConfig, VertexProvider,
AzureConfig, AzureProvider, BedrockConfig, BedrockProvider, ClientHeaders, GoogleConfig,
GoogleProvider, MistralConfig, MistralProvider, OpenAICompatibleEndpoint, OpenAIConfig,
OpenAIProvider, OpenAIResponsesProvider, Provider, VertexConfig, VertexProvider,
};
pub use retry::{RetryPolicy, RetryStrategy};
pub use router::{create_provider, extract_request_hints, RequestHints, Router, RouterBuilder};
Expand Down
43 changes: 23 additions & 20 deletions crates/braintrust-llm-router/src/providers/anthropic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,25 @@ use std::time::Duration;

use async_trait::async_trait;
use bytes::Bytes;
use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
use reqwest::header::{HeaderMap, HeaderValue};
use reqwest::{Client, StatusCode, Url};

use crate::auth::AuthConfig;
use crate::catalog::ModelSpec;
use crate::client::{default_client, ClientSettings};
use crate::error::{Error, Result, UpstreamHttpError};
use crate::providers::ClientHeaders;
use crate::streaming::{single_bytes_stream, sse_stream, RawResponseStream};
use lingua::ProviderFormat;

const ANTHROPIC_VERSION: &str = "anthropic-version";
const ANTHROPIC_BETA: &str = "anthropic-beta";
const STRUCTURED_OUTPUTS_BETA: &str = "structured-outputs-2025-11-13";

#[derive(Debug, Clone)]
pub struct AnthropicConfig {
pub endpoint: Url,
pub version: String,
pub beta: Option<String>,
pub timeout: Option<Duration>,
}

Expand All @@ -26,7 +30,6 @@ impl Default for AnthropicConfig {
endpoint: Url::parse("https://api.anthropic.com/v1/")
.expect("valid Anthropic endpoint"),
version: "2023-06-01".to_string(),
beta: None,
timeout: None,
}
}
Expand Down Expand Up @@ -56,7 +59,6 @@ impl AnthropicProvider {
///
/// Extracts Anthropic-specific options from metadata:
/// - `version`: Anthropic API version (defaults to "2023-06-01")
/// - `beta`: Beta feature flag
pub fn from_config(
endpoint: Option<&Url>,
timeout: Option<Duration>,
Expand All @@ -74,9 +76,6 @@ impl AnthropicProvider {
if let Some(version) = metadata.get("version").and_then(Value::as_str) {
config.version = version.to_string();
}
if let Some(beta) = metadata.get("beta").and_then(Value::as_str) {
config.beta = Some(beta.to_string());
}

Self::new(config)
}
Expand All @@ -88,18 +87,23 @@ impl AnthropicProvider {
.expect("join messages path")
}

fn apply_headers(&self, headers: &mut HeaderMap) {
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
fn build_headers(&self, client_headers: &ClientHeaders) -> HeaderMap {
let mut headers = client_headers.to_json_headers();

headers.insert(
"anthropic-version",
ANTHROPIC_VERSION,
HeaderValue::from_str(&self.config.version).expect("version header"),
);
if let Some(beta) = &self.config.beta {

// Respect caller override: only set default if missing.
if !headers.contains_key(ANTHROPIC_BETA) {
Comment on lines +98 to +99
Copy link
Contributor Author

@knjiang knjiang Jan 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

auto-adding structured outputs since it seems useful esp. for chat completions response_format -> anthropic.

also seems hard coded in their sdk -> https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/resources/beta/messages/messages.py#L1106

headers.insert(
"anthropic-beta",
HeaderValue::from_str(beta).unwrap_or_else(|_| HeaderValue::from_static("")),
ANTHROPIC_BETA,
HeaderValue::from_static(STRUCTURED_OUTPUTS_BETA),
);
}

headers
}
}

Expand All @@ -118,6 +122,7 @@ impl crate::providers::Provider for AnthropicProvider {
payload: Bytes,
auth: &AuthConfig,
_spec: &ModelSpec,
client_headers: &ClientHeaders,
) -> Result<Bytes> {
let url = self.messages_url();

Expand All @@ -129,8 +134,7 @@ impl crate::providers::Provider for AnthropicProvider {
"sending request to Anthropic"
);

let mut headers = HeaderMap::new();
self.apply_headers(&mut headers);
let mut headers = self.build_headers(client_headers);
auth.apply_headers(&mut headers)?;

let response = self
Expand Down Expand Up @@ -176,9 +180,10 @@ impl crate::providers::Provider for AnthropicProvider {
payload: Bytes,
auth: &AuthConfig,
spec: &ModelSpec,
client_headers: &ClientHeaders,
) -> Result<RawResponseStream> {
if !spec.supports_streaming {
let response = self.complete(payload, auth, spec).await?;
let response = self.complete(payload, auth, spec, client_headers).await?;
return Ok(single_bytes_stream(response));
}

Expand All @@ -194,8 +199,7 @@ impl crate::providers::Provider for AnthropicProvider {
"sending streaming request to Anthropic"
);

let mut headers = HeaderMap::new();
self.apply_headers(&mut headers);
let mut headers = self.build_headers(client_headers);
auth.apply_headers(&mut headers)?;

let response = self
Expand Down Expand Up @@ -243,8 +247,7 @@ impl crate::providers::Provider for AnthropicProvider {
.endpoint
.join("models")
.expect("join models path");
let mut headers = HeaderMap::new();
self.apply_headers(&mut headers);
let mut headers = self.build_headers(&ClientHeaders::default());
auth.apply_headers(&mut headers)?;

let response = self.client.get(url).headers(headers).send().await?;
Expand Down
20 changes: 13 additions & 7 deletions crates/braintrust-llm-router/src/providers/azure.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,14 @@ use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use lingua::serde_json::Value;
use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
use reqwest::header::HeaderMap;
use reqwest::{Client, StatusCode, Url};

use crate::auth::AuthConfig;
use crate::catalog::ModelSpec;
use crate::client::{default_client, ClientSettings};
use crate::error::{Error, Result, UpstreamHttpError};
use crate::providers::ClientHeaders;
use crate::streaming::{single_bytes_stream, sse_stream, RawResponseStream};
use lingua::ProviderFormat;

Expand Down Expand Up @@ -143,7 +144,13 @@ impl crate::providers::Provider for AzureProvider {
ProviderFormat::OpenAI
}

async fn complete(&self, payload: Bytes, auth: &AuthConfig, spec: &ModelSpec) -> Result<Bytes> {
async fn complete(
&self,
payload: Bytes,
auth: &AuthConfig,
spec: &ModelSpec,
client_headers: &ClientHeaders,
) -> Result<Bytes> {
let url = self.chat_url(&spec.model)?;

#[cfg(feature = "tracing")]
Expand All @@ -154,8 +161,7 @@ impl crate::providers::Provider for AzureProvider {
"sending request to Azure"
);

let mut headers = HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
let mut headers = self.build_headers(client_headers);
auth.apply_headers(&mut headers)?;

let response = self
Expand Down Expand Up @@ -201,9 +207,10 @@ impl crate::providers::Provider for AzureProvider {
payload: Bytes,
auth: &AuthConfig,
spec: &ModelSpec,
client_headers: &ClientHeaders,
) -> Result<RawResponseStream> {
if !spec.supports_streaming {
let response = self.complete(payload, auth, spec).await?;
let response = self.complete(payload, auth, spec, client_headers).await?;
return Ok(single_bytes_stream(response));
}

Expand All @@ -219,8 +226,7 @@ impl crate::providers::Provider for AzureProvider {
"sending streaming request to Azure"
);

let mut headers = HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
let mut headers = self.build_headers(client_headers);
auth.apply_headers(&mut headers)?;

let response = self
Expand Down
Loading