semtree-rag 0.5.0

use std::sync::Arc;

use semtree_store::Hit;
use serde::{Deserialize, Serialize};

use crate::{ChunkRegistry, RagError, SearchEngine};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextSnippet {
    pub chunk_id: String,
    pub score: f32,
    pub path: String,
    pub name: Option<String>,
    /// 1-based start line in the source file.
    pub start_line: usize,
    /// Raw source text of the chunk.
    pub content: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextWindow {
    pub query: String,
    pub snippets: Vec<ContextSnippet>,
    pub prompt: String,
}

impl ContextWindow {
    /// Assemble a window from hits that have already been ranked, resolving each
    /// against `registry` so the prompt carries real code and not just chunk ids.
    ///
    /// Taking hits rather than running the search means a caller can feed this
    /// any ranking - vector, BM25, or the two fused - instead of being tied to
    /// plain vector similarity.
    pub fn from_hits(query: &str, hits: &[Hit], registry: &ChunkRegistry) -> Self {
        let snippets: Vec<ContextSnippet> = hits
            .iter()
            .filter_map(|h| {
                registry.get(&h.id).map(|c| ContextSnippet {
                    chunk_id: h.id.clone(),
                    score: h.score,
                    path: c.path.display().to_string(),
                    name: c.name.clone(),
                    start_line: c.span.start_line + 1,
                    content: c.content.clone(),
                })
            })
            .collect();

        let context_block = snippets
            .iter()
            .enumerate()
            .map(|(i, s)| {
                let header = match &s.name {
                    Some(name) => format!("[{}] {}:{} - {name}", i + 1, s.path, s.start_line),
                    None => format!("[{}] {}:{}", i + 1, s.path, s.start_line),
                };
                format!("{header}\n```\n{}\n```", s.content)
            })
            .collect::<Vec<_>>()
            .join("\n\n");

        let prompt = format!(
            "Use the following code context to answer the question.\n\n{context_block}\n\nQuestion: {query}"
        );

        Self {
            query: query.to_string(),
            snippets,
            prompt,
        }
    }
}

pub struct ContextBuilder {
    engine: Arc<SearchEngine>,
    max_chunks: usize,
}

impl ContextBuilder {
    pub fn new(engine: Arc<SearchEngine>) -> Self {
        Self {
            engine,
            max_chunks: 5,
        }
    }

    pub fn with_max_chunks(mut self, n: usize) -> Self {
        self.max_chunks = n;
        self
    }

    /// Builds a context window for `query` from a plain vector search.
    pub async fn build(
        &self,
        query: &str,
        registry: &ChunkRegistry,
    ) -> Result<ContextWindow, RagError> {
        let hits = self.engine.search(query, self.max_chunks).await?;
        Ok(ContextWindow::from_hits(query, &hits, registry))
    }
}

#[cfg(test)]
mod tests {
    use semtree_core::{Chunk, ChunkKind, Language, Span};

    use super::*;

    #[test]
    fn window_carries_code_and_locations_into_the_prompt() {
        let mut registry = ChunkRegistry::default();
        registry.insert(Chunk {
            id: "c1".into(),
            path: "src/auth.rs".into(),
            language: Language::Rust,
            kind: ChunkKind::Function,
            name: Some("verify_token".into()),
            content: "fn verify_token() {}".into(),
            span: Span::new(0, 20, 41, 43),
            doc: None,
        });

        let hits = vec![
            Hit {
                id: "c1".into(),
                score: 0.9,
            },
            // A hit with no chunk behind it contributes nothing rather than
            // leaving an empty code fence in the prompt.
            Hit {
                id: "gone".into(),
                score: 0.5,
            },
        ];

        let window = ContextWindow::from_hits("how are tokens checked", &hits, &registry);

        assert_eq!(window.snippets.len(), 1);
        assert_eq!(
            window.snippets[0].start_line, 42,
            "line numbers are 1-based"
        );
        assert!(window.prompt.contains("src/auth.rs:42 - verify_token"));
        assert!(window.prompt.contains("fn verify_token() {}"));
        assert!(window.prompt.contains("how are tokens checked"));
    }
}