Skip to main content

rullama_reasoning/
lib.rs

1#![deny(missing_docs)]
2//! # rullama reasoning
3//!
4//! Layer 3 — Intelligence. Provider-agnostic reasoning primitives for the
5//! rullama.
6//!
7//! This crate owns:
8//!
9//! - **`plan_parser`** — extract numbered task steps from LLM plan output.
10//! - **`output_parser`** — parse structured data (JSON, regex) from raw LLM
11//!   text.
12//! - **Local inference scorers** — provider-agnostic TIER 1/2 components for
13//!   routing, validation, complexity scoring, summarisation, retrieval
14//!   gating, relevance ranking, strategy selection, and entity enhancement.
15//!   All accept `Arc<dyn Provider>` and fall back to pattern-based logic
16//!   when the provider is unavailable.
17//!
18//! ## Scope note — `prompting` stays in `rullama-knowledge`
19//!
20//! The original architectural plan (`sleepy-popping-falcon.md`) called for
21//! `prompting` to move here too. It didn't — `prompting` is tightly coupled
22//! to `bks_pks` inside `rullama-knowledge` and a move would have pulled
23//! an entire knowledge-store dependency into this crate. The deviation is
24//! intentional; tests and consumers of prompting should continue to target
25//! `rullama_prompting`.
26//!
27//! ## Configuration
28//!
29//! Enable via `LocalInferenceConfig`:
30//! ```toml
31//! [local_llm]
32//! enabled = true
33//! use_for_routing = true
34//! use_for_validation = true
35//! use_for_complexity = true
36//! use_for_summarization = true
37//! ```
38
39// ── Parsers (moved from rullama-core) ─────────────────────────────────
40/// Structured output parsers for LLM responses.
41pub mod output_parser;
42/// Plan text parser for extracting steps from LLM output.
43pub mod plan_parser;
44/// JSON Schema validation + retry-on-violation orchestration. Behind the
45/// `schema-validation` feature.
46#[cfg(feature = "schema-validation")]
47pub mod schema_validator;
48
49// Flat re-exports for convenience.
50pub 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
55// ── Scorers (moved from rullama-agent::reasoning) ────────────────────
56mod complexity;
57mod entity_enhancer;
58mod relevance_scorer;
59mod retrieval_classifier;
60mod router;
61/// Reasoning strategies: CoT, ReAct, Reflexion, Tree-of-Thoughts.
62pub 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/// Configuration for local inference components.
94#[derive(Clone, Debug)]
95pub struct LocalInferenceConfig {
96    // TIER 1: Quick Wins
97    /// Enable local routing
98    pub routing_enabled: bool,
99    /// Enable local validation
100    pub validation_enabled: bool,
101    /// Enable complexity scoring
102    pub complexity_enabled: bool,
103
104    // TIER 2: Context & Retrieval
105    /// Enable local summarization for tiered memory
106    pub summarization_enabled: bool,
107    /// Enable local retrieval gating
108    pub retrieval_gating_enabled: bool,
109    /// Enable local relevance scoring
110    pub relevance_scoring_enabled: bool,
111    /// Enable local strategy selection
112    pub strategy_selection_enabled: bool,
113    /// Enable local entity enhancement
114    pub entity_enhancement_enabled: bool,
115
116    // Model selection per task
117    /// Model ID to use for routing (fast model preferred)
118    pub routing_model: Option<String>,
119    /// Model ID to use for validation (fast model preferred)
120    pub validation_model: Option<String>,
121    /// Model ID to use for complexity scoring (fast model preferred)
122    pub complexity_model: Option<String>,
123    /// Model ID to use for summarization (larger model preferred)
124    pub summarization_model: Option<String>,
125    /// Model ID to use for retrieval classification (fast model preferred)
126    pub retrieval_model: Option<String>,
127    /// Model ID to use for relevance scoring (fast model preferred)
128    pub relevance_model: Option<String>,
129    /// Model ID to use for strategy selection (larger model preferred)
130    pub strategy_model: Option<String>,
131    /// Model ID to use for entity enhancement (fast model preferred)
132    pub entity_model: Option<String>,
133
134    /// Log all local inference calls
135    pub log_inference: bool,
136}
137
138impl Default for LocalInferenceConfig {
139    fn default() -> Self {
140        Self {
141            // TIER 1
142            routing_enabled: false,
143            validation_enabled: false,
144            complexity_enabled: false,
145            // TIER 2
146            summarization_enabled: false,
147            retrieval_gating_enabled: false,
148            relevance_scoring_enabled: false,
149            strategy_selection_enabled: false,
150            entity_enhancement_enabled: false,
151            // Model selection - TIER 1
152            routing_model: Some("lfm2-350m".to_string()),
153            validation_model: Some("lfm2-350m".to_string()),
154            complexity_model: Some("lfm2-350m".to_string()),
155            // Model selection - TIER 2
156            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    /// Create a config with all TIER 1 features enabled
168    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    /// Create a config with all TIER 2 features enabled
178    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    /// Create a config with all features enabled
190    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    /// Create a config with only routing enabled
205    pub fn routing_only() -> Self {
206        Self {
207            routing_enabled: true,
208            ..Default::default()
209        }
210    }
211
212    /// Create a config with only validation enabled
213    pub fn validation_only() -> Self {
214        Self {
215            validation_enabled: true,
216            ..Default::default()
217        }
218    }
219
220    /// Create a config with only summarization enabled
221    pub fn summarization_only() -> Self {
222        Self {
223            summarization_enabled: true,
224            ..Default::default()
225        }
226    }
227}
228
229/// Log a local inference event.
230pub 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
250/// Measure inference latency.
251pub struct InferenceTimer {
252    start: Instant,
253    task: String,
254    model: String,
255}
256
257impl InferenceTimer {
258    /// Create a new inference timer for the given task and model.
259    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    /// Stop the timer and log the inference event.
268    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    /// Return the elapsed time in milliseconds since the timer was created.
274    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}