semtree_rag/
context.rs

1use std::sync::Arc;
2
3use semtree_store::Hit;
4use serde::{Deserialize, Serialize};
5
6use crate::{ChunkRegistry, RagError, SearchEngine};
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ContextSnippet {
10    pub chunk_id: String,
11    pub score: f32,
12    pub path: String,
13    pub name: Option<String>,
14    /// 1-based start line in the source file.
15    pub start_line: usize,
16    /// Raw source text of the chunk.
17    pub content: String,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct ContextWindow {
22    pub query: String,
23    pub snippets: Vec<ContextSnippet>,
24    pub prompt: String,
25}
26
27impl ContextWindow {
28    /// Assemble a window from hits that have already been ranked, resolving each
29    /// against `registry` so the prompt carries real code and not just chunk ids.
30    ///
31    /// Taking hits rather than running the search means a caller can feed this
32    /// any ranking - vector, BM25, or the two fused - instead of being tied to
33    /// plain vector similarity.
34    pub fn from_hits(query: &str, hits: &[Hit], registry: &ChunkRegistry) -> Self {
35        let snippets: Vec<ContextSnippet> = hits
36            .iter()
37            .filter_map(|h| {
38                registry.get(&h.id).map(|c| ContextSnippet {
39                    chunk_id: h.id.clone(),
40                    score: h.score,
41                    path: c.path.display().to_string(),
42                    name: c.name.clone(),
43                    start_line: c.span.start_line + 1,
44                    content: c.content.clone(),
45                })
46            })
47            .collect();
48
49        let context_block = snippets
50            .iter()
51            .enumerate()
52            .map(|(i, s)| {
53                let header = match &s.name {
54                    Some(name) => format!("[{}] {}:{} - {name}", i + 1, s.path, s.start_line),
55                    None => format!("[{}] {}:{}", i + 1, s.path, s.start_line),
56                };
57                format!("{header}\n```\n{}\n```", s.content)
58            })
59            .collect::<Vec<_>>()
60            .join("\n\n");
61
62        let prompt = format!(
63            "Use the following code context to answer the question.\n\n{context_block}\n\nQuestion: {query}"
64        );
65
66        Self {
67            query: query.to_string(),
68            snippets,
69            prompt,
70        }
71    }
72}
73
74pub struct ContextBuilder {
75    engine: Arc<SearchEngine>,
76    max_chunks: usize,
77}
78
79impl ContextBuilder {
80    pub fn new(engine: Arc<SearchEngine>) -> Self {
81        Self {
82            engine,
83            max_chunks: 5,
84        }
85    }
86
87    pub fn with_max_chunks(mut self, n: usize) -> Self {
88        self.max_chunks = n;
89        self
90    }
91
92    /// Builds a context window for `query` from a plain vector search.
93    pub async fn build(
94        &self,
95        query: &str,
96        registry: &ChunkRegistry,
97    ) -> Result<ContextWindow, RagError> {
98        let hits = self.engine.search(query, self.max_chunks).await?;
99        Ok(ContextWindow::from_hits(query, &hits, registry))
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use semtree_core::{Chunk, ChunkKind, Language, Span};
106
107    use super::*;
108
109    #[test]
110    fn window_carries_code_and_locations_into_the_prompt() {
111        let mut registry = ChunkRegistry::default();
112        registry.insert(Chunk {
113            id: "c1".into(),
114            path: "src/auth.rs".into(),
115            language: Language::Rust,
116            kind: ChunkKind::Function,
117            name: Some("verify_token".into()),
118            content: "fn verify_token() {}".into(),
119            span: Span::new(0, 20, 41, 43),
120            doc: None,
121        });
122
123        let hits = vec![
124            Hit {
125                id: "c1".into(),
126                score: 0.9,
127            },
128            // A hit with no chunk behind it contributes nothing rather than
129            // leaving an empty code fence in the prompt.
130            Hit {
131                id: "gone".into(),
132                score: 0.5,
133            },
134        ];
135
136        let window = ContextWindow::from_hits("how are tokens checked", &hits, &registry);
137
138        assert_eq!(window.snippets.len(), 1);
139        assert_eq!(
140            window.snippets[0].start_line, 42,
141            "line numbers are 1-based"
142        );
143        assert!(window.prompt.contains("src/auth.rs:42 - verify_token"));
144        assert!(window.prompt.contains("fn verify_token() {}"));
145        assert!(window.prompt.contains("how are tokens checked"));
146    }
147}
semtree_rag/context.rs

semtree_rag/
context.rs
