semtree_rag/
context.rs

1use std::sync::Arc;
2
3use serde::{Deserialize, Serialize};
4
5use crate::{ChunkRegistry, RagError, SearchEngine};
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ContextSnippet {
9    pub chunk_id: String,
10    pub score: f32,
11    pub path: String,
12    pub name: Option<String>,
13    /// 1-based start line in the source file.
14    pub start_line: usize,
15    /// Raw source text of the chunk.
16    pub content: String,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct ContextWindow {
21    pub query: String,
22    pub snippets: Vec<ContextSnippet>,
23    pub prompt: String,
24}
25
26pub struct ContextBuilder {
27    engine: Arc<SearchEngine>,
28    max_chunks: usize,
29}
30
31impl ContextBuilder {
32    pub fn new(engine: Arc<SearchEngine>) -> Self {
33        Self {
34            engine,
35            max_chunks: 5,
36        }
37    }
38
39    pub fn with_max_chunks(mut self, n: usize) -> Self {
40        self.max_chunks = n;
41        self
42    }
43
44    /// Builds a context window for `query`, resolving each hit against
45    /// `registry` so the prompt carries real code (not just chunk ids).
46    pub async fn build(
47        &self,
48        query: &str,
49        registry: &ChunkRegistry,
50    ) -> Result<ContextWindow, RagError> {
51        let hits = self.engine.search(query, self.max_chunks).await?;
52
53        let snippets: Vec<ContextSnippet> = hits
54            .iter()
55            .filter_map(|h| {
56                registry.get(&h.id).map(|c| ContextSnippet {
57                    chunk_id: h.id.clone(),
58                    score: h.score,
59                    path: c.path.display().to_string(),
60                    name: c.name.clone(),
61                    start_line: c.span.start_line + 1,
62                    content: c.content.clone(),
63                })
64            })
65            .collect();
66
67        let context_block = snippets
68            .iter()
69            .enumerate()
70            .map(|(i, s)| {
71                let header = match &s.name {
72                    Some(name) => format!("[{}] {}:{} - {name}", i + 1, s.path, s.start_line),
73                    None => format!("[{}] {}:{}", i + 1, s.path, s.start_line),
74                };
75                format!("{header}\n```\n{}\n```", s.content)
76            })
77            .collect::<Vec<_>>()
78            .join("\n\n");
79
80        let prompt = format!(
81            "Use the following code context to answer the question.\n\n{context_block}\n\nQuestion: {query}"
82        );
83
84        Ok(ContextWindow {
85            query: query.to_string(),
86            snippets,
87            prompt,
88        })
89    }
90}
semtree_rag/context.rs

semtree_rag/
context.rs
