simple_agents_core/
routing.rs1use simple_agents_router::{
4 CostRouter, CostRouterConfig, FallbackRouter, FallbackRouterConfig, LatencyRouter,
5 LatencyRouterConfig, RoundRobinRouter,
6};
7use simple_agent_type::prelude::{
8 CompletionChunk, 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 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}