chat_router/
lib.rs

1mod circuit_breaker;
2mod router;
3pub mod strategy;
4pub use circuit_breaker::CircuitBreakerConfig;
5pub use router::Router;
6pub use strategy::{RoutingStrategy, StrategyError};
7
8#[cfg(feature = "stream")]
9mod stream_router;
10#[cfg(feature = "stream")]
11pub use stream_router::StreamRouter;
12
13use std::marker::PhantomData;
14
15use chat_core::traits::CompletionProvider;
16
17#[cfg(feature = "stream")]
18use chat_core::traits::ChatProvider;
19
20use crate::circuit_breaker::CircuitBreaker;
21
22pub struct WithoutProvider;
23pub struct WithProvider;
24
25// ---------------------------------------------------------------------------
26// RouterBuilder — completion-only router
27// ---------------------------------------------------------------------------
28
29pub struct RouterBuilder<M = WithoutProvider> {
30    providers: Vec<Box<dyn CompletionProvider>>,
31    strategy: Option<Box<dyn RoutingStrategy>>,
32    circuit_breaker_config: Option<CircuitBreakerConfig>,
33    _m: PhantomData<M>,
34}
35
36impl Default for RouterBuilder<WithoutProvider> {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl RouterBuilder<WithoutProvider> {
43    pub fn new() -> Self {
44        Self {
45            providers: Vec::new(),
46            strategy: None,
47            circuit_breaker_config: None,
48            _m: PhantomData,
49        }
50    }
51
52    pub fn add_provider(
53        self,
54        provider: impl CompletionProvider + 'static,
55    ) -> RouterBuilder<WithProvider> {
56        let mut providers = self.providers;
57        providers.push(Box::new(provider));
58        RouterBuilder {
59            providers,
60            strategy: self.strategy,
61            circuit_breaker_config: self.circuit_breaker_config,
62            _m: PhantomData,
63        }
64    }
65}
66
67impl RouterBuilder<WithProvider> {
68    pub fn add_provider(mut self, provider: impl CompletionProvider + 'static) -> Self {
69        self.providers.push(Box::new(provider));
70        self
71    }
72
73    pub fn with_strategy(mut self, strategy: impl RoutingStrategy + 'static) -> Self {
74        self.strategy = Some(Box::new(strategy));
75        self
76    }
77
78    pub fn circuit_breaker(mut self, config: CircuitBreakerConfig) -> Self {
79        self.circuit_breaker_config = Some(config);
80        self
81    }
82
83    pub fn build(self) -> Router {
84        let count = self.providers.len();
85        Router {
86            providers: self.providers,
87            strategy: self.strategy,
88            circuit_breaker: self
89                .circuit_breaker_config
90                .map(|config| CircuitBreaker::new(config, count)),
91        }
92    }
93}
94
95// ---------------------------------------------------------------------------
96// StreamRouterBuilder — streaming router (all providers must support stream)
97// ---------------------------------------------------------------------------
98
99#[cfg(feature = "stream")]
100pub struct StreamRouterBuilder<M = WithoutProvider> {
101    providers: Vec<Box<dyn ChatProvider>>,
102    strategy: Option<Box<dyn RoutingStrategy>>,
103    circuit_breaker_config: Option<CircuitBreakerConfig>,
104    _m: PhantomData<M>,
105}
106
107#[cfg(feature = "stream")]
108impl Default for StreamRouterBuilder<WithoutProvider> {
109    fn default() -> Self {
110        Self::new()
111    }
112}
113
114#[cfg(feature = "stream")]
115impl StreamRouterBuilder<WithoutProvider> {
116    pub fn new() -> Self {
117        Self {
118            providers: Vec::new(),
119            strategy: None,
120            circuit_breaker_config: None,
121            _m: PhantomData,
122        }
123    }
124
125    pub fn add_provider(
126        self,
127        provider: impl ChatProvider + 'static,
128    ) -> StreamRouterBuilder<WithProvider> {
129        let mut providers = self.providers;
130        providers.push(Box::new(provider));
131        StreamRouterBuilder {
132            providers,
133            strategy: self.strategy,
134            circuit_breaker_config: self.circuit_breaker_config,
135            _m: PhantomData,
136        }
137    }
138}
139
140#[cfg(feature = "stream")]
141impl StreamRouterBuilder<WithProvider> {
142    pub fn add_provider(mut self, provider: impl ChatProvider + 'static) -> Self {
143        self.providers.push(Box::new(provider));
144        self
145    }
146
147    pub fn with_strategy(mut self, strategy: impl RoutingStrategy + 'static) -> Self {
148        self.strategy = Some(Box::new(strategy));
149        self
150    }
151
152    pub fn circuit_breaker(mut self, config: CircuitBreakerConfig) -> Self {
153        self.circuit_breaker_config = Some(config);
154        self
155    }
156
157    pub fn build(self) -> StreamRouter {
158        let count = self.providers.len();
159        StreamRouter {
160            providers: self.providers,
161            strategy: self.strategy,
162            circuit_breaker: self
163                .circuit_breaker_config
164                .map(|config| CircuitBreaker::new(config, count)),
165            last_used: None,
166        }
167    }
168}
chat_router/lib.rs

chat_router/
lib.rs
