1use anyhow::Result;
19use oxi_sdk::{
20 CatalogConfig, FileModelCatalog, ModelCatalog, Oxi, OxiBuilder, ProviderPool, RateLimitPolicy,
21};
22use std::sync::Arc;
23
24use crate::credential::{CredentialStore, discover_auth_store_providers};
25
26pub struct OxiosEngine {
40 oxi: Oxi,
41 default_model_id: String,
42 routing_control: Option<oxi_sdk::RoutingControl>,
44 pools: parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
47 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
52 tracer: Option<Arc<oxi_sdk::Tracer>>,
53 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
54}
55
56impl OxiosEngine {
57 pub fn new(default_model_id: impl Into<String>) -> Self {
62 let model_id = default_model_id.into();
63 let oxi = OxiBuilder::new().with_builtins().build();
64 Self {
65 oxi,
66 default_model_id: model_id,
67 routing_control: None,
68 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
69 authorizer: None,
71 tracer: None,
72 cost_tracker: None,
73 }
74 }
75
76 pub fn from_config(default_model_id: impl Into<String>, config_api_key: Option<&str>) -> Self {
88 Self::from_config_with_catalog_opt(default_model_id, config_api_key, None)
89 }
90
91 pub fn from_config_with_catalog(
100 default_model_id: impl Into<String>,
101 config_api_key: Option<&str>,
102 catalog: Arc<dyn ModelCatalog>,
103 ) -> Self {
104 Self::from_config_with_catalog_opt(default_model_id, config_api_key, Some(catalog))
105 }
106
107 fn from_config_with_catalog_opt(
108 default_model_id: impl Into<String>,
109 config_api_key: Option<&str>,
110 catalog: Option<Arc<dyn ModelCatalog>>,
111 ) -> Self {
112 let model_id = default_model_id.into();
113
114 let primary_provider = model_id
116 .split_once('/')
117 .map(|(p, _)| p)
118 .unwrap_or("anthropic");
119
120 let mut builder = OxiBuilder::new().with_builtins();
121
122 let mut providers_to_try: Vec<String> = vec![
127 "anthropic".into(),
128 "openai".into(),
129 "google".into(),
130 "deepseek".into(),
131 "xai".into(),
132 "groq".into(),
133 "openrouter".into(),
134 "mistral".into(),
135 "cerebras".into(),
136 "fireworks".into(),
137 "github-copilot".into(),
138 "huggingface".into(),
139 "together".into(),
140 "minimax".into(),
141 "moonshotai".into(),
142 "kimi-coding".into(),
143 "zai".into(),
144 "opencode".into(),
145 ];
146
147 if let Ok(extra) = discover_auth_store_providers() {
150 for p in extra {
151 if !providers_to_try.contains(&p) {
152 providers_to_try.push(p);
153 }
154 }
155 }
156
157 let primary_owned = primary_provider.to_string();
159 if !providers_to_try.contains(&primary_owned) {
160 providers_to_try.push(primary_owned);
161 }
162
163 for provider in &providers_to_try {
164 let config_key = if provider == primary_provider {
167 config_api_key
168 } else {
169 None
170 };
171
172 if let Some((key, source)) = CredentialStore::resolve(provider, config_key) {
173 tracing::debug!(
174 provider,
175 source = ?source,
176 "Injected credential into engine"
177 );
178 builder = builder.api_key(provider, key);
179 }
180 }
181
182 let builder = match catalog {
183 Some(cat) => builder.with_catalog(cat),
184 None => builder,
185 };
186 let oxi = builder.build();
187 Self {
188 oxi,
189 default_model_id: model_id,
190 routing_control: None,
191 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
192 authorizer: None,
194 tracer: None,
195 cost_tracker: None,
196 }
197 }
198
199 pub fn builder() -> OxiosEngineBuilder {
228 OxiosEngineBuilder {
229 inner: OxiBuilder::new().with_builtins(),
230 default_model_id: "anthropic/claude-sonnet-4-20250514".to_string(),
231 authorizer: None,
233 tracer: None,
234 cost_tracker: None,
235 }
236 }
237
238 pub fn catalog_config() -> CatalogConfig {
246 let home = dirs::home_dir()
247 .unwrap_or_else(|| std::path::PathBuf::from("."))
248 .join(".oxios");
249 CatalogConfig {
250 cache_path: home.join("cache/models-dev.json"),
251 etag_path: home.join("cache/models-dev.json.etag"),
252 override_path: home.join("catalog/overrides.toml"),
253 snapshot_path: home.join("cache/models-dev.json"),
254 local_discovery_urls: Vec::new(),
256 ..CatalogConfig::default()
257 }
258 }
259
260 pub async fn init_file_catalog() -> Result<Arc<dyn ModelCatalog>> {
270 let catalog: Arc<dyn ModelCatalog> = FileModelCatalog::init(Self::catalog_config())
271 .await
272 .map_err(|e| anyhow::anyhow!("Failed to initialize model catalog: {e}"))?;
273 Ok(catalog)
274 }
275
276 pub fn oxi(&self) -> &Oxi {
281 &self.oxi
282 }
283
284 pub fn authorizer(&self) -> Option<&Arc<oxi_sdk::Authorizer>> {
289 self.authorizer.as_ref()
290 }
291
292 pub fn tracer(&self) -> Option<&Arc<oxi_sdk::Tracer>> {
297 self.tracer.as_ref()
298 }
299
300 pub fn cost_tracker(&self) -> Option<&Arc<oxi_sdk::CostTracker>> {
305 self.cost_tracker.as_ref()
306 }
307
308 pub fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
310 self.oxi.resolve_model(model_id)
311 }
312
313 pub fn create_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
315 self.oxi.create_provider(name)
316 }
317
318 pub fn default_model_id(&self) -> &str {
320 &self.default_model_id
321 }
322
323 pub fn routing_control(&self) -> Option<&oxi_sdk::RoutingControl> {
325 self.routing_control.as_ref()
326 }
327
328 pub fn pooled_provider(&self, name: &str, rpm: u32) -> Result<Arc<dyn oxi_sdk::Provider>> {
336 {
338 let pools = self.pools.read();
339 if let Some(pooled) = pools.get(name) {
340 return Ok(pooled.clone());
341 }
342 }
343
344 let base = self.create_provider(name)?;
346 let policy = RateLimitPolicy::rpm(rpm);
347 let pool = ProviderPool::new(base, policy, name);
348 let pooled: Arc<dyn oxi_sdk::Provider> = Arc::new(pool);
349
350 {
352 let mut pools = self.pools.write();
353 pools.insert(name.to_string(), pooled.clone());
354 }
355
356 tracing::info!(provider = name, rpm, "Created provider pool");
357 Ok(pooled)
358 }
359}
360
361pub struct OxiosEngineBuilder {
367 inner: OxiBuilder,
368 default_model_id: String,
369 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
372 tracer: Option<Arc<oxi_sdk::Tracer>>,
373 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
374}
375
376impl OxiosEngineBuilder {
377 pub fn default_model(mut self, model_id: impl Into<String>) -> Self {
379 self.default_model_id = model_id.into();
380 self
381 }
382
383 pub fn api_key(self, provider: &str, key: impl Into<String>) -> Self {
385 Self {
386 inner: self.inner.api_key(provider, key),
387 default_model_id: self.default_model_id,
388 authorizer: self.authorizer,
389 tracer: self.tracer,
390 cost_tracker: self.cost_tracker,
391 }
392 }
393
394 pub fn credential(
396 self,
397 provider: &str,
398 api_key: impl Into<String>,
399 base_url: Option<&str>,
400 ) -> Self {
401 Self {
402 inner: self.inner.credential(provider, api_key, base_url),
403 default_model_id: self.default_model_id,
404 authorizer: self.authorizer,
405 tracer: self.tracer,
406 cost_tracker: self.cost_tracker,
407 }
408 }
409
410 pub fn provider(self, name: &str, p: impl oxi_sdk::Provider + 'static) -> Self {
412 Self {
413 inner: self.inner.provider(name, p),
414 default_model_id: self.default_model_id,
415 authorizer: self.authorizer,
416 tracer: self.tracer,
417 cost_tracker: self.cost_tracker,
418 }
419 }
420
421 pub fn build(self) -> OxiosEngine {
423 OxiosEngine {
424 oxi: self.inner.build(),
425 default_model_id: self.default_model_id,
426 routing_control: None,
427 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
428 authorizer: self.authorizer,
430 tracer: self.tracer,
431 cost_tracker: self.cost_tracker,
432 }
433 }
434
435 pub fn build_with_routing(self) -> (OxiosEngine, oxi_sdk::RoutingControl) {
439 use oxi_sdk::RoutingControl;
440
441 let routing_config = oxi_sdk::routing::RoutingConfig::default();
442 let routing_control = RoutingControl::new(routing_config);
443 let engine = OxiosEngine {
444 oxi: self.inner.build(),
445 default_model_id: self.default_model_id,
446 routing_control: Some(routing_control.clone()),
447 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
448 authorizer: self.authorizer,
450 tracer: self.tracer,
451 cost_tracker: self.cost_tracker,
452 };
453 (engine, routing_control)
454 }
455
456 pub fn with_authorizer(mut self, authorizer: Arc<oxi_sdk::Authorizer>) -> Self {
468 self.authorizer = Some(authorizer);
469 self
470 }
471
472 pub fn with_tracer(mut self, tracer: Arc<oxi_sdk::Tracer>) -> Self {
475 self.tracer = Some(tracer);
476 self
477 }
478
479 pub fn with_cost_tracker(mut self, cost_tracker: Arc<oxi_sdk::CostTracker>) -> Self {
482 self.cost_tracker = Some(cost_tracker);
483 self
484 }
485
486 pub fn with_catalog(mut self, catalog: Arc<dyn oxi_sdk::ModelCatalog>) -> Self {
497 self.inner = self.inner.with_catalog(catalog);
498 self
499 }
500}
501
502pub trait EngineProvider: Send + Sync {
510 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>>;
512
513 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model>;
515
516 fn default_model_id(&self) -> &str;
518}
519
520impl EngineProvider for OxiosEngine {
521 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
522 self.create_provider(provider_name)
523 }
524
525 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
526 self.resolve_model(model_id)
527 }
528
529 fn default_model_id(&self) -> &str {
530 &self.default_model_id
531 }
532}
533
534impl std::fmt::Debug for OxiosEngine {
535 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
536 f.debug_struct("OxiosEngine")
537 .field("default_model_id", &self.default_model_id)
538 .field("routing_enabled", &self.routing_control.is_some())
539 .finish()
540 }
541}
542
543pub struct EngineHandle {
565 inner: parking_lot::RwLock<Arc<OxiosEngine>>,
566}
567
568impl EngineHandle {
569 pub fn new(engine: Arc<OxiosEngine>) -> Self {
571 Self {
572 inner: parking_lot::RwLock::new(engine),
573 }
574 }
575
576 pub fn get(&self) -> Arc<OxiosEngine> {
581 Arc::clone(&self.inner.read())
582 }
583
584 pub fn swap(&self, new_engine: OxiosEngine) {
589 let mut guard = self.inner.write();
590 let old_id = guard.default_model_id().to_string();
591 *guard = Arc::new(new_engine);
592 tracing::info!(
593 old_model = %old_id,
594 new_model = %guard.default_model_id(),
595 "Engine hot-swapped"
596 );
597 }
598}
599
600impl std::fmt::Debug for EngineHandle {
601 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
602 let engine = self.inner.read();
603 f.debug_struct("EngineHandle")
604 .field("current_model", &engine.default_model_id())
605 .finish()
606 }
607}
608
609#[cfg(test)]
614mod tests {
615 use super::*;
616
617 #[test]
618 fn test_resolve_model_with_provider_prefix() {
619 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
620 let model = engine.resolve_model("openai/gpt-4o").unwrap();
621 assert_eq!(model.provider, "openai");
622 assert_eq!(model.id, "gpt-4o");
623 }
624
625 #[test]
626 fn test_resolve_model_without_provider_prefix() {
627 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
628 let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
629 assert_eq!(model.provider, "anthropic");
630 }
631
632 #[test]
633 fn test_default_model_id() {
634 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
635 assert_eq!(
636 engine.default_model_id(),
637 "anthropic/claude-sonnet-4-20250514"
638 );
639 }
640
641 #[test]
642 fn test_resolve_model_not_found() {
643 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
644 let result = engine.resolve_model("nonexistent/model-xyz");
645 assert!(result.is_err());
646 }
647
648 #[test]
649 fn test_create_provider_anthropic() {
650 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
651 let provider = engine.create_provider("anthropic");
652 assert!(provider.is_ok());
653 }
654
655 #[test]
656 fn test_create_provider_not_found() {
657 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
658 let result = engine.create_provider("nonexistent_provider");
659 assert!(result.is_err());
660 }
661
662 #[test]
663 fn test_builder_with_credential() {
664 let engine = OxiosEngine::builder()
665 .default_model("openai/gpt-4o")
666 .credential("openai", "sk-test", None)
667 .build();
668 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
669 }
670
671 #[test]
672 fn test_engine_provider_trait_on_engine() {
673 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
674 let provider: &dyn EngineProvider = &engine;
675 assert!(provider.create_provider("anthropic").is_ok());
676 assert!(provider.resolve_model("openai/gpt-4o").is_ok());
677 }
678
679 #[test]
682 fn test_engine_handle_get_returns_current() {
683 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
684 let handle = EngineHandle::new(Arc::new(engine));
685 let e = handle.get();
686 assert_eq!(e.default_model_id(), "anthropic/claude-sonnet-4-20250514");
687 }
688
689 #[test]
690 fn test_engine_handle_swap_updates() {
691 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
692 let handle = EngineHandle::new(Arc::new(engine));
693
694 let new_engine = OxiosEngine::new("openai/gpt-4o");
695 handle.swap(new_engine);
696
697 let e = handle.get();
698 assert_eq!(e.default_model_id(), "openai/gpt-4o");
699 }
700
701 #[test]
702 fn test_engine_handle_swap_preserves_old_arc() {
703 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
705 let handle = EngineHandle::new(Arc::new(engine));
706
707 let old = handle.get();
708 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
709
710 handle.swap(OxiosEngine::new("openai/gpt-4o"));
711
712 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
714
715 let current = handle.get();
717 assert_eq!(current.default_model_id(), "openai/gpt-4o");
718 }
719
720 #[test]
723 fn test_rfc014_phase_d_default_fields_are_none() {
724 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
728 assert!(engine.authorizer().is_none());
729 assert!(engine.tracer().is_none());
730 assert!(engine.cost_tracker().is_none());
731
732 let engine = OxiosEngine::from_config("anthropic/claude-sonnet-4-20250514", None);
733 assert!(engine.authorizer().is_none());
734 assert!(engine.tracer().is_none());
735 assert!(engine.cost_tracker().is_none());
736
737 let engine = OxiosEngine::builder()
738 .default_model("openai/gpt-4o")
739 .build();
740 assert!(engine.authorizer().is_none());
741 assert!(engine.tracer().is_none());
742 assert!(engine.cost_tracker().is_none());
743
744 let (engine, _rc) = OxiosEngine::builder()
745 .default_model("openai/gpt-4o")
746 .build_with_routing();
747 assert!(engine.authorizer().is_none());
748 assert!(engine.tracer().is_none());
749 assert!(engine.cost_tracker().is_none());
750 }
751
752 #[test]
753 fn test_rfc014_phase_d_with_tracer() {
754 let tracer = Arc::new(oxi_sdk::Tracer::new());
756 let engine = OxiosEngine::builder()
757 .default_model("openai/gpt-4o")
758 .with_tracer(tracer.clone())
759 .build();
760 assert!(engine.tracer().is_some());
761 assert!(engine.authorizer().is_none());
762 assert!(engine.cost_tracker().is_none());
763 }
764
765 #[test]
766 fn test_rfc014_phase_d_with_cost_tracker() {
767 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
772 let model_registry = oxi_for_registry.models_arc();
773 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
774 model_registry,
775 oxi_sdk::CostTrackerConfig::default(),
776 ));
777 let engine = OxiosEngine::builder()
778 .default_model("openai/gpt-4o")
779 .with_cost_tracker(cost_tracker)
780 .build();
781 assert!(engine.cost_tracker().is_some());
782 assert!(engine.authorizer().is_none());
783 assert!(engine.tracer().is_none());
784 }
785
786 #[test]
787 fn test_rfc014_phase_d_with_authorizer() {
788 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
790 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
791 let engine = OxiosEngine::builder()
792 .default_model("openai/gpt-4o")
793 .with_authorizer(authorizer)
794 .build();
795 assert!(engine.authorizer().is_some());
796 assert!(engine.tracer().is_none());
797 assert!(engine.cost_tracker().is_none());
798 }
799
800 #[test]
801 fn test_rfc014_phase_d_all_three_handles() {
802 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
806 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
807 let tracer = Arc::new(oxi_sdk::Tracer::new());
808 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
809 let model_registry = oxi_for_registry.models_arc();
810 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
811 model_registry,
812 oxi_sdk::CostTrackerConfig::default(),
813 ));
814
815 let engine = OxiosEngine::builder()
816 .default_model("openai/gpt-4o")
817 .api_key("openai", "sk-test")
818 .with_authorizer(authorizer)
819 .with_tracer(tracer)
820 .with_cost_tracker(cost_tracker)
821 .build();
822
823 assert!(engine.authorizer().is_some());
824 assert!(engine.tracer().is_some());
825 assert!(engine.cost_tracker().is_some());
826 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
827 }
828
829 #[tokio::test]
835 #[ignore]
836 async fn catalog_integration_init_and_resolve() {
837 let catalog = OxiosEngine::init_file_catalog()
839 .await
840 .expect("catalog init should succeed (SNAP is always embedded)");
841
842 assert!(
845 catalog.model_count_sync() > 0,
846 "catalog should expose models from the embedded snapshot"
847 );
848 assert!(!catalog.list_providers_sync().is_empty());
849
850 let engine = OxiosEngine::builder()
852 .default_model("anthropic/claude-sonnet-4-20250514")
853 .with_catalog(catalog)
854 .build();
855 let model = engine
856 .resolve_model("openai/gpt-4o")
857 .expect("catalog-backed resolve_model should succeed");
858 assert_eq!(model.provider, "openai");
859 assert_eq!(model.id, "gpt-4o");
860 }
861}