spire_ai/code/
context.rs

1//! Code context assembly for LLM consumption.
2
3use crate::code::CodeIndex;
4use crate::code::symbols::CodeChunk;
5use crate::error::Result;
6
7/// Assembled code context ready for LLM consumption.
8#[derive(Debug, Clone)]
9pub struct CodeContext {
10    /// Formatted context string
11    pub text: String,
12    /// Source chunks used
13    pub chunks: Vec<CodeChunk>,
14    /// Approximate token count
15    pub tokens: usize,
16}
17
18impl CodeContext {
19    /// Format as a system prompt.
20    pub fn as_system_prompt(&self) -> String {
21        format!(
22            "You have access to the following code context:\n\n{}",
23            self.text
24        )
25    }
26
27    /// Format as user context for an LLM message.
28    pub fn as_user_context(&self) -> String {
29        format!("Relevant code:\n\n{}\n\nQuestion: ", self.text)
30    }
31}
32
33/// Builder for assembling code context.
34pub struct ContextBuilder<'a> {
35    code_index: &'a CodeIndex,
36    question: String,
37    max_tokens: usize,
38    max_chunks: usize,
39}
40
41impl<'a> ContextBuilder<'a> {
42    pub(crate) fn new(code_index: &'a CodeIndex, question: String) -> Self {
43        Self {
44            code_index,
45            question,
46            max_tokens: 4000,
47            max_chunks: 10,
48        }
49    }
50
51    /// Set the maximum approximate token count (default: 4000).
52    pub fn max_tokens(mut self, n: usize) -> Self {
53        self.max_tokens = n;
54        self
55    }
56
57    /// Set the maximum number of chunks to include (default: 10).
58    pub fn max_chunks(mut self, n: usize) -> Self {
59        self.max_chunks = n;
60        self
61    }
62
63    /// Build the context.
64    pub async fn build(self) -> Result<CodeContext> {
65        let hits = self
66            .code_index
67            .collection
68            .search(&self.question)
69            .limit(self.max_chunks)
70            .run()
71            .await?;
72
73        let mut chunks = Vec::new();
74        let mut text = String::new();
75        let mut approx_tokens = 0;
76
77        for hit in hits {
78            let chunk = hit.doc;
79            // Rough token estimate: 1 token ~= 4 chars
80            let chunk_tokens = chunk.code.len() / 4;
81
82            if approx_tokens + chunk_tokens > self.max_tokens {
83                break;
84            }
85
86            // Format the chunk
87            text.push_str(&format!(
88                "// File: {} (lines {}-{})\n",
89                chunk.file, chunk.start_line, chunk.end_line
90            ));
91            if let Some(ref name) = chunk.name {
92                text.push_str(&format!("// Symbol: {} ({:?})\n", name, chunk.kind));
93            }
94            text.push_str(&chunk.code);
95            text.push_str("\n\n");
96
97            approx_tokens += chunk_tokens;
98            chunks.push(chunk);
99        }
100
101        Ok(CodeContext {
102            text,
103            chunks,
104            tokens: approx_tokens,
105        })
106    }
107}
spire_ai/code/context.rs

spire_ai/code/
context.rs
