Skip to main content

simple_agents_core/
routing.rs

1//! Routing integration for SimpleAgents core.
2
3use simple_agents_router::{
4    CostRouter, CostRouterConfig, FallbackRouter, FallbackRouterConfig, LatencyRouter,
5    LatencyRouterConfig, RoundRobinRouter,
6};
7use simple_agent_type::prelude::{
8    CompletionRequest, CompletionResponse, Provider, Result, SimpleAgentsError,
9};
10use std::sync::Arc;
11
12/// Routing modes supported by the core client.
13#[derive(Debug, Clone, Default)]
14pub enum RoutingMode {
15    /// Direct execution against the first provider.
16    Direct,
17    /// Round-robin routing across providers.
18    #[default]
19    RoundRobin,
20    /// Latency-based routing with configurable smoothing.
21    Latency(LatencyRouterConfig),
22    /// Cost-based routing with per-provider costs.
23    Cost(CostRouterConfig),
24    /// Fallback routing (try providers in order).
25    Fallback(FallbackRouterConfig),
26}
27
28pub(crate) enum RouterEngine {
29    Direct(Arc<dyn Provider>),
30    RoundRobin(RoundRobinRouter),
31    Latency(LatencyRouter),
32    Cost(CostRouter),
33    Fallback(FallbackRouter),
34}
35
36impl RouterEngine {
37    pub(crate) async fn complete(&self, request: &CompletionRequest) -> Result<CompletionResponse> {
38        match self {
39            Self::Direct(provider) => {
40                let provider_request = provider.transform_request(request)?;
41                let provider_response = provider.execute(provider_request).await?;
42                provider.transform_response(provider_response)
43            }
44            Self::RoundRobin(router) => router.complete(request).await,
45            Self::Latency(router) => router.complete(request).await,
46            Self::Cost(router) => router.complete(request).await,
47            Self::Fallback(router) => router.complete(request).await,
48        }
49    }
50}
51
52impl RoutingMode {
53    pub(crate) fn build_router(&self, providers: Vec<Arc<dyn Provider>>) -> Result<RouterEngine> {
54        if providers.is_empty() {
55            return Err(SimpleAgentsError::Routing(
56                "no providers configured".to_string(),
57            ));
58        }
59
60        match self {
61            RoutingMode::Direct => Ok(RouterEngine::Direct(
62                providers
63                    .first()
64                    .ok_or_else(|| {
65                        SimpleAgentsError::Routing("no providers configured".to_string())
66                    })?
67                    .clone(),
68            )),
69            RoutingMode::RoundRobin => {
70                Ok(RouterEngine::RoundRobin(RoundRobinRouter::new(providers)?))
71            }
72            RoutingMode::Latency(config) => Ok(RouterEngine::Latency(LatencyRouter::with_config(
73                providers,
74                config.clone(),
75            )?)),
76            RoutingMode::Cost(config) => Ok(RouterEngine::Cost(CostRouter::new(
77                providers,
78                config.clone(),
79            )?)),
80            RoutingMode::Fallback(config) => Ok(RouterEngine::Fallback(
81                FallbackRouter::with_config(providers, *config)?,
82            )),
83        }
84    }
85}