use serde::{Deserialize, Serialize};
use super::context::{PruningStrategy, TokenEstimation};
use crate::document::NodeId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QueryComplexity {
Simple,
Medium,
Complex,
}
impl Default for QueryComplexity {
fn default() -> Self {
Self::Medium
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum StrategyPreference {
Auto,
ForceKeyword,
ForceLlm,
ForceHybrid,
ForceCrossDocument,
ForcePageRange,
}
impl Default for StrategyPreference {
fn default() -> Self {
Self::Auto
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SufficiencyLevel {
Sufficient,
PartialSufficient,
Insufficient,
}
impl Default for SufficiencyLevel {
fn default() -> Self {
Self::Insufficient
}
}
#[derive(Debug, Clone)]
pub struct RetrieveOptions {
pub top_k: usize,
pub beam_width: usize,
pub max_iterations: usize,
pub include_content: bool,
pub include_summaries: bool,
pub min_score: f32,
pub strategy: StrategyPreference,
pub sufficiency_check: bool,
pub max_tokens: usize,
pub enable_cache: bool,
pub pruning_strategy: super::PruningStrategy,
pub token_estimation: super::TokenEstimation,
pub use_async_context: bool,
pub streaming: bool,
pub document_graph: Option<std::sync::Arc<crate::graph::DocumentGraph>>,
pub fallback_chain: Vec<String>,
}
impl Default for RetrieveOptions {
fn default() -> Self {
Self {
top_k: 5,
beam_width: 3,
max_iterations: 10,
include_content: true,
include_summaries: true,
min_score: 0.1,
strategy: StrategyPreference::Auto,
sufficiency_check: true,
max_tokens: 4000,
enable_cache: true,
pruning_strategy: super::PruningStrategy::default(),
token_estimation: super::TokenEstimation::default(),
use_async_context: false,
streaming: false,
document_graph: None,
fallback_chain: vec!["beam".into(), "mcts".into(), "pure_pilot".into()],
}
}
}
impl RetrieveOptions {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_top_k(mut self, top_k: usize) -> Self {
self.top_k = top_k;
self
}
#[must_use]
pub fn with_beam_width(mut self, beam_width: usize) -> Self {
self.beam_width = beam_width;
self
}
#[must_use]
pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
self.max_iterations = max_iterations;
self
}
#[must_use]
pub fn with_include_content(mut self, include: bool) -> Self {
self.include_content = include;
self
}
#[must_use]
pub fn with_include_summaries(mut self, include: bool) -> Self {
self.include_summaries = include;
self
}
#[must_use]
pub fn with_min_score(mut self, min_score: f32) -> Self {
self.min_score = min_score;
self
}
#[must_use]
pub fn with_strategy(mut self, strategy: StrategyPreference) -> Self {
self.strategy = strategy;
self
}
#[must_use]
pub fn with_sufficiency_check(mut self, enable: bool) -> Self {
self.sufficiency_check = enable;
self
}
#[must_use]
pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
self.max_tokens = max_tokens;
self
}
#[must_use]
pub fn with_enable_cache(mut self, enable: bool) -> Self {
self.enable_cache = enable;
self
}
#[must_use]
pub fn with_pruning_strategy(mut self, strategy: PruningStrategy) -> Self {
self.pruning_strategy = strategy;
self
}
#[must_use]
pub fn with_token_estimation(mut self, mode: TokenEstimation) -> Self {
self.token_estimation = mode;
self
}
#[must_use]
pub fn with_async_context(mut self, enable: bool) -> Self {
self.use_async_context = enable;
self
}
#[must_use]
pub fn with_streaming(mut self, enable: bool) -> Self {
self.streaming = enable;
self
}
#[must_use]
pub fn with_document_graph(
mut self,
graph: std::sync::Arc<crate::graph::DocumentGraph>,
) -> Self {
self.document_graph = Some(graph);
self
}
#[must_use]
pub fn with_fallback_chain(mut self, chain: Vec<String>) -> Self {
self.fallback_chain = chain;
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrievalResult {
pub node_id: Option<String>,
pub title: String,
pub content: Option<String>,
pub summary: Option<String>,
pub score: f32,
pub depth: usize,
pub page_range: Option<(usize, usize)>,
}
impl RetrievalResult {
#[must_use]
pub fn new(title: impl Into<String>) -> Self {
Self {
node_id: None,
title: title.into(),
content: None,
summary: None,
score: 1.0,
depth: 0,
page_range: None,
}
}
#[must_use]
pub fn with_node_id(mut self, id: impl Into<String>) -> Self {
self.node_id = Some(id.into());
self
}
#[must_use]
pub fn with_content(mut self, content: impl Into<String>) -> Self {
self.content = Some(content.into());
self
}
#[must_use]
pub fn with_summary(mut self, summary: impl Into<String>) -> Self {
self.summary = Some(summary.into());
self
}
#[must_use]
pub fn with_score(mut self, score: f32) -> Self {
self.score = score;
self
}
#[must_use]
pub fn with_depth(mut self, depth: usize) -> Self {
self.depth = depth;
self
}
#[must_use]
pub fn with_page_range(mut self, start: usize, end: usize) -> Self {
self.page_range = Some((start, end));
self
}
}
#[derive(Debug, Clone)]
pub struct RetrieveResponse {
pub results: Vec<RetrievalResult>,
pub content: String,
pub confidence: f32,
pub is_sufficient: bool,
pub strategy_used: String,
pub complexity: QueryComplexity,
pub reasoning_chain: ReasoningChain,
pub tokens_used: usize,
}
impl Default for RetrieveResponse {
fn default() -> Self {
Self {
results: Vec::new(),
content: String::new(),
confidence: 0.0,
is_sufficient: false,
strategy_used: String::new(),
complexity: QueryComplexity::Medium,
reasoning_chain: ReasoningChain::default(),
tokens_used: 0,
}
}
}
impl RetrieveResponse {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.results.is_empty()
}
#[must_use]
pub fn len(&self) -> usize {
self.results.len()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NavigationStep {
pub node_id: String,
pub title: String,
pub score: f32,
pub decision: NavigationDecision,
pub depth: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum NavigationDecision {
GoToChild(usize),
ThisIsTheAnswer,
ExploreMore,
Skip,
BacktrackFrom(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum StageName {
Analyze,
Plan,
Search,
Evaluate,
}
impl std::fmt::Display for StageName {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Analyze => write!(f, "analyze"),
Self::Plan => write!(f, "plan"),
Self::Search => write!(f, "search"),
Self::Evaluate => write!(f, "evaluate"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmCallSummary {
pub prompt_summary: String,
pub tokens_used: usize,
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningCandidate {
pub node_id: String,
pub title: String,
pub score: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningStep {
pub stage: StageName,
pub node_id: Option<String>,
pub title: Option<String>,
pub score: f32,
pub decision: NavigationDecision,
pub depth: usize,
pub reasoning: String,
pub candidates: Vec<ReasoningCandidate>,
pub strategy_used: Option<String>,
pub llm_call: Option<LlmCallSummary>,
pub references_followed: Vec<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ReasoningChain {
pub steps: Vec<ReasoningStep>,
}
impl ReasoningChain {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn push(&mut self, step: ReasoningStep) {
self.steps.push(step);
}
#[must_use]
pub fn len(&self) -> usize {
self.steps.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.steps.is_empty()
}
#[must_use]
pub fn summary(&self) -> String {
self.steps
.iter()
.map(|s| {
let node_info = s.title.as_deref().unwrap_or("(no node)");
format!(
"[{}] {} (score={:.2}): {}",
s.stage, node_info, s.score, s.reasoning
)
})
.collect::<Vec<_>>()
.join("\n")
}
}
#[derive(Debug, Clone)]
pub struct SearchPath {
pub nodes: Vec<NodeId>,
pub score: f32,
pub leaf: Option<NodeId>,
pub step_reasons: Vec<Option<String>>,
}
impl SearchPath {
#[must_use]
pub fn new() -> Self {
Self {
nodes: Vec::new(),
score: 0.0,
leaf: None,
step_reasons: Vec::new(),
}
}
#[must_use]
pub fn from_node(node_id: NodeId, score: f32) -> Self {
Self {
nodes: vec![node_id],
score,
leaf: Some(node_id),
step_reasons: vec![None],
}
}
#[must_use]
pub fn extend(&self, node_id: NodeId, score: f32) -> Self {
let mut nodes = self.nodes.clone();
let mut step_reasons = self.step_reasons.clone();
nodes.push(node_id);
step_reasons.push(None);
Self {
nodes,
score: self.score + score,
leaf: Some(node_id),
step_reasons,
}
}
#[must_use]
pub fn extend_with_reason(
&self,
node_id: NodeId,
score: f32,
reason: impl Into<String>,
) -> Self {
let mut nodes = self.nodes.clone();
let mut step_reasons = self.step_reasons.clone();
nodes.push(node_id);
step_reasons.push(Some(reason.into()));
Self {
nodes,
score: self.score + score,
leaf: Some(node_id),
step_reasons,
}
}
}
impl Default for SearchPath {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Default)]
pub struct RetrievalStats {
pub nodes_visited: usize,
pub llm_calls: usize,
pub time_ms: u64,
pub tokens_used: usize,
pub cache_hits: usize,
pub cache_misses: usize,
}