use async_trait::async_trait;
use http::HeaderMap;
use lellm_core::{ChatRequest, ChatResponse, LlmError};
use secrecy::ExposeSecret;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio_stream::StreamExt;
use crate::{LlmProvider, ProviderEvent, ProviderStream};
use super::codec::{
Capabilities, CodecRequest, ProviderBuildError, ProviderEnvError, ProviderExtension,
ProviderMeta, validate_capabilities,
};
use super::stream::{EventSink, StreamEvent};
#[derive(Debug)]
pub(crate) struct RawResponse {
pub status: u16,
#[allow(dead_code)]
pub headers: HeaderMap,
pub body: bytes::Bytes,
}
struct ChannelSink {
tx: tokio::sync::mpsc::Sender<Result<ProviderEvent, LlmError>>,
dropped: Arc<AtomicU64>,
}
impl EventSink for ChannelSink {
async fn emit(&mut self, event: StreamEvent) -> bool {
if event.is_critical() {
let mapped = map_stream_event(event);
match self.tx.send(mapped).await {
Ok(()) => true,
Err(e) => {
tracing::error!(
error = %e,
"critical stream event lost: channel closed"
);
false
}
}
} else {
let mapped = map_stream_event(event);
match self.tx.try_send(mapped) {
Ok(()) => true,
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
self.dropped.fetch_add(1, Ordering::Relaxed);
true
}
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => false,
}
}
}
fn is_closed(&self) -> bool {
self.tx.is_closed()
}
}
fn map_stream_event(event: StreamEvent) -> Result<ProviderEvent, LlmError> {
match event {
StreamEvent::Start { model } => Ok(ProviderEvent::Start { model }),
StreamEvent::Token { token } => Ok(ProviderEvent::Token { token }),
StreamEvent::ThinkingDelta { thinking, redacted } => {
Ok(ProviderEvent::ThinkingDelta { thinking, redacted })
}
StreamEvent::Error(e) => Err(e),
StreamEvent::ResponseComplete { tool_calls, usage } => {
Ok(ProviderEvent::ResponseComplete { tool_calls, usage })
}
}
}
#[allow(private_bounds)]
pub struct CodecProvider<C: ProviderExtension> {
codec: Arc<C>,
config: ProviderConfig,
client: reqwest::Client,
extra_headers: HeaderMap,
}
#[allow(private_bounds)]
impl<C: ProviderExtension> CodecProvider<C> {
pub(crate) fn new(codec: C, config: ProviderConfig) -> Self {
Self::builder(codec)
.base_url(config.base_url.as_str())
.auth(config.auth)
.connect_timeout(config.connect_timeout)
.timeout(config.timeout)
.idle_timeout(config.idle_timeout)
.build()
.expect("ProviderConfig base_url was already validated")
}
pub fn load(codec: C) -> Result<Self, ProviderBuildError> {
let config = ProviderConfig::load(&codec)?;
Ok(Self::new(codec, config))
}
pub fn builder(codec: C) -> ProviderBuilder<C> {
ProviderBuilder::new(codec)
}
fn validate_request(&self, req: &ChatRequest) -> Result<(), LlmError> {
for msg in &req.messages {
msg.validate()
.map_err(|e| LlmError::Parse { detail: e.detail })?;
}
let caps = self.codec.capabilities_for(&req.model);
validate_capabilities(req, &caps)?;
Ok(())
}
fn build_request_builder(
&self,
req: &CodecRequest,
) -> Result<reqwest::RequestBuilder, LlmError> {
let url = self
.config
.base_url
.join(&req.path as &str)
.map_err(|e| LlmError::Network {
detail: format!("Invalid URL: {}", e),
})?;
let builder = self.client.post(url);
let builder = self.config.auth.apply(builder);
let mut merged = self.codec.default_headers().clone();
merged.extend(self.extra_headers.clone());
merged.extend(req.headers.clone());
let builder = merged
.iter()
.fold(builder, |b, (key, value)| b.header(key, value));
Ok(builder.body(req.body.clone()))
}
async fn send(&self, req: CodecRequest) -> Result<RawResponse, LlmError> {
let builder = self.build_request_builder(&req)?;
let resp = builder.send().await.map_err(|e| LlmError::Network {
detail: e.to_string(),
})?;
let status = resp.status().as_u16();
let headers = resp.headers().clone();
let body = resp.bytes().await.map_err(|e| LlmError::Network {
detail: e.to_string(),
})?;
Ok(RawResponse {
status,
headers,
body,
})
}
fn handle_error(&self, resp: &RawResponse) -> LlmError {
let body_str = String::from_utf8_lossy(&resp.body);
match resp.status {
401 => LlmError::Provider {
provider: self.codec.provider_id().to_string(),
status: Some(401),
code: None,
message: body_str.into_owned(),
},
429 => LlmError::Provider {
provider: self.codec.provider_id().to_string(),
status: Some(429),
code: None,
message: "rate limited".into(),
},
status @ (400..=599) => LlmError::Provider {
provider: self.codec.provider_id().to_string(),
status: Some(status),
code: None,
message: body_str.into_owned(),
},
_ => LlmError::Provider {
provider: self.codec.provider_id().to_string(),
status: Some(resp.status),
code: None,
message: format!("Unexpected status: {}", resp.status),
},
}
}
}
#[async_trait]
#[allow(private_bounds)]
impl<C: ProviderExtension + 'static> LlmProvider for CodecProvider<C> {
async fn call(&self, request: &ChatRequest) -> Result<ChatResponse, LlmError> {
self.validate_request(request)?;
let http_req = self.codec.encode(request, false)?;
let resp = match tokio::time::timeout(self.config.timeout, self.send(http_req)).await {
Ok(result) => result?,
Err(_elapsed) => {
return Err(LlmError::Timeout {
detail: format!("request timed out after {:?}", self.config.timeout),
});
}
};
if (200..=299).contains(&resp.status) {
self.codec.decode(&resp.body)
} else {
Err(self.handle_error(&resp))
}
}
async fn stream(&self, request: &ChatRequest) -> Result<ProviderStream, LlmError> {
self.validate_request(request)?;
let http_req = self.codec.encode(request, true)?;
let resp = self
.build_request_builder(&http_req)?
.send()
.await
.map_err(|e| LlmError::Network {
detail: e.to_string(),
})?;
let status = resp.status().as_u16();
if status >= 400 {
let body = resp.bytes().await.map_err(|e| LlmError::Network {
detail: e.to_string(),
})?;
let raw_resp = RawResponse {
status,
headers: HeaderMap::new(),
body,
};
return Err(self.handle_error(&raw_resp));
}
let byte_stream = resp.bytes_stream();
let mut first_token_guard = byte_stream;
let first_chunk = match tokio::time::timeout(
self.config.timeout,
tokio_stream::StreamExt::next(&mut first_token_guard),
)
.await
{
Ok(Some(result)) => result,
Ok(None) => return Err(LlmError::UnexpectedEof),
Err(_elapsed) => {
return Err(LlmError::Timeout {
detail: format!("first token timed out after {:?}", self.config.timeout),
});
}
};
let first_bytes = first_chunk.map_err(|e| LlmError::Network {
detail: e.to_string(),
})?;
let model = request.model.clone();
let codec = Arc::clone(&self.codec);
let idle_timeout = self.config.idle_timeout;
let remaining = first_token_guard
.map(move |item| {
item.map_err(|e| LlmError::Network {
detail: e.to_string(),
})
})
.timeout(idle_timeout)
.map(move |item| match item {
Ok(Ok(bytes)) => Ok(bytes),
Ok(Err(e)) => Err(e),
Err(_elapsed) => {
tracing::error!(?idle_timeout, "stream idle timeout triggered");
Err(LlmError::Timeout {
detail: format!("stream idle timed out after {:?}", idle_timeout),
})
}
});
let byte_stream =
futures_util::stream::once(async move { Ok(first_bytes) }).chain(remaining);
let boxed_stream = Box::pin(byte_stream);
let (tx, rx) = tokio::sync::mpsc::channel(32);
let dropped = Arc::new(AtomicU64::new(0));
let dropped_clone = Arc::clone(&dropped);
let mut sink = ChannelSink { tx, dropped };
tokio::spawn(async move {
super::stream::stream_processor::process_stream(
&mut sink,
&*codec,
model,
boxed_stream,
)
.await;
let n = dropped_clone.load(Ordering::Relaxed);
if n > 0 {
tracing::warn!(
dropped_events = n,
"non-critical stream events were dropped due to full channel"
);
}
});
let rx_stream = tokio_stream::wrappers::ReceiverStream::new(rx);
Ok(Box::pin(rx_stream))
}
fn provider_id(&self) -> &str {
self.codec.provider_id()
}
fn capabilities_for(&self, model: &str) -> Capabilities {
self.codec.capabilities_for(model)
}
}
#[derive(Clone, Debug)]
pub(crate) struct ProviderConfig {
base_url: url::Url,
auth: AuthConfig,
connect_timeout: std::time::Duration,
timeout: std::time::Duration,
idle_timeout: std::time::Duration,
}
impl ProviderConfig {
pub(crate) fn bearer(
base_url: impl AsRef<str>,
api_key: impl Into<String>,
) -> Result<Self, url::ParseError> {
Ok(Self {
base_url: url::Url::parse(base_url.as_ref())?,
auth: AuthConfig::Bearer {
api_key: secrecy::SecretString::new(api_key.into()),
},
connect_timeout: std::time::Duration::from_secs(10),
timeout: std::time::Duration::from_secs(120),
idle_timeout: std::time::Duration::from_secs(30),
})
}
pub(crate) fn header(
base_url: impl AsRef<str>,
header: impl Into<String>,
value: impl Into<String>,
) -> Result<Self, url::ParseError> {
Ok(Self {
base_url: url::Url::parse(base_url.as_ref())?,
auth: AuthConfig::Header {
header: header.into(),
value: secrecy::SecretString::new(value.into()),
},
connect_timeout: std::time::Duration::from_secs(10),
timeout: std::time::Duration::from_secs(120),
idle_timeout: std::time::Duration::from_secs(30),
})
}
pub(crate) fn none(base_url: impl AsRef<str>) -> Result<Self, url::ParseError> {
Ok(Self {
base_url: url::Url::parse(base_url.as_ref())?,
auth: AuthConfig::None,
connect_timeout: std::time::Duration::from_secs(10),
timeout: std::time::Duration::from_secs(120),
idle_timeout: std::time::Duration::from_secs(30),
})
}
pub(crate) fn load(meta: &dyn ProviderMeta) -> Result<Self, ProviderBuildError> {
let provider_id = meta.provider_id();
let env_prefix = provider_id.to_ascii_uppercase();
let default_url = meta.default_base_url();
let auth_style = meta.auth_style();
let api_key_env = meta.api_key_env();
let base_url = std::env::var(format!("{}_BASE_URL", env_prefix)).unwrap_or_else(|_| {
tracing::debug!(
provider = provider_id,
url = default_url,
"{}_BASE_URL not set, using default",
env_prefix
);
default_url.to_string()
});
let api_key_name = api_key_env.into_owned();
let api_key = std::env::var(&api_key_name).map_err(|_| {
ProviderBuildError::Env(ProviderEnvError::MissingEnv {
name: api_key_name.clone(),
})
})?;
if api_key.is_empty() {
return Err(ProviderBuildError::Env(ProviderEnvError::EmptyEnv {
name: api_key_name,
}));
}
match auth_style {
super::codec::AuthStyle::Bearer => Self::bearer(&base_url, api_key).map_err(Into::into),
super::codec::AuthStyle::CustomHeader(header) => {
Self::header(&base_url, header, api_key).map_err(Into::into)
}
super::codec::AuthStyle::None => Self::none(&base_url).map_err(Into::into),
}
}
}
#[derive(Clone, Debug)]
pub(crate) enum AuthConfig {
Bearer {
api_key: secrecy::SecretString,
},
Header {
header: String,
value: secrecy::SecretString,
},
None,
}
impl AuthConfig {
pub(crate) fn apply(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
match self {
AuthConfig::Bearer { api_key } => builder.bearer_auth(api_key.expose_secret()),
AuthConfig::Header { header, value } => builder.header(header, value.expose_secret()),
AuthConfig::None => builder,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProviderProfile {
OpenRouter,
Anthropic,
Groq,
DeepSeek,
SGLang,
Ollama,
MiMo,
Zhipu,
DashScope,
}
impl ProviderProfile {
fn base_url(self) -> &'static str {
match self {
Self::OpenRouter => "https://openrouter.ai/api/v1",
Self::Anthropic => "https://api.anthropic.com/v1",
Self::Groq => "https://api.groq.com/openai/v1",
Self::DeepSeek => "https://api.deepseek.com/v1",
Self::SGLang => "http://localhost:30000/v1",
Self::Ollama => "http://localhost:11434/v1",
Self::MiMo => "https://api.xiaomimimo.com/v1",
Self::Zhipu => "https://open.bigmodel.cn/api/paas/v4",
Self::DashScope => "https://dashscope.aliyuncs.com/compatible-mode/v1",
}
}
fn env_prefix(self) -> &'static str {
match self {
Self::OpenRouter => "OPENROUTER",
Self::Anthropic => "ANTHROPIC",
Self::Groq => "GROQ",
Self::DeepSeek => "DEEPSEEK",
Self::SGLang => "SGLANG",
Self::Ollama => "OLLAMA",
Self::MiMo => "MIMO",
Self::Zhipu => "ZAI",
Self::DashScope => "DASHSCOPE",
}
}
}
pub struct ProviderBuilder<C> {
codec: C,
base_url: String,
auth: Option<AuthConfig>,
connect_timeout: std::time::Duration,
timeout: std::time::Duration,
idle_timeout: std::time::Duration,
extra_headers: HeaderMap,
error: Option<ProviderBuildError>,
}
impl<C> ProviderBuilder<C> {
pub(crate) fn new(codec: C) -> Self {
Self {
codec,
base_url: String::new(),
auth: None,
connect_timeout: std::time::Duration::from_secs(10),
timeout: std::time::Duration::from_secs(120),
idle_timeout: std::time::Duration::from_secs(30),
extra_headers: HeaderMap::new(),
error: None,
}
}
pub fn profile(mut self, profile: ProviderProfile) -> Self {
if self.base_url.is_empty() {
self.base_url = profile.base_url().to_string();
}
self
}
pub fn base_url(mut self, url: impl AsRef<str>) -> Self {
self.base_url = url.as_ref().to_string();
self
}
pub fn api_key(mut self, key: impl Into<String>) -> Self {
self.auth = Some(AuthConfig::Bearer {
api_key: secrecy::SecretString::new(key.into()),
});
self
}
pub fn auth_header(mut self, header: impl Into<String>, value: impl Into<String>) -> Self {
self.auth = Some(AuthConfig::Header {
header: header.into(),
value: secrecy::SecretString::new(value.into()),
});
self
}
pub(crate) fn auth(mut self, auth: AuthConfig) -> Self {
self.auth = Some(auth);
self
}
pub fn header(mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> Self {
if self.error.is_some() {
return self;
}
let name: http::HeaderName = match key.as_ref().parse() {
Ok(v) => v,
Err(_) => {
self.error = Some(ProviderBuildError::InvalidHeader {
field: "name".to_string(),
value: key.as_ref().to_string(),
});
return self;
}
};
let val: http::HeaderValue = match value.as_ref().parse() {
Ok(v) => v,
Err(_) => {
self.error = Some(ProviderBuildError::InvalidHeader {
field: "value".to_string(),
value: value.as_ref().to_string(),
});
return self;
}
};
self.extra_headers.insert(name, val);
self
}
pub fn extra_headers(mut self, headers: impl IntoIterator<Item = (String, String)>) -> Self {
if self.error.is_some() {
return self;
}
for (key, value) in headers {
let name: http::HeaderName = match key.parse() {
Ok(v) => v,
Err(_) => {
self.error = Some(ProviderBuildError::InvalidHeader {
field: "name".to_string(),
value: key,
});
return self;
}
};
let val: http::HeaderValue = match value.parse() {
Ok(v) => v,
Err(_) => {
self.error = Some(ProviderBuildError::InvalidHeader {
field: "value".to_string(),
value,
});
return self;
}
};
self.extra_headers.insert(name, val);
}
self
}
pub fn connect_timeout(mut self, timeout: std::time::Duration) -> Self {
self.connect_timeout = timeout;
self
}
pub fn timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
pub fn idle_timeout(mut self, timeout: std::time::Duration) -> Self {
self.idle_timeout = timeout;
self
}
pub fn build(self) -> Result<CodecProvider<C>, ProviderBuildError>
where
C: ProviderExtension,
{
if let Some(e) = self.error {
return Err(e);
}
let base_url = url::Url::parse(&self.base_url)?;
let auth = self.auth.unwrap_or(AuthConfig::None);
let network_idle_guard = std::cmp::max(
self.idle_timeout.saturating_mul(3),
std::time::Duration::from_secs(120),
);
let client = reqwest::Client::builder()
.connect_timeout(self.connect_timeout)
.read_timeout(network_idle_guard)
.user_agent(format!("LeLLM/{}", env!("CARGO_PKG_VERSION")))
.build()
.unwrap_or_default();
Ok(CodecProvider {
codec: Arc::new(self.codec),
config: ProviderConfig {
base_url,
auth,
connect_timeout: self.connect_timeout,
timeout: self.timeout,
idle_timeout: self.idle_timeout,
},
client,
extra_headers: self.extra_headers,
})
}
}
impl<C> CodecProvider<C>
where
C: ProviderExtension,
{
pub fn openrouter(codec: C) -> Result<Self, ProviderBuildError> {
let profile = ProviderProfile::OpenRouter;
let base_url =
std::env::var(format!("{}_BASE_URL", profile.env_prefix())).unwrap_or_else(|_| {
tracing::debug!("{}_BASE_URL not set, using default", profile.env_prefix());
profile.base_url().to_string()
});
let api_key_env = format!("{}_API_KEY", profile.env_prefix());
let api_key = std::env::var(&api_key_env).map_err(|_| {
ProviderBuildError::Env(ProviderEnvError::MissingEnv {
name: api_key_env.clone(),
})
})?;
if api_key.is_empty() {
return Err(ProviderBuildError::Env(ProviderEnvError::EmptyEnv {
name: api_key_env,
}));
}
Self::builder(codec)
.base_url(&base_url)
.api_key(api_key)
.build()
}
}