lean_ctx/core/knowledge_router/
planner.rs

1use super::{
2    reference_resolver::{ReferenceType, ResolvedReference},
3    source_manifest::SourceManifestEntry,
4};
5#[derive(Debug, Clone, Default, PartialEq, Eq)]
6pub struct ContextCandidate {
7    pub candidate_id: String,
8    pub source_id: String,
9    pub kind: String,
10    pub relevance_milli: u16,
11    pub estimated_tokens: u64,
12    pub reference: Option<String>,
13}
14#[derive(Debug, Clone, Copy, Default)]
15pub struct QueryPlanner;
16impl QueryPlanner {
17    pub fn plan(
18        references: &[ResolvedReference],
19        manifests: &[SourceManifestEntry],
20        budget_tokens: u64,
21    ) -> Vec<ContextCandidate> {
22        let mut candidates = references
23            .iter()
24            .enumerate()
25            .filter_map(|(index, reference)| {
26                manifests
27                    .iter()
28                    .find(|manifest| manifest.source_id == reference.source_id)
29                    .map(|manifest| ContextCandidate {
30                        candidate_id: format!("candidate-{index}-{}", manifest.source_id),
31                        source_id: manifest.source_id.clone(),
32                        kind: kind_for(reference.ref_type).into(),
33                        relevance_milli: reference.confidence_milli,
34                        estimated_tokens: estimated_tokens(reference.ref_type),
35                        reference: Some(reference.identifier.clone()),
36                    })
37            })
38            .collect::<Vec<_>>();
39        candidates.sort_by(|left, right| {
40            right
41                .relevance_milli
42                .cmp(&left.relevance_milli)
43                .then_with(|| left.candidate_id.cmp(&right.candidate_id))
44        });
45        let mut used: u64 = 0;
46        candidates
47            .into_iter()
48            .filter(|candidate| {
49                if used.saturating_add(candidate.estimated_tokens) > budget_tokens {
50                    return false;
51                }
52                used += candidate.estimated_tokens;
53                true
54            })
55            .collect()
56    }
57}
58fn kind_for(reference: ReferenceType) -> &'static str {
59    match reference {
60        ReferenceType::JiraIssue | ReferenceType::GitHubIssue => "issue",
61        ReferenceType::GitHubPR => "pull_request",
62        ReferenceType::FilePath => "file",
63        ReferenceType::Function => "function",
64        ReferenceType::Url => "url",
65    }
66}
67fn estimated_tokens(reference: ReferenceType) -> u64 {
68    match reference {
69        ReferenceType::FilePath => 256,
70        _ => 384,
71    }
72}
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::core::knowledge_router::{
77        reference_resolver::{ReferenceType, ResolvedReference},
78        source_manifest::builtin_manifests,
79    };
80    fn reference(source: &str, confidence: u16) -> ResolvedReference {
81        ResolvedReference {
82            ref_type: ReferenceType::JiraIssue,
83            identifier: "LEAN-1".into(),
84            source_id: source.into(),
85            confidence_milli: confidence,
86        }
87    }
88    #[test]
89    fn plans_matching_sources() {
90        assert_eq!(
91            QueryPlanner::plan(&[reference("jira", 900)], &builtin_manifests(), 400).len(),
92            1
93        );
94    }
95    #[test]
96    fn respects_budget() {
97        {
98            let found = &QueryPlanner::plan(&[reference("jira", 900)], &builtin_manifests(), 383);
99            assert!(found.is_empty(), "{found:?}");
100        }
101    }
102    #[test]
103    fn sorts_by_relevance() {
104        let candidates = QueryPlanner::plan(
105            &[reference("jira", 800), reference("github", 950)],
106            &builtin_manifests(),
107            1_000,
108        );
109        assert_eq!(candidates[0].source_id, "github");
110    }
111    #[test]
112    fn skips_unknown_sources() {
113        {
114            let found = &QueryPlanner::plan(&[reference("other", 900)], &builtin_manifests(), 400);
115            assert!(found.is_empty(), "{found:?}");
116        }
117    }
118}
lean_ctx/core/knowledge_router/planner.rs

lean_ctx/core/knowledge_router/
planner.rs
