1use anyhow::Result;
19use oxi_sdk::{Oxi, OxiBuilder, ProviderPool, RateLimitPolicy};
20use std::sync::Arc;
21
22use crate::credential::{CredentialStore, discover_auth_store_providers};
23
24pub struct OxiosEngine {
38 oxi: Oxi,
39 default_model_id: String,
40 routing_control: Option<oxi_sdk::RoutingControl>,
42 pools: parking_lot::RwLock<std::collections::HashMap<String, Arc<dyn oxi_sdk::Provider>>>,
45 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
50 tracer: Option<Arc<oxi_sdk::Tracer>>,
51 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
52}
53
54impl OxiosEngine {
55 pub fn new(default_model_id: impl Into<String>) -> Self {
60 let model_id = default_model_id.into();
61 let oxi = OxiBuilder::new().with_builtins().build();
62 Self {
63 oxi,
64 default_model_id: model_id,
65 routing_control: None,
66 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
67 authorizer: None,
69 tracer: None,
70 cost_tracker: None,
71 }
72 }
73
74 pub fn from_config(default_model_id: impl Into<String>, config_api_key: Option<&str>) -> Self {
82 let model_id = default_model_id.into();
83
84 let primary_provider = model_id
86 .split_once('/')
87 .map(|(p, _)| p)
88 .unwrap_or("anthropic");
89
90 let mut builder = OxiBuilder::new().with_builtins();
91
92 let mut providers_to_try: Vec<String> = vec![
97 "anthropic".into(),
98 "openai".into(),
99 "google".into(),
100 "deepseek".into(),
101 "xai".into(),
102 "groq".into(),
103 "openrouter".into(),
104 "mistral".into(),
105 "cerebras".into(),
106 "fireworks".into(),
107 "github-copilot".into(),
108 "huggingface".into(),
109 "together".into(),
110 "minimax".into(),
111 "moonshotai".into(),
112 "kimi-coding".into(),
113 "zai".into(),
114 "opencode".into(),
115 ];
116
117 if let Ok(extra) = discover_auth_store_providers() {
120 for p in extra {
121 if !providers_to_try.contains(&p) {
122 providers_to_try.push(p);
123 }
124 }
125 }
126
127 let primary_owned = primary_provider.to_string();
129 if !providers_to_try.contains(&primary_owned) {
130 providers_to_try.push(primary_owned);
131 }
132
133 for provider in &providers_to_try {
134 let config_key = if provider == primary_provider {
137 config_api_key
138 } else {
139 None
140 };
141
142 if let Some((key, source)) = CredentialStore::resolve(provider, config_key) {
143 tracing::debug!(
144 provider,
145 source = ?source,
146 "Injected credential into engine"
147 );
148 builder = builder.api_key(provider, key);
149 }
150 }
151
152 let oxi = builder.build();
153 Self {
154 oxi,
155 default_model_id: model_id,
156 routing_control: None,
157 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
158 authorizer: None,
160 tracer: None,
161 cost_tracker: None,
162 }
163 }
164
165 pub fn builder() -> OxiosEngineBuilder {
187 OxiosEngineBuilder {
188 inner: OxiBuilder::new().with_builtins(),
189 default_model_id: "anthropic/claude-sonnet-4-20250514".to_string(),
190 authorizer: None,
192 tracer: None,
193 cost_tracker: None,
194 }
195 }
196
197 pub fn oxi(&self) -> &Oxi {
202 &self.oxi
203 }
204
205 pub fn authorizer(&self) -> Option<&Arc<oxi_sdk::Authorizer>> {
210 self.authorizer.as_ref()
211 }
212
213 pub fn tracer(&self) -> Option<&Arc<oxi_sdk::Tracer>> {
218 self.tracer.as_ref()
219 }
220
221 pub fn cost_tracker(&self) -> Option<&Arc<oxi_sdk::CostTracker>> {
226 self.cost_tracker.as_ref()
227 }
228
229 pub fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
231 self.oxi.resolve_model(model_id)
232 }
233
234 pub fn create_provider(&self, name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
236 self.oxi.create_provider(name)
237 }
238
239 pub fn default_model_id(&self) -> &str {
241 &self.default_model_id
242 }
243
244 pub fn routing_control(&self) -> Option<&oxi_sdk::RoutingControl> {
246 self.routing_control.as_ref()
247 }
248
249 pub fn pooled_provider(&self, name: &str, rpm: u32) -> Result<Arc<dyn oxi_sdk::Provider>> {
257 {
259 let pools = self.pools.read();
260 if let Some(pooled) = pools.get(name) {
261 return Ok(pooled.clone());
262 }
263 }
264
265 let base = self.create_provider(name)?;
267 let policy = RateLimitPolicy::rpm(rpm);
268 let pool = ProviderPool::new(base, policy, name);
269 let pooled: Arc<dyn oxi_sdk::Provider> = Arc::new(pool);
270
271 {
273 let mut pools = self.pools.write();
274 pools.insert(name.to_string(), pooled.clone());
275 }
276
277 tracing::info!(provider = name, rpm, "Created provider pool");
278 Ok(pooled)
279 }
280}
281
282pub struct OxiosEngineBuilder {
288 inner: OxiBuilder,
289 default_model_id: String,
290 authorizer: Option<Arc<oxi_sdk::Authorizer>>,
293 tracer: Option<Arc<oxi_sdk::Tracer>>,
294 cost_tracker: Option<Arc<oxi_sdk::CostTracker>>,
295}
296
297impl OxiosEngineBuilder {
298 pub fn default_model(mut self, model_id: impl Into<String>) -> Self {
300 self.default_model_id = model_id.into();
301 self
302 }
303
304 pub fn api_key(self, provider: &str, key: impl Into<String>) -> Self {
306 Self {
307 inner: self.inner.api_key(provider, key),
308 default_model_id: self.default_model_id,
309 authorizer: self.authorizer,
310 tracer: self.tracer,
311 cost_tracker: self.cost_tracker,
312 }
313 }
314
315 pub fn credential(
317 self,
318 provider: &str,
319 api_key: impl Into<String>,
320 base_url: Option<&str>,
321 ) -> Self {
322 Self {
323 inner: self.inner.credential(provider, api_key, base_url),
324 default_model_id: self.default_model_id,
325 authorizer: self.authorizer,
326 tracer: self.tracer,
327 cost_tracker: self.cost_tracker,
328 }
329 }
330
331 pub fn provider(self, name: &str, p: impl oxi_sdk::Provider + 'static) -> Self {
333 Self {
334 inner: self.inner.provider(name, p),
335 default_model_id: self.default_model_id,
336 authorizer: self.authorizer,
337 tracer: self.tracer,
338 cost_tracker: self.cost_tracker,
339 }
340 }
341
342 pub fn build(self) -> OxiosEngine {
344 OxiosEngine {
345 oxi: self.inner.build(),
346 default_model_id: self.default_model_id,
347 routing_control: None,
348 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
349 authorizer: self.authorizer,
351 tracer: self.tracer,
352 cost_tracker: self.cost_tracker,
353 }
354 }
355
356 pub fn build_with_routing(self) -> (OxiosEngine, oxi_sdk::RoutingControl) {
360 use oxi_sdk::RoutingControl;
361
362 let routing_config = oxi_sdk::routing::RoutingConfig::default();
363 let routing_control = RoutingControl::new(routing_config);
364 let engine = OxiosEngine {
365 oxi: self.inner.build(),
366 default_model_id: self.default_model_id,
367 routing_control: Some(routing_control.clone()),
368 pools: parking_lot::RwLock::new(std::collections::HashMap::new()),
369 authorizer: self.authorizer,
371 tracer: self.tracer,
372 cost_tracker: self.cost_tracker,
373 };
374 (engine, routing_control)
375 }
376
377 pub fn with_authorizer(mut self, authorizer: Arc<oxi_sdk::Authorizer>) -> Self {
389 self.authorizer = Some(authorizer);
390 self
391 }
392
393 pub fn with_tracer(mut self, tracer: Arc<oxi_sdk::Tracer>) -> Self {
396 self.tracer = Some(tracer);
397 self
398 }
399
400 pub fn with_cost_tracker(mut self, cost_tracker: Arc<oxi_sdk::CostTracker>) -> Self {
403 self.cost_tracker = Some(cost_tracker);
404 self
405 }
406}
407
408pub trait EngineProvider: Send + Sync {
416 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>>;
418
419 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model>;
421
422 fn default_model_id(&self) -> &str;
424}
425
426impl EngineProvider for OxiosEngine {
427 fn create_provider(&self, provider_name: &str) -> Result<Arc<dyn oxi_sdk::Provider>> {
428 self.create_provider(provider_name)
429 }
430
431 fn resolve_model(&self, model_id: &str) -> Result<oxi_sdk::Model> {
432 self.resolve_model(model_id)
433 }
434
435 fn default_model_id(&self) -> &str {
436 &self.default_model_id
437 }
438}
439
440impl std::fmt::Debug for OxiosEngine {
441 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442 f.debug_struct("OxiosEngine")
443 .field("default_model_id", &self.default_model_id)
444 .field("routing_enabled", &self.routing_control.is_some())
445 .finish()
446 }
447}
448
449pub struct EngineHandle {
471 inner: parking_lot::RwLock<Arc<OxiosEngine>>,
472}
473
474impl EngineHandle {
475 pub fn new(engine: Arc<OxiosEngine>) -> Self {
477 Self {
478 inner: parking_lot::RwLock::new(engine),
479 }
480 }
481
482 pub fn get(&self) -> Arc<OxiosEngine> {
487 Arc::clone(&self.inner.read())
488 }
489
490 pub fn swap(&self, new_engine: OxiosEngine) {
495 let mut guard = self.inner.write();
496 let old_id = guard.default_model_id().to_string();
497 *guard = Arc::new(new_engine);
498 tracing::info!(
499 old_model = %old_id,
500 new_model = %guard.default_model_id(),
501 "Engine hot-swapped"
502 );
503 }
504}
505
506impl std::fmt::Debug for EngineHandle {
507 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
508 let engine = self.inner.read();
509 f.debug_struct("EngineHandle")
510 .field("current_model", &engine.default_model_id())
511 .finish()
512 }
513}
514
515#[cfg(test)]
520mod tests {
521 use super::*;
522
523 #[test]
524 fn test_resolve_model_with_provider_prefix() {
525 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
526 let model = engine.resolve_model("openai/gpt-4o").unwrap();
527 assert_eq!(model.provider, "openai");
528 assert_eq!(model.id, "gpt-4o");
529 }
530
531 #[test]
532 fn test_resolve_model_without_provider_prefix() {
533 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
534 let model = engine.resolve_model("claude-sonnet-4-20250514").unwrap();
535 assert_eq!(model.provider, "anthropic");
536 }
537
538 #[test]
539 fn test_default_model_id() {
540 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
541 assert_eq!(
542 engine.default_model_id(),
543 "anthropic/claude-sonnet-4-20250514"
544 );
545 }
546
547 #[test]
548 fn test_resolve_model_not_found() {
549 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
550 let result = engine.resolve_model("nonexistent/model-xyz");
551 assert!(result.is_err());
552 }
553
554 #[test]
555 fn test_create_provider_anthropic() {
556 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
557 let provider = engine.create_provider("anthropic");
558 assert!(provider.is_ok());
559 }
560
561 #[test]
562 fn test_create_provider_not_found() {
563 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
564 let result = engine.create_provider("nonexistent_provider");
565 assert!(result.is_err());
566 }
567
568 #[test]
569 fn test_builder_with_credential() {
570 let engine = OxiosEngine::builder()
571 .default_model("openai/gpt-4o")
572 .credential("openai", "sk-test", None)
573 .build();
574 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
575 }
576
577 #[test]
578 fn test_engine_provider_trait_on_engine() {
579 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
580 let provider: &dyn EngineProvider = &engine;
581 assert!(provider.create_provider("anthropic").is_ok());
582 assert!(provider.resolve_model("openai/gpt-4o").is_ok());
583 }
584
585 #[test]
588 fn test_engine_handle_get_returns_current() {
589 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
590 let handle = EngineHandle::new(Arc::new(engine));
591 let e = handle.get();
592 assert_eq!(e.default_model_id(), "anthropic/claude-sonnet-4-20250514");
593 }
594
595 #[test]
596 fn test_engine_handle_swap_updates() {
597 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
598 let handle = EngineHandle::new(Arc::new(engine));
599
600 let new_engine = OxiosEngine::new("openai/gpt-4o");
601 handle.swap(new_engine);
602
603 let e = handle.get();
604 assert_eq!(e.default_model_id(), "openai/gpt-4o");
605 }
606
607 #[test]
608 fn test_engine_handle_swap_preserves_old_arc() {
609 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
611 let handle = EngineHandle::new(Arc::new(engine));
612
613 let old = handle.get();
614 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
615
616 handle.swap(OxiosEngine::new("openai/gpt-4o"));
617
618 assert_eq!(old.default_model_id(), "anthropic/claude-sonnet-4-20250514");
620
621 let current = handle.get();
623 assert_eq!(current.default_model_id(), "openai/gpt-4o");
624 }
625
626 #[test]
629 fn test_rfc014_phase_d_default_fields_are_none() {
630 let engine = OxiosEngine::new("anthropic/claude-sonnet-4-20250514");
634 assert!(engine.authorizer().is_none());
635 assert!(engine.tracer().is_none());
636 assert!(engine.cost_tracker().is_none());
637
638 let engine = OxiosEngine::from_config("anthropic/claude-sonnet-4-20250514", None);
639 assert!(engine.authorizer().is_none());
640 assert!(engine.tracer().is_none());
641 assert!(engine.cost_tracker().is_none());
642
643 let engine = OxiosEngine::builder()
644 .default_model("openai/gpt-4o")
645 .build();
646 assert!(engine.authorizer().is_none());
647 assert!(engine.tracer().is_none());
648 assert!(engine.cost_tracker().is_none());
649
650 let (engine, _rc) = OxiosEngine::builder()
651 .default_model("openai/gpt-4o")
652 .build_with_routing();
653 assert!(engine.authorizer().is_none());
654 assert!(engine.tracer().is_none());
655 assert!(engine.cost_tracker().is_none());
656 }
657
658 #[test]
659 fn test_rfc014_phase_d_with_tracer() {
660 let tracer = Arc::new(oxi_sdk::Tracer::new());
662 let engine = OxiosEngine::builder()
663 .default_model("openai/gpt-4o")
664 .with_tracer(tracer.clone())
665 .build();
666 assert!(engine.tracer().is_some());
667 assert!(engine.authorizer().is_none());
668 assert!(engine.cost_tracker().is_none());
669 }
670
671 #[test]
672 fn test_rfc014_phase_d_with_cost_tracker() {
673 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
678 let model_registry = oxi_for_registry.models_arc();
679 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
680 model_registry,
681 oxi_sdk::CostTrackerConfig::default(),
682 ));
683 let engine = OxiosEngine::builder()
684 .default_model("openai/gpt-4o")
685 .with_cost_tracker(cost_tracker)
686 .build();
687 assert!(engine.cost_tracker().is_some());
688 assert!(engine.authorizer().is_none());
689 assert!(engine.tracer().is_none());
690 }
691
692 #[test]
693 fn test_rfc014_phase_d_with_authorizer() {
694 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
696 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
697 let engine = OxiosEngine::builder()
698 .default_model("openai/gpt-4o")
699 .with_authorizer(authorizer)
700 .build();
701 assert!(engine.authorizer().is_some());
702 assert!(engine.tracer().is_none());
703 assert!(engine.cost_tracker().is_none());
704 }
705
706 #[test]
707 fn test_rfc014_phase_d_all_three_handles() {
708 let audit = Arc::new(oxi_sdk::AuditLog::new(16));
712 let authorizer = Arc::new(oxi_sdk::Authorizer::new(audit));
713 let tracer = Arc::new(oxi_sdk::Tracer::new());
714 let oxi_for_registry = oxi_sdk::OxiBuilder::new().with_builtins().build();
715 let model_registry = oxi_for_registry.models_arc();
716 let cost_tracker = Arc::new(oxi_sdk::CostTracker::new(
717 model_registry,
718 oxi_sdk::CostTrackerConfig::default(),
719 ));
720
721 let engine = OxiosEngine::builder()
722 .default_model("openai/gpt-4o")
723 .api_key("openai", "sk-test")
724 .with_authorizer(authorizer)
725 .with_tracer(tracer)
726 .with_cost_tracker(cost_tracker)
727 .build();
728
729 assert!(engine.authorizer().is_some());
730 assert!(engine.tracer().is_some());
731 assert!(engine.cost_tracker().is_some());
732 assert_eq!(engine.default_model_id(), "openai/gpt-4o");
733 }
734}