use async_trait::async_trait;
use super::super::RetrievalContext;
use super::super::types::{NavigationDecision, QueryComplexity};
use crate::document::{DocumentTree, NodeId};
#[derive(Debug, Clone)]
pub struct NodeEvaluation {
pub score: f32,
pub decision: NavigationDecision,
pub reasoning: Option<String>,
}
impl Default for NodeEvaluation {
fn default() -> Self {
Self {
score: 0.5,
decision: NavigationDecision::ExploreMore,
reasoning: None,
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct StrategyCapabilities {
pub uses_llm: bool,
pub uses_embeddings: bool,
pub supports_sufficiency: bool,
pub typical_latency_ms: u64,
}
#[async_trait]
pub trait RetrievalStrategy: Send + Sync {
async fn evaluate_node(
&self,
tree: &DocumentTree,
node_id: NodeId,
context: &RetrievalContext,
) -> NodeEvaluation;
async fn evaluate_nodes(
&self,
tree: &DocumentTree,
node_ids: &[NodeId],
context: &RetrievalContext,
) -> Vec<NodeEvaluation> {
let mut results = Vec::with_capacity(node_ids.len());
for node_id in node_ids {
results.push(self.evaluate_node(tree, *node_id, context).await);
}
results
}
fn name(&self) -> &str;
fn capabilities(&self) -> StrategyCapabilities;
fn suitable_for_complexity(&self, complexity: QueryComplexity) -> bool;
fn estimate_cost(&self, node_count: usize) -> StrategyCost {
StrategyCost {
llm_calls: if self.capabilities().uses_llm {
node_count
} else {
0
},
tokens: if self.capabilities().uses_llm {
node_count * 200
} else {
0
},
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct StrategyCost {
pub llm_calls: usize,
pub tokens: usize,
}