simple_agents_core/
routing.rs1use 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#[derive(Debug, Clone, Default)]
14pub enum RoutingMode {
15 Direct,
17 #[default]
19 RoundRobin,
20 Latency(LatencyRouterConfig),
22 Cost(CostRouterConfig),
24 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}