simple_agents_core/
routing.rs1use 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#[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 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}