use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::typed_id::ProviderId;
#[cfg(feature = "openapi")]
use utoipa::ToSchema;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DriverId(std::borrow::Cow<'static, str>);
impl std::fmt::Display for DriverId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
impl DriverId {
#[allow(non_upper_case_globals)]
pub const OpenAI: Self = Self(std::borrow::Cow::Borrowed("openai"));
#[allow(non_upper_case_globals)]
pub const OpenRouter: Self = Self(std::borrow::Cow::Borrowed("openrouter"));
#[allow(non_upper_case_globals)]
pub const AzureOpenAI: Self = Self(std::borrow::Cow::Borrowed("azure_openai"));
#[allow(non_upper_case_globals)]
pub const OpenAICompletions: Self = Self(std::borrow::Cow::Borrowed("openai_completions"));
#[allow(non_upper_case_globals)]
pub const Anthropic: Self = Self(std::borrow::Cow::Borrowed("anthropic"));
#[allow(non_upper_case_globals)]
pub const Gemini: Self = Self(std::borrow::Cow::Borrowed("gemini"));
#[allow(non_upper_case_globals)]
pub const LlmSim: Self = Self(std::borrow::Cow::Borrowed("llmsim"));
#[allow(non_upper_case_globals)]
pub const Bedrock: Self = Self(std::borrow::Cow::Borrowed("bedrock"));
#[allow(non_upper_case_globals)]
pub const Mai: Self = Self(std::borrow::Cow::Borrowed("mai"));
#[allow(non_upper_case_globals)]
pub const Fireworks: Self = Self(std::borrow::Cow::Borrowed("fireworks"));
#[allow(non_upper_case_globals)]
pub const Meta: Self = Self(std::borrow::Cow::Borrowed("meta"));
pub fn external(id: impl AsRef<str>) -> Self {
Self(std::borrow::Cow::Owned(
id.as_ref().trim().to_ascii_lowercase(),
))
}
pub fn as_str(&self) -> &str {
self.0.as_ref()
}
pub fn default_trace_templates(&self) -> (Option<String>, Option<String>) {
if self == &DriverId::OpenRouter {
(
Some("https://openrouter.ai/logs?id={response_id}".to_string()),
Some("https://openrouter.ai/logs".to_string()),
)
} else {
(None, None)
}
}
}
impl std::str::FromStr for DriverId {
type Err = std::convert::Infallible;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(DriverId::external(s))
}
}
impl Serialize for DriverId {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for DriverId {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
let s = String::deserialize(d)?;
if s.trim().is_empty() {
return Err(serde::de::Error::custom("provider type cannot be empty"));
}
if s != s.trim() {
return Err(serde::de::Error::custom(
"provider type cannot have leading or trailing whitespace",
));
}
Ok(s.parse().unwrap_or_else(|_| unreachable!()))
}
}
#[cfg(test)]
mod tests {
use super::DriverId;
#[test]
fn built_in_driver_ids_share_literal_wire_and_display_contracts() {
for (driver, wire) in [
(DriverId::OpenAI, "openai"),
(DriverId::OpenRouter, "openrouter"),
(DriverId::AzureOpenAI, "azure_openai"),
(DriverId::OpenAICompletions, "openai_completions"),
(DriverId::Anthropic, "anthropic"),
(DriverId::Gemini, "gemini"),
(DriverId::LlmSim, "llmsim"),
(DriverId::Bedrock, "bedrock"),
(DriverId::Mai, "mai"),
(DriverId::Fireworks, "fireworks"),
(DriverId::Meta, "meta"),
] {
assert_eq!(driver.as_str(), wire);
assert_eq!(driver.to_string(), wire);
assert_eq!(wire.parse::<DriverId>().unwrap(), driver);
assert_eq!(
serde_json::to_value(&driver).unwrap(),
serde_json::json!(wire)
);
assert_eq!(
serde_json::from_value::<DriverId>(serde_json::json!(wire)).unwrap(),
driver
);
}
}
#[test]
fn wire_ids_reject_empty_and_padded_values() {
for value in ["", " ", "\t\n"] {
let error = serde_json::from_value::<DriverId>(serde_json::json!(value)).unwrap_err();
assert!(error.to_string().contains("provider type cannot be empty"));
}
let error = serde_json::from_value::<DriverId>(serde_json::json!(" openai ")).unwrap_err();
assert!(error.to_string().contains("leading or trailing whitespace"));
}
#[test]
fn wire_ids_normalize_case_and_accept_extensions() {
assert_eq!("OpenAI".parse::<DriverId>().unwrap(), DriverId::OpenAI);
assert_eq!(DriverId::external("OpenAI-Codex").as_str(), "openai-codex");
let driver = serde_json::from_str::<DriverId>(r#""Custom-Driver""#).unwrap();
assert_eq!(driver.as_str(), "custom-driver");
assert_eq!(
serde_json::to_string(&driver).unwrap(),
r#""custom-driver""#
);
assert_eq!(" Custom-Driver ".parse::<DriverId>().unwrap(), driver);
let mut ids = std::collections::HashSet::new();
ids.insert(driver);
ids.insert(DriverId::external("CUSTOM-DRIVER"));
assert_eq!(ids.len(), 1, "normalization must preserve hash identity");
}
}
#[cfg(feature = "openapi")]
impl utoipa::ToSchema for DriverId {
fn name() -> std::borrow::Cow<'static, str> {
std::borrow::Cow::Borrowed("DriverId")
}
}
#[cfg(feature = "openapi")]
impl utoipa::PartialSchema for DriverId {
fn schema() -> utoipa::openapi::RefOr<utoipa::openapi::Schema> {
utoipa::openapi::ObjectBuilder::new()
.schema_type(utoipa::openapi::schema::SchemaType::new(
utoipa::openapi::schema::Type::String,
))
.description(Some(
"LLM provider type. Built-in: openai, openrouter, azure_openai, \
openai_completions, anthropic, gemini, llmsim, bedrock, mai, fireworks, meta. \
Any other string is treated as an embedder-defined external provider.",
))
.build()
.into()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum ProviderStatus {
Active,
Disabled,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct ProviderTraceConfig {
pub enabled: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub generation_url_template: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_url_template: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct ProviderRequestHeader {
pub name: String,
pub value: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct ProviderRequestOptions {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<ProviderRequestHeader>,
#[serde(default)]
pub cache_diagnostics: bool,
}
impl ProviderRequestOptions {
pub fn is_empty(&self) -> bool {
self.headers.is_empty() && !self.cache_diagnostics
}
pub fn header_pairs(&self) -> Vec<(String, String)> {
self.headers
.iter()
.map(|header| (header.name.clone(), header.value.clone()))
.collect()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct Provider {
#[cfg_attr(feature = "openapi", schema(value_type = String, example = "provider_01933b5a00007000800000000000001"))]
pub id: ProviderId,
pub name: String,
pub provider_type: DriverId,
#[serde(skip_serializing_if = "Option::is_none")]
pub base_url: Option<String>,
pub api_key_set: bool,
pub status: ProviderStatus,
pub managed: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_synced_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
#[serde(skip_serializing_if = "Option::is_none")]
pub trace: Option<ProviderTraceConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub request_options: Option<ProviderRequestOptions>,
}