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                let provider_request = provider.transform_request(request)?;
58                provider.execute_stream(provider_request).await
59            }
60            Self::RoundRobin(router) => router.stream(request).await,
61            Self::Latency(router) => router.stream(request).await,
62            Self::Cost(router) => router.stream(request).await,
63            Self::Fallback(router) => router.stream(request).await,
64        }
65    }
66}
67
68impl RoutingMode {
69    pub(crate) fn build_router(&self, providers: Vec<Arc<dyn Provider>>) -> Result<RouterEngine> {
70        if providers.is_empty() {
71            return Err(SimpleAgentsError::Routing(
72                "no providers configured".to_string(),
73            ));
74        }
75
76        match self {
77            RoutingMode::Direct => Ok(RouterEngine::Direct(
78                providers
79                    .first()
80                    .ok_or_else(|| {
81                        SimpleAgentsError::Routing("no providers configured".to_string())
82                    })?
83                    .clone(),
84            )),
85            RoutingMode::RoundRobin => {
86                Ok(RouterEngine::RoundRobin(RoundRobinRouter::new(providers)?))
87            }
88            RoutingMode::Latency(config) => Ok(RouterEngine::Latency(LatencyRouter::with_config(
89                providers,
90                config.clone(),
91            )?)),
92            RoutingMode::Cost(config) => Ok(RouterEngine::Cost(CostRouter::new(
93                providers,
94                config.clone(),
95            )?)),
96            RoutingMode::Fallback(config) => Ok(RouterEngine::Fallback(
97                FallbackRouter::with_config(providers, *config)?,
98            )),
99        }
100    }
101}