feat: add memoization support for LLM operations in retrieval pipeline

- Add MemoOpType variants for NodeEvaluation, SufficiencyCheck,
  ComplexityDetection, and QueryDecomposition
- Implement caching for complexity detection with new ComplexityDetector
  memo store integration
- Add query decomposition caching with serialization/deserialization
 logic for DecompositionResult
- Integrate memo stores across retrieval stages (analyze, evaluate,
  search) to cache LLM evaluations
- Add NodeEvaluation caching for both single and batch LLM evaluations
- Implement sufficiency check caching in LlmJudge component
- Update module exports to include MemoOpType in public API

Author: zTgx

rust/src/llm/memo/mod.rs

@@ -11,4 +11,4 @@ mod store;
 mod types;
 
 pub use store::MemoStore;
-pub use types::{MemoKey, MemoValue, PilotDecisionValue};
+pub use types::{MemoKey, MemoOpType, MemoValue, PilotDecisionValue};

rust/src/llm/memo/types.rs

@@ -23,6 +23,18 @@ pub enum MemoOpType {
     /// Content extraction result.
     Extraction,
 
+    /// LLM node evaluation during retrieval.
+    NodeEvaluation,
+
+    /// Sufficiency check result.
+    SufficiencyCheck,
+
+    /// Query complexity detection.
+    ComplexityDetection,
+
+    /// Query decomposition.
+    QueryDecomposition,
+
     /// Custom operation type.
     Custom(u8),
 }
@@ -35,6 +47,10 @@ impl MemoOpType {
             MemoOpType::PilotDecision => 1,
             MemoOpType::QueryAnalysis => 2,
             MemoOpType::Extraction => 3,
+            MemoOpType::NodeEvaluation => 4,
+            MemoOpType::SufficiencyCheck => 5,
+            MemoOpType::ComplexityDetection => 6,
+            MemoOpType::QueryDecomposition => 7,
             MemoOpType::Custom(n) => 100 + n,
         }
     }

rust/src/retrieval/complexity/detector.rs

@@ -9,40 +9,105 @@
 use std::collections::HashSet;
 
 use super::QueryComplexity;
+use crate::llm::memo::{MemoKey, MemoOpType, MemoStore, MemoValue};
+use crate::utils::fingerprint::Fingerprint;
 
 /// Query complexity detector.
 ///
 /// Uses LLM for classification when available; falls back to heuristic rules.
 pub struct ComplexityDetector {
     /// Optional LLM client for LLM-based detection.
     llm_client: Option<crate::llm::LlmClient>,
+    /// Memo store for caching complexity detection results.
+    memo_store: Option<MemoStore>,
 }
 
 impl ComplexityDetector {
     /// Create a new complexity detector (heuristic only).
     pub fn new() -> Self {
-        Self { llm_client: None }
+        Self {
+            llm_client: None,
+            memo_store: None,
+        }
     }
 
     /// Create with LLM client for accurate detection.
     pub fn with_llm_client(client: crate::llm::LlmClient) -> Self {
         Self {
             llm_client: Some(client),
+            memo_store: None,
         }
     }
 
+    /// Add memo store for caching complexity detection results.
+    pub fn with_memo_store(mut self, store: MemoStore) -> Self {
+        self.memo_store = Some(store);
+        self
+    }
+
     /// Detect the complexity of a query.
     ///
     /// Uses LLM when available; falls back to heuristic rules.
     pub async fn detect(&self, query: &str) -> QueryComplexity {
-        if let Some(ref client) = self.llm_client {
-            if let Some(complexity) = crate::retrieval::pilot::detect_with_llm(client, query).await
+        // Check memo cache
+        if let Some(ref store) = self.memo_store {
+            let cache_key = Self::build_cache_key(query);
+            if let Some(cached) = store.get(&cache_key) {
+                if let Some(complexity) = Self::deserialize_complexity(&cached) {
+                    return complexity;
+                }
+            }
+        }
+
+        let result = if let Some(ref client) = self.llm_client {
+            if let Some(complexity) =
+                crate::retrieval::pilot::detect_with_llm(client, query).await
             {
-                return complexity;
+                complexity
+            } else {
+                tracing::warn!("LLM complexity detection failed, falling back to heuristic");
+                self.detect_heuristic(query)
             }
-            tracing::warn!("LLM complexity detection failed, falling back to heuristic");
+        } else {
+            self.detect_heuristic(query)
+        };
+
+        // Cache the result
+        if let Some(ref store) = self.memo_store {
+            let cache_key = Self::build_cache_key(query);
+            store.put_with_tokens(
+                cache_key,
+                MemoValue::Text(format!("{:?}", result)),
+                (query.len() / 4) as u64,
+            );
+        }
+
+        result
+    }
+
+    /// Build a cache key for complexity detection.
+    fn build_cache_key(query: &str) -> MemoKey {
+        let fp = Fingerprint::from_str(query);
+        MemoKey {
+            op_type: MemoOpType::ComplexityDetection,
+            input_fp: fp,
+            model_id: None,
+            version: 1,
+            context_fp: Fingerprint::zero(),
+        }
+    }
+
+    /// Deserialize a QueryComplexity from a MemoValue.
+    fn deserialize_complexity(value: &MemoValue) -> Option<QueryComplexity> {
+        match value {
+            MemoValue::Text(s) => match s.as_str() {
+                "Simple" => Some(QueryComplexity::Simple),
+                "Medium" => Some(QueryComplexity::Medium),
+                "Complex" => Some(QueryComplexity::Complex),
+                _ => None,
+            },
+            _ => None,
         }
-        self.detect_heuristic(query)
     }
 
     /// Heuristic-based fallback: keyword matching + word count.

rust/src/retrieval/decompose.rs

@@ -48,6 +48,8 @@ use serde::{Deserialize, Serialize};
 use tracing::{debug, info};
 
 use crate::llm::{LlmClient, LlmExecutor};
+use crate::llm::memo::{MemoKey, MemoOpType, MemoStore, MemoValue};
+use crate::utils::fingerprint::Fingerprint;
 
 /// Sub-query resulting from decomposition.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -204,6 +206,8 @@ pub struct QueryDecomposer {
     llm_client: Option<LlmClient>,
     /// LLM executor for unified execution (optional).
     llm_executor: Option<LlmExecutor>,
+    /// Memo store for caching decomposition results.
+    memo_store: Option<MemoStore>,
 }
 
 impl Default for QueryDecomposer {
@@ -219,6 +223,7 @@ impl QueryDecomposer {
             config,
             llm_client: None,
             llm_executor: None,
+            memo_store: None,
         }
     }
 
@@ -234,6 +239,12 @@ impl QueryDecomposer {
         self
     }
 
+    /// Add memo store for caching decomposition results.
+    pub fn with_memo_store(mut self, store: MemoStore) -> Self {
+        self.memo_store = Some(store);
+        self
+    }
+
     /// Decompose a query into sub-queries.
     pub async fn decompose(&self, query: &str) -> crate::error::Result<DecompositionResult> {
         // Check if decomposition is needed
@@ -244,23 +255,71 @@ impl QueryDecomposer {
             ));
         }
 
+        // Check memo cache
+        if let Some(ref store) = self.memo_store {
+            let cache_key = Self::build_cache_key(query);
+            if let Some(cached) = store.get(&cache_key) {
+                if let Some(result) = Self::deserialize_decomposition(&cached) {
+                    tracing::debug!("Memo cache hit for query decomposition");
+                    return Ok(result);
+                }
+            }
+        }
+
         info!("Decomposing complex query: '{}'", query);
 
         // Try LLM-based decomposition if available
-        if self.config.use_llm && (self.llm_client.is_some() || self.llm_executor.is_some()) {
+        let result = if self.config.use_llm && (self.llm_client.is_some() || self.llm_executor.is_some()) {
             match self.llm_decompose(query).await {
-                Ok(result) => return Ok(result),
+                Ok(result) => result,
                 Err(e) => {
                     debug!(
                         "LLM decomposition failed, falling back to rule-based: {}",
                         e
                     );
+                    self.rule_based_decompose(query)?
                 }
             }
+        } else {
+            self.rule_based_decompose(query)?
+        };
+
+        // Cache the result
+        if let Some(ref store) = self.memo_store {
+            let cache_key = Self::build_cache_key(query);
+            if let Ok(json) = serde_json::to_value(&CachedDecomposition::from_result(&result)) {
+                store.put_with_tokens(
+                    cache_key,
+                    MemoValue::Json(json),
+                    (query.len() / 4) as u64,
+                );
+            }
         }
 
-        // Fall back to rule-based decomposition
-        self.rule_based_decompose(query)
+        Ok(result)
+    }
+
+    /// Build a cache key for query decomposition.
+    fn build_cache_key(query: &str) -> MemoKey {
+        let fp = Fingerprint::from_str(query);
+        MemoKey {
+            op_type: MemoOpType::QueryDecomposition,
+            input_fp: fp,
+            model_id: None,
+            version: 1,
+            context_fp: Fingerprint::zero(),
+        }
+    }
+
+    /// Deserialize a DecompositionResult from a MemoValue.
+    fn deserialize_decomposition(value: &MemoValue) -> Option<DecompositionResult> {
+        match value {
+            MemoValue::Json(json) => {
+                let cached: CachedDecomposition = serde_json::from_value(json.clone()).ok()?;
+                Some(cached.into_result())
+            }
+            _ => None,
+        }
     }
 
     /// Check if a query should be decomposed.
@@ -536,6 +595,72 @@ fn extract_json(text: &str) -> String {
     text.to_string()
 }
 
+/// Serializable decomposition result for caching.
+///
+/// Only caches the essential fields needed to reconstruct a DecompositionResult.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct CachedSubQuery {
+    text: String,
+    priority: u8,
+    query_type: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct CachedDecomposition {
+    original: String,
+    sub_queries: Vec<CachedSubQuery>,
+    was_decomposed: bool,
+    reason: String,
+}
+
+impl CachedDecomposition {
+    fn from_result(result: &DecompositionResult) -> Self {
+        Self {
+            original: result.original.clone(),
+            sub_queries: result
+                .sub_queries
+                .iter()
+                .map(|sq| CachedSubQuery {
+                    text: sq.text.clone(),
+                    priority: sq.priority,
+                    query_type: format!("{:?}", sq.query_type),
+                })
+                .collect(),
+            was_decomposed: result.was_decomposed,
+            reason: result.reason.clone(),
+        }
+    }
+
+    fn into_result(self) -> DecompositionResult {
+        let sub_queries: Vec<SubQuery> = self
+            .sub_queries
+            .into_iter()
+            .map(|csq| SubQuery {
+                text: csq.text,
+                priority: csq.priority,
+                query_type: match csq.query_type.as_str() {
+                    "Fact" => SubQueryType::Fact,
+                    "Explanation" => SubQueryType::Explanation,
+                    "Comparison" => SubQueryType::Comparison,
+                    "Synthesis" => SubQueryType::Synthesis,
+                    "Navigation" => SubQueryType::Navigation,
+                    _ => SubQueryType::Fact,
+                },
+                complexity: SubQueryComplexity::Simple,
+                depends_on: vec![],
+                path_constraint: None,
+            })
+            .collect();
+        DecompositionResult {
+            original: self.original,
+            sub_queries,
+            was_decomposed: self.was_decomposed,
+            reason: self.reason,
+            total_complexity: 0.5,
+        }
+    }
+}
+
 /// Result aggregator for multi-turn retrieval.
 #[derive(Debug, Clone)]
 pub struct SubQueryResult {

rust/src/retrieval/pipeline_retriever.rs

@@ -13,6 +13,7 @@ use super::content::ContentAggregatorConfig;
 use super::pipeline::RetrievalOrchestrator;
 use super::retriever::{CostEstimate, Retriever, RetrieverError, RetrieverResult};
 use super::stages::{AnalyzeStage, EvaluateStage, PlanStage, SearchStage};
+use super::strategy::LlmStrategy;
 use super::stream::RetrieveEventReceiver;
 use super::types::{RetrieveOptions, RetrieveResponse};
 use crate::document::{DocumentTree, ReasoningIndex};
@@ -112,6 +113,9 @@ impl PipelineRetriever {
         if let Some(ref client) = self.llm_client {
             analyze_stage = analyze_stage.with_llm_client(client.clone());
         }
+        if let Some(ref store) = self.memo_store {
+            analyze_stage = analyze_stage.with_memo_store(store.clone());
+        }
         orchestrator = orchestrator.stage(analyze_stage);
 
         // Add plan stage
@@ -133,11 +137,21 @@ impl PipelineRetriever {
             }
 
             search_stage = search_stage.with_pilot(Arc::new(pilot));
+
+            // Create LLM strategy with memo store for node evaluation
+            let mut llm_strategy = LlmStrategy::new(client.clone());
+            if let Some(ref store) = self.memo_store {
+                llm_strategy = llm_strategy.with_memo_store(store.clone());
+            }
+            search_stage = search_stage.with_llm_strategy(llm_strategy);
         }
         orchestrator = orchestrator.stage(search_stage);
 
         // Add evaluate stage with optional content aggregator
         let mut evaluate_stage = EvaluateStage::new();
+        if let Some(ref store) = self.memo_store {
+            evaluate_stage = evaluate_stage.with_memo_store(store.clone());
+        }
         if let Some(ref client) = self.llm_client {
             evaluate_stage = evaluate_stage.with_llm_judge(client.clone());
         }