spire_ai/code/
mod.rs

1//! Code indexing and semantic search.
2//!
3//! Index source code repositories and search them semantically.
4//! Uses tree-sitter for language-aware parsing when the `code` feature is enabled.
5
6mod context;
7#[cfg(feature = "code")]
8mod parser;
9mod symbols;
10#[cfg(feature = "code")]
11mod walker;
12
13pub use context::{CodeContext, ContextBuilder};
14pub use symbols::{CodeChunk, CodeHit, SymbolKind};
15
16use crate::client::Spire;
17use crate::collection::Collection;
18use crate::error::Result;
19use crate::types::IndexResult;
20
21/// Index and search a codebase.
22#[derive(Clone)]
23#[allow(dead_code)] // spire, name used by index_dir/index_file/search methods
24pub struct CodeIndex {
25    spire: Spire,
26    name: String,
27    pub(crate) collection: Collection<CodeChunk>,
28}
29
30impl CodeIndex {
31    pub(crate) fn new(spire: Spire, name: String) -> Self {
32        let collection_name = format!("code_{name}");
33        let collection = spire.collection::<CodeChunk>(&collection_name);
34        Self {
35            spire,
36            name,
37            collection,
38        }
39    }
40
41    /// Ensure the backing collection exists.
42    pub async fn ensure(&self) -> Result<()> {
43        self.collection.ensure().await
44    }
45
46    /// Index a directory (respects .gitignore).
47    #[cfg(feature = "code")]
48    pub async fn index_dir(&self, path: &str) -> Result<IndexResult> {
49        let files = walker::walk_dir(path)?;
50        let mut total_chunks = 0;
51        let mut total_symbols = 0;
52        let mut errors = Vec::new();
53
54        for file_path in &files {
55            match self.index_file_internal(file_path).await {
56                Ok((chunks, symbols)) => {
57                    total_chunks += chunks;
58                    total_symbols += symbols;
59                }
60                Err(e) => {
61                    errors.push(format!("{file_path}: {e}"));
62                }
63            }
64        }
65
66        Ok(IndexResult {
67            files: files.len(),
68            chunks: total_chunks,
69            symbols: total_symbols,
70            errors,
71        })
72    }
73
74    /// Index a single file.
75    #[cfg(feature = "code")]
76    pub async fn index_file(&self, path: &str) -> Result<IndexResult> {
77        let (chunks, symbols) = self.index_file_internal(path).await?;
78        Ok(IndexResult {
79            files: 1,
80            chunks,
81            symbols,
82            errors: Vec::new(),
83        })
84    }
85
86    #[cfg(feature = "code")]
87    async fn index_file_internal(&self, path: &str) -> Result<(usize, usize)> {
88        let content = tokio::fs::read_to_string(path).await?;
89        let language = parser::detect_language(path);
90        let chunks = parser::parse_file(path, &content, &language);
91        let symbols = chunks.iter().filter(|c| c.name.is_some()).count();
92        let count = chunks.len();
93        self.collection.insert_many(&chunks).await?;
94        Ok((count, symbols))
95    }
96
97    /// Semantic code search.
98    pub async fn search(&self, query: &str) -> Result<Vec<CodeHit>> {
99        let hits = self.collection.search(query).limit(10).run().await?;
100        Ok(hits
101            .into_iter()
102            .map(|h| CodeHit {
103                chunk: h.doc,
104                score: h.score,
105            })
106            .collect())
107    }
108
109    /// Find symbol by name (SQL filter).
110    pub async fn find_symbol(&self, name: &str) -> Result<Vec<CodeChunk>> {
111        // Use semantic search with the symbol name as query
112        let hits = self.collection.search(name).limit(20).docs().await?;
113        Ok(hits
114            .into_iter()
115            .filter(|c| c.name.as_ref().is_some_and(|n| n.contains(name)))
116            .collect())
117    }
118
119    /// Get relevant context for an LLM question.
120    pub async fn context_for(&self, question: &str) -> Result<CodeContext> {
121        self.context(question).build().await
122    }
123
124    /// Build a context query with options.
125    pub fn context(&self, question: &str) -> ContextBuilder<'_> {
126        ContextBuilder::new(self, question.to_string())
127    }
128}
spire_ai/code/mod.rs

spire_ai/code/
mod.rs
