use anyhow::Result;
use oxi_sdk::{Oxi, OxiBuilder, ProviderPool, RateLimitPolicy};
use std::sync::Arc;
use crate::credential::CredentialStore;
pub struct OxiosEngine {
oxi: Oxi,
default_model_id: String,
routing_control: Option<oxi_sdk::RoutingControl>,
pools: parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
authorizer: Option<Arc<oxi_sdk::Authorizer>>,
tracer: Option<Arc<oxi_sdk::Tracer>>,
cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
}
impl OxiosEngine {
pub fn new(default_model_id: impl Into<String>) -> Self {
let model_id = default_model_id.into();
let oxi = OxiBuilder::new().with_builtins().build();
Self {
oxi,
default_model_id: model_id,
routing_control: None,
pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
authorizer: None,
tracer: None,
cost_tracker: None,
}
}
pub fn from_config(default_model_id: impl Into<String>, config_api_key: Option<&str>) -> Self {
let model_id = default_model_id.into();
let primary_provider = model_id
.split_once('/')
.map(|(p, _)| p)
.unwrap_or("anthropic");
let mut builder = OxiBuilder::new().with_builtins();
let providers = ["anthropic", "openai", "google", "deepseek", "xai"];
for provider in providers {
let config_key = if provider == primary_provider {
config_api_key
} else {
None
};
if let Some((key, source)) = CredentialStore::resolve(provider, config_key) {
tracing::debug!(
provider,
source = ?source,
"Injected credential into engine"
);
builder = builder.api_key(provider, key);
}
}
let oxi = builder.build();
Self {
oxi,
default_model_id: model_id,
routing_control: None,
pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
authorizer: None,
tracer: None,
cost_tracker: None,
}
}
pub fn builder() -> OxiosEngineBuilder {
OxiosEngineBuilder {
inner: OxiBuilder::new().with_builtins(),
default_model_id: "anthropic/claude-sonnet-4-20250514".to_string(),
authorizer: None,
tracer: None,
cost_tracker: None,
}
}
pub fn oxi(&self) -> &Oxi {
&self.oxi
}
pub fn authorizer(&self) -> Option<&Arc<oxi_sdk::Authorizer>> {
self.authorizer.as_ref()
}
pub fn tracer(&self) -> Option<&Arc<oxi_sdk::Tracer>> {
self.tracer.as_ref()
}
pub fn cost_tracker(&self) -> Option<&Arc<oxi_sdk::CostTracker>> {
self.cost_tracker.as_ref()
}
pub fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
self.oxi.resolve_model(model_id)
}
pub fn create_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
self.oxi.create_provider(name)
}
pub fn default_model_id(&self) -> &str {
&self.default_model_id
}
pub fn routing_control(&self) -> Option<&oxi_sdk::RoutingControl> {
self.routing_control.as_ref()
}
pub fn pooled_provider(&self, name: &str, rpm: u32) -> Result<Arc<dyn oxi_sdk::Provider>> {
{
let pools = self.pools.read();
if let Some(pooled) = pools.get(name) {
return Ok(pooled.clone());
}
}
let base = self.create_provider(name)?;
let policy = RateLimitPolicy::rpm(rpm);
let pool = ProviderPool::new(base, policy, name);
let pooled: Arc<dyn oxi_sdk::Provider> = Arc::new(pool);
{
let mut pools = self.pools.write();
pools.insert(name.to_string(), pooled.clone());
}
tracing::info!(provider = name, rpm, "Created provider pool");
Ok(pooled)
}
}
pub struct OxiosEngineBuilder {
inner: OxiBuilder,
default_model_id: String,
authorizer: Option<Arc<oxi_sdk::Authorizer>>,
tracer: Option<Arc<oxi_sdk::Tracer>>,
cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
}
impl OxiosEngineBuilder {
pub fn default_model(mut self, model_id: impl Into<String>) -> Self {
self.default_model_id = model_id.into();
self
}
pub fn api_key(self, provider: &str, key: impl Into<String>) -> Self {
Self {
inner: self.inner.api_key(provider, key),
default_model_id: self.default_model_id,
authorizer: self.authorizer,
tracer: self.tracer,
cost_tracker: self.cost_tracker,
}
}
pub fn credential(
self,
provider: &str,
api_key: impl Into<String>,
base_url: Option<&str>,
) -> Self {
Self {
inner: self.inner.credential(provider, api_key, base_url),
default_model_id: self.default_model_id,
authorizer: self.authorizer,
tracer: self.tracer,
cost_tracker: self.cost_tracker,
}
}
pub fn provider(self, name: &str, p: impl oxi_sdk::Provider + 'static) -> Self {
Self {
inner: self.inner.provider(name, p),
default_model_id: self.default_model_id,
authorizer: self.authorizer,
tracer: self.tracer,
cost_tracker: self.cost_tracker,
}
}
pub fn build(self) -> OxiosEngine {
OxiosEngine {
oxi: self.inner.build(),
default_model_id: self.default_model_id,
routing_control: None,
pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
authorizer: self.authorizer,
tracer: self.tracer,
cost_tracker: self.cost_tracker,
}
}
pub fn build_with_routing(self) -> (OxiosEngine, oxi_sdk::RoutingControl) {
use oxi_sdk::RoutingControl;
let routing_config = oxi_sdk::routing::RoutingConfig::default();
let routing_control = RoutingControl::new(routing_config);
let engine = OxiosEngine {
oxi: self.inner.build(),
default_model_id: self.default_model_id,
routing_control: Some(routing_control.clone()),
pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
authorizer: self.authorizer,
tracer: self.tracer,
cost_tracker: self.cost_tracker,
};
(engine, routing_control)
}
pub fn with_authorizer(mut self, authorizer: Arc<oxi_sdk::Authorizer>) -> Self {
self.authorizer = Some(authorizer);
self
}
pub fn with_tracer(mut self, tracer: Arc<oxi_sdk::Tracer>) -> Self {
self.tracer = Some(tracer);
self
}
pub fn with_cost_tracker(mut self, cost_tracker: Arc<oxi_sdk::CostTracker>) -> Self {
self.cost_tracker = Some(cost_tracker);
self
}
}
pub trait EngineProvider: Send + Sync {
fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>>;
fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model>;
fn default_model_id(&self) -> &str;
}
impl EngineProvider for OxiosEngine {
fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
self.create_provider(provider_name)
}
fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
self.resolve_model(model_id)
}
fn default_model_id(&self) -> &str {
&self.default_model_id
}
}
impl std::fmt::Debug for OxiosEngine {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OxiosEngine")
.field("default_model_id", &self.default_model_id)
.field("routing_enabled", &self.routing_control.is_some())
.finish()
}
}
pub struct EngineHandle {
inner: parking_lot::RwLock<Arc<OxiosEngine>>,
}
impl EngineHandle {
pub fn new(engine: Arc<OxiosEngine>) -> Self {
Self {
inner: parking_lot::RwLock::new(engine),
}
}
pub fn get(&self) -> Arc<OxiosEngine> {
Arc::clone(&self.inner.read())
}
pub fn swap(&self, new_engine: OxiosEngine) {
let mut guard = self.inner.write();
let old_id = guard.default_model_id().to_string();
*guard = Arc::new(new_engine);
tracing::info!(
old_model = %old_id,
new_model = %guard.default_model_id(),
"Engine hot-swapped"
);
}
}
impl std::fmt::Debug for EngineHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let engine = self.inner.read();
f.debug_struct("EngineHandle")
.field("current_model", &engine.default_model_id())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolve_model_with_provider_prefix() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let model = engine.resolve_model("openai/gpt-4o").unwrap();
assert_eq!(model.provider, "openai");
assert_eq!(model.id, "gpt-4o");
}
#[test]
fn test_resolve_model_without_provider_prefix() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
assert_eq!(model.provider, "anthropic");
}
#[test]
fn test_default_model_id() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
assert_eq!(
engine.default_model_id(),
"anthropic/claude-sonnet-4-20250514"
);
}
#[test]
fn test_resolve_model_not_found() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let result = engine.resolve_model("nonexistent/model-xyz");
assert!(result.is_err());
}
#[test]
fn test_create_provider_anthropic() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let provider = engine.create_provider("anthropic");
assert!(provider.is_ok());
}
#[test]
fn test_create_provider_not_found() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let result = engine.create_provider("nonexistent_provider");
assert!(result.is_err());
}
#[test]
fn test_builder_with_credential() {
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.credential("openai", "sk-test", None)
.build();
assert_eq!(engine.default_model_id(), "openai/gpt-4o");
}
#[test]
fn test_engine_provider_trait_on_engine() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let provider: &dyn EngineProvider = &engine;
assert!(provider.create_provider("anthropic").is_ok());
assert!(provider.resolve_model("openai/gpt-4o").is_ok());
}
#[test]
fn test_engine_handle_get_returns_current() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let handle = EngineHandle::new(Arc::new(engine));
let e = handle.get();
assert_eq!(e.default_model_id(), "anthropic/claude-sonnet-4-20250514");
}
#[test]
fn test_engine_handle_swap_updates() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let handle = EngineHandle::new(Arc::new(engine));
let new_engine = OxiosEngine::new("openai/gpt-4o");
handle.swap(new_engine);
let e = handle.get();
assert_eq!(e.default_model_id(), "openai/gpt-4o");
}
#[test]
fn test_engine_handle_swap_preserves_old_arc() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
let handle = EngineHandle::new(Arc::new(engine));
let old = handle.get();
assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
handle.swap(OxiosEngine::new("openai/gpt-4o"));
assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
let current = handle.get();
assert_eq!(current.default_model_id(), "openai/gpt-4o");
}
#[test]
fn test_rfc014_phase_d_default_fields_are_none() {
let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
assert!(engine.authorizer().is_none());
assert!(engine.tracer().is_none());
assert!(engine.cost_tracker().is_none());
let engine = OxiosEngine::from_config("anthropic/claude-sonnet-4-20250514", None);
assert!(engine.authorizer().is_none());
assert!(engine.tracer().is_none());
assert!(engine.cost_tracker().is_none());
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.build();
assert!(engine.authorizer().is_none());
assert!(engine.tracer().is_none());
assert!(engine.cost_tracker().is_none());
let (engine, _rc) = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.build_with_routing();
assert!(engine.authorizer().is_none());
assert!(engine.tracer().is_none());
assert!(engine.cost_tracker().is_none());
}
#[test]
fn test_rfc014_phase_d_with_tracer() {
let tracer = Arc::new(oxi_sdk::Tracer::new());
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.with_tracer(tracer.clone())
.build();
assert!(engine.tracer().is_some());
assert!(engine.authorizer().is_none());
assert!(engine.cost_tracker().is_none());
}
#[test]
fn test_rfc014_phase_d_with_cost_tracker() {
let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
let model_registry = oxi_for_registry.models_arc();
let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
model_registry,
oxi_sdk::CostTrackerConfig::default(),
));
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.with_cost_tracker(cost_tracker)
.build();
assert!(engine.cost_tracker().is_some());
assert!(engine.authorizer().is_none());
assert!(engine.tracer().is_none());
}
#[test]
fn test_rfc014_phase_d_with_authorizer() {
let audit = Arc::new(oxi_sdk::AuditLog::new(16));
let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.with_authorizer(authorizer)
.build();
assert!(engine.authorizer().is_some());
assert!(engine.tracer().is_none());
assert!(engine.cost_tracker().is_none());
}
#[test]
fn test_rfc014_phase_d_all_three_handles() {
let audit = Arc::new(oxi_sdk::AuditLog::new(16));
let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
let tracer = Arc::new(oxi_sdk::Tracer::new());
let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
let model_registry = oxi_for_registry.models_arc();
let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
model_registry,
oxi_sdk::CostTrackerConfig::default(),
));
let engine = OxiosEngine::builder()
.default_model("openai/gpt-4o")
.api_key("openai", "sk-test")
.with_authorizer(authorizer)
.with_tracer(tracer)
.with_cost_tracker(cost_tracker)
.build();
assert!(engine.authorizer().is_some());
assert!(engine.tracer().is_some());
assert!(engine.cost_tracker().is_some());
assert_eq!(engine.default_model_id(), "openai/gpt-4o");
}
}