1#![deny(missing_docs)]
2pub mod output_parser;
42pub mod plan_parser;
44#[cfg(feature = "schema-validation")]
47pub mod schema_validator;
48
49pub use output_parser::{JsonListParser, JsonOutputParser, OutputParser, RegexOutputParser};
51pub use plan_parser::{ParsedStep, parse_plan_steps, steps_to_tasks};
52#[cfg(feature = "schema-validation")]
53pub use schema_validator::{SchemaValidator, retry_until_valid};
54
55mod complexity;
57mod entity_enhancer;
58mod relevance_scorer;
59mod retrieval_classifier;
60mod router;
61pub mod strategies;
63mod strategy_selector;
64mod summarizer;
65mod validator;
66
67pub use complexity::{ComplexityResult, ComplexityScorer, ComplexityScorerBuilder};
68pub use entity_enhancer::{
69 EnhancedEntity, EnhancedRelationship, EnhancementResult, EntityEnhancer, EntityEnhancerBuilder,
70 RelationType, SemanticEntityType,
71};
72pub use relevance_scorer::{RelevanceResult, RelevanceScorer, RelevanceScorerBuilder};
73pub use retrieval_classifier::{
74 ClassificationResult, RetrievalClassifier, RetrievalClassifierBuilder,
75 RetrievalNeed as LocalRetrievalNeed,
76};
77pub use router::{LocalRouter, LocalRouterBuilder, RouteResult};
78pub use strategies::{
79 ChainOfThoughtStrategy, ReActStrategy, ReasoningStrategy, ReflexionStrategy, StrategyPreset,
80 StrategyStep, TreeOfThoughtsStrategy,
81};
82pub use strategy_selector::{
83 RecommendedStrategy, StrategyResult, StrategySelector, StrategySelectorBuilder, TaskType,
84};
85pub use summarizer::{
86 ExtractedFact, FactCategory, LocalSummarizer, LocalSummarizerBuilder, SummarizationResult,
87};
88pub use validator::{LocalValidator, LocalValidatorBuilder, ValidationResult};
89
90use std::time::Instant;
91use tracing::{info, warn};
92
93#[derive(Clone, Debug)]
95pub struct LocalInferenceConfig {
96 pub routing_enabled: bool,
99 pub validation_enabled: bool,
101 pub complexity_enabled: bool,
103
104 pub summarization_enabled: bool,
107 pub retrieval_gating_enabled: bool,
109 pub relevance_scoring_enabled: bool,
111 pub strategy_selection_enabled: bool,
113 pub entity_enhancement_enabled: bool,
115
116 pub routing_model: Option<String>,
119 pub validation_model: Option<String>,
121 pub complexity_model: Option<String>,
123 pub summarization_model: Option<String>,
125 pub retrieval_model: Option<String>,
127 pub relevance_model: Option<String>,
129 pub strategy_model: Option<String>,
131 pub entity_model: Option<String>,
133
134 pub log_inference: bool,
136}
137
138impl Default for LocalInferenceConfig {
139 fn default() -> Self {
140 Self {
141 routing_enabled: false,
143 validation_enabled: false,
144 complexity_enabled: false,
145 summarization_enabled: false,
147 retrieval_gating_enabled: false,
148 relevance_scoring_enabled: false,
149 strategy_selection_enabled: false,
150 entity_enhancement_enabled: false,
151 routing_model: Some("lfm2-350m".to_string()),
153 validation_model: Some("lfm2-350m".to_string()),
154 complexity_model: Some("lfm2-350m".to_string()),
155 summarization_model: Some("lfm2-1.2b".to_string()),
157 retrieval_model: Some("lfm2-350m".to_string()),
158 relevance_model: Some("lfm2-350m".to_string()),
159 strategy_model: Some("lfm2-1.2b".to_string()),
160 entity_model: Some("lfm2-350m".to_string()),
161 log_inference: true,
162 }
163 }
164}
165
166impl LocalInferenceConfig {
167 pub fn tier1_enabled() -> Self {
169 Self {
170 routing_enabled: true,
171 validation_enabled: true,
172 complexity_enabled: true,
173 ..Default::default()
174 }
175 }
176
177 pub fn tier2_enabled() -> Self {
179 Self {
180 summarization_enabled: true,
181 retrieval_gating_enabled: true,
182 relevance_scoring_enabled: true,
183 strategy_selection_enabled: true,
184 entity_enhancement_enabled: true,
185 ..Default::default()
186 }
187 }
188
189 pub fn all_enabled() -> Self {
191 Self {
192 routing_enabled: true,
193 validation_enabled: true,
194 complexity_enabled: true,
195 summarization_enabled: true,
196 retrieval_gating_enabled: true,
197 relevance_scoring_enabled: true,
198 strategy_selection_enabled: true,
199 entity_enhancement_enabled: true,
200 ..Default::default()
201 }
202 }
203
204 pub fn routing_only() -> Self {
206 Self {
207 routing_enabled: true,
208 ..Default::default()
209 }
210 }
211
212 pub fn validation_only() -> Self {
214 Self {
215 validation_enabled: true,
216 ..Default::default()
217 }
218 }
219
220 pub fn summarization_only() -> Self {
222 Self {
223 summarization_enabled: true,
224 ..Default::default()
225 }
226 }
227}
228
229pub fn log_inference(task: &str, model: &str, latency_ms: u64, success: bool) {
231 if success {
232 info!(
233 target: "local_llm",
234 task = task,
235 model = model,
236 latency_ms = latency_ms,
237 "Local inference completed"
238 );
239 } else {
240 warn!(
241 target: "local_llm",
242 task = task,
243 model = model,
244 latency_ms = latency_ms,
245 "Local inference failed, falling back to pattern-based"
246 );
247 }
248}
249
250pub struct InferenceTimer {
252 start: Instant,
253 task: String,
254 model: String,
255}
256
257impl InferenceTimer {
258 pub fn new(task: impl Into<String>, model: impl Into<String>) -> Self {
260 Self {
261 start: Instant::now(),
262 task: task.into(),
263 model: model.into(),
264 }
265 }
266
267 pub fn finish(self, success: bool) {
269 let latency_ms = self.start.elapsed().as_millis() as u64;
270 log_inference(&self.task, &self.model, latency_ms, success);
271 }
272
273 pub fn elapsed_ms(&self) -> u64 {
275 self.start.elapsed().as_millis() as u64
276 }
277}
278
279#[cfg(test)]
280mod tests {
281 use super::*;
282
283 #[test]
284 fn test_config_default() {
285 let config = LocalInferenceConfig::default();
286 assert!(!config.routing_enabled);
287 assert!(!config.validation_enabled);
288 assert!(!config.complexity_enabled);
289 assert!(!config.summarization_enabled);
290 assert!(!config.retrieval_gating_enabled);
291 assert!(!config.relevance_scoring_enabled);
292 }
293
294 #[test]
295 fn test_config_tier1_enabled() {
296 let config = LocalInferenceConfig::tier1_enabled();
297 assert!(config.routing_enabled);
298 assert!(config.validation_enabled);
299 assert!(config.complexity_enabled);
300 assert!(!config.summarization_enabled);
301 }
302
303 #[test]
304 fn test_config_tier2_enabled() {
305 let config = LocalInferenceConfig::tier2_enabled();
306 assert!(!config.routing_enabled);
307 assert!(config.summarization_enabled);
308 assert!(config.retrieval_gating_enabled);
309 assert!(config.relevance_scoring_enabled);
310 assert!(config.strategy_selection_enabled);
311 assert!(config.entity_enhancement_enabled);
312 }
313
314 #[test]
315 fn test_config_all_enabled() {
316 let config = LocalInferenceConfig::all_enabled();
317 assert!(config.routing_enabled);
318 assert!(config.validation_enabled);
319 assert!(config.complexity_enabled);
320 assert!(config.summarization_enabled);
321 assert!(config.retrieval_gating_enabled);
322 assert!(config.relevance_scoring_enabled);
323 assert!(config.strategy_selection_enabled);
324 assert!(config.entity_enhancement_enabled);
325 }
326
327 #[test]
328 fn test_config_summarization_only() {
329 let config = LocalInferenceConfig::summarization_only();
330 assert!(!config.routing_enabled);
331 assert!(config.summarization_enabled);
332 assert_eq!(config.summarization_model, Some("lfm2-1.2b".to_string()));
333 }
334
335 #[test]
336 fn test_inference_timer() {
337 let timer = InferenceTimer::new("test_task", "test_model");
338 std::thread::sleep(std::time::Duration::from_millis(10));
339 assert!(timer.elapsed_ms() >= 10);
340 }
341}