Skip to main content

simple_agents_core/
routing.rs

1//! Routing integration for SimpleAgents core.
2
3use simple_agent_type::prelude::{
4    CompletionChunk, CompletionRequest, CompletionResponse, Provider, Result, SimpleAgentsError,
5};
6use simple_agents_router::{
7    CostRouter, CostRouterConfig, FallbackRouter, FallbackRouterConfig, LatencyRouter,
8    LatencyRouterConfig, RoundRobinRouter,
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    pub(crate) async fn stream(
52        &self,
53        request: &CompletionRequest,
54    ) -> Result<Box<dyn futures_core::Stream<Item = Result<CompletionChunk>> + Send + Unpin>> {
55        match self {
56            Self::Direct(provider) => {
57                eprintln!(
58                    "RouterEngine.stream: provider={}, stream={:?}",
59                    provider.name(),
60                    request.stream
61                );
62                let provider_request = provider.transform_request(request)?;
63                provider.execute_stream(provider_request).await
64            }
65            Self::RoundRobin(router) => router.stream(request).await,
66            Self::Latency(router) => router.stream(request).await,
67            Self::Cost(router) => router.stream(request).await,
68            Self::Fallback(router) => router.stream(request).await,
69        }
70    }
71}
72
73impl RoutingMode {
74    pub(crate) fn build_router(&self, providers: Vec<Arc<dyn Provider>>) -> Result<RouterEngine> {
75        if providers.is_empty() {
76            return Err(SimpleAgentsError::Routing(
77                "no providers configured".to_string(),
78            ));
79        }
80
81        match self {
82            RoutingMode::Direct => Ok(RouterEngine::Direct(
83                providers
84                    .first()
85                    .ok_or_else(|| {
86                        SimpleAgentsError::Routing("no providers configured".to_string())
87                    })?
88                    .clone(),
89            )),
90            RoutingMode::RoundRobin => {
91                Ok(RouterEngine::RoundRobin(RoundRobinRouter::new(providers)?))
92            }
93            RoutingMode::Latency(config) => Ok(RouterEngine::Latency(LatencyRouter::with_config(
94                providers,
95                config.clone(),
96            )?)),
97            RoutingMode::Cost(config) => Ok(RouterEngine::Cost(CostRouter::new(
98                providers,
99                config.clone(),
100            )?)),
101            RoutingMode::Fallback(config) => Ok(RouterEngine::Fallback(
102                FallbackRouter::with_config(providers, *config)?,
103            )),
104        }
105    }
106}