use crate::db::models::CodeElement;
use crate::graph::GraphEngine;
use serde::{Deserialize, Serialize};
const DEFAULT_MAX_TOKENS: usize = 4000;
const CHARS_PER_TOKEN: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum ContextPriority {
Contained = 1,
Imported = 2,
RecentlyChanged = 3,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextElement {
pub element: CodeElement,
pub priority: ContextPriority,
pub token_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextResult {
pub elements: Vec<ContextElement>,
pub total_tokens: usize,
pub max_tokens: usize,
pub truncated: bool,
}
impl ContextResult {
pub fn to_prompt(&self) -> String {
let mut prompt = String::new();
prompt.push_str("# Code Context\n\n");
for ctx_elem in &self.elements {
let elem = &ctx_elem.element;
prompt.push_str(&format!(
"## {}: {}\nFile: {}:{}:{}\n",
elem.element_type,
elem.qualified_name,
elem.file_path,
elem.line_start,
elem.line_end
));
if let Some(parent) = &elem.parent_qualified {
prompt.push_str(&format!("Parent: {}\n", parent));
}
prompt.push('\n');
}
if self.truncated {
prompt.push_str(&format!(
"\n*Context truncated: {} tokens total (max: {})*\n",
self.total_tokens, self.max_tokens
));
}
prompt
}
}
pub struct ContextProvider<'a> {
graph: &'a GraphEngine,
max_tokens: usize,
}
impl<'a> ContextProvider<'a> {
pub fn new(graph: &'a GraphEngine) -> Self {
Self {
graph,
max_tokens: DEFAULT_MAX_TOKENS,
}
}
pub fn with_max_tokens(graph: &'a GraphEngine, max_tokens: usize) -> Self {
Self { graph, max_tokens }
}
pub fn estimate_tokens(text: &str) -> usize {
text.len().div_ceil(CHARS_PER_TOKEN)
}
pub fn element_tokens(element: &CodeElement) -> usize {
let base = format!(
"{} {} {} {}",
element.element_type, element.name, element.qualified_name, element.file_path
);
let metadata_len = serde_json::to_string(&element.metadata)
.map(|s| s.len())
.unwrap_or(0);
Self::estimate_tokens(&base) + metadata_len / CHARS_PER_TOKEN
}
pub fn get_context_for_file(
&self,
file_path: &str,
) -> Result<ContextResult, Box<dyn std::error::Error>> {
let file_element = self.graph.find_element(file_path)?;
let mut context_elements = Vec::new();
if let Some(file) = file_element {
let priority = self.determine_priority(&file);
context_elements.push(ContextElement {
element: file.clone(),
priority,
token_count: Self::element_tokens(&file),
});
}
let relationships = self.graph.get_relationships(file_path)?;
for rel in relationships {
if let Some(element) = self.graph.find_element(&rel.target_qualified)? {
let priority = match rel.rel_type.as_str() {
"imports" => ContextPriority::Imported,
"contains" | "defines" => ContextPriority::Contained,
_ => ContextPriority::Contained,
};
let token_count = Self::element_tokens(&element);
context_elements.push(ContextElement {
element,
priority,
token_count,
});
}
}
let children = self.get_child_elements(file_path)?;
for child in children {
let token_count = Self::element_tokens(&child);
context_elements.push(ContextElement {
element: child,
priority: ContextPriority::Contained,
token_count,
});
}
context_elements.sort_by(|a, b| {
b.priority
.cmp(&a.priority)
.then_with(|| b.token_count.cmp(&a.token_count))
});
let mut total_tokens = 0;
let mut selected_elements = Vec::new();
let mut truncated = false;
for elem in context_elements {
if total_tokens + elem.token_count <= self.max_tokens {
total_tokens += elem.token_count;
selected_elements.push(elem);
} else {
truncated = true;
break;
}
}
Ok(ContextResult {
elements: selected_elements,
total_tokens,
max_tokens: self.max_tokens,
truncated,
})
}
fn get_child_elements(
&self,
parent_qualified: &str,
) -> Result<Vec<CodeElement>, Box<dyn std::error::Error>> {
let result = self.graph.get_children(parent_qualified)?;
Ok(result)
}
fn determine_priority(&self, element: &CodeElement) -> ContextPriority {
if let Some(changed) = element
.metadata
.get("recently_changed")
.and_then(|v| v.as_bool())
{
if changed {
return ContextPriority::RecentlyChanged;
}
}
if element.metadata.get("tested_by").is_some() {
return ContextPriority::RecentlyChanged;
}
ContextPriority::Contained
}
}
impl GraphEngine {
pub fn get_context(
&self,
file_path: &str,
max_tokens: usize,
) -> Result<ContextResult, Box<dyn std::error::Error>> {
let provider = ContextProvider::with_max_tokens(self, max_tokens);
provider.get_context_for_file(file_path)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_estimate_tokens() {
assert_eq!(ContextProvider::estimate_tokens("hello"), 2);
assert_eq!(ContextProvider::estimate_tokens("hello world"), 3);
assert_eq!(ContextProvider::estimate_tokens(""), 0);
assert_eq!(ContextProvider::estimate_tokens("0123456789"), 3);
}
#[test]
fn test_priority_ordering() {
assert!(ContextPriority::RecentlyChanged > ContextPriority::Imported);
assert!(ContextPriority::Imported > ContextPriority::Contained);
assert!(ContextPriority::RecentlyChanged > ContextPriority::Contained);
}
#[test]
fn test_element_tokens_calculation() {
let elem = CodeElement {
qualified_name: "test.rs::foo".to_string(),
element_type: "function".to_string(),
name: "foo".to_string(),
file_path: "test.rs".to_string(),
line_start: 1,
line_end: 10,
language: "rust".to_string(),
parent_qualified: None,
metadata: serde_json::json!({}),
};
let tokens = ContextProvider::element_tokens(&elem);
assert!(tokens > 0, "Should calculate some tokens");
}
#[test]
fn test_context_result_to_prompt_empty() {
let result = ContextResult {
elements: vec![],
total_tokens: 0,
max_tokens: 4000,
truncated: false,
};
let prompt = result.to_prompt();
assert!(prompt.contains("# Code Context"));
}
#[test]
fn test_context_result_to_prompt_with_elements() {
let elem = CodeElement {
qualified_name: "test.rs::main".to_string(),
element_type: "function".to_string(),
name: "main".to_string(),
file_path: "test.rs".to_string(),
line_start: 1,
line_end: 5,
language: "rust".to_string(),
parent_qualified: None,
metadata: serde_json::json!({}),
};
let ctx_elem = ContextElement {
element: elem,
priority: ContextPriority::RecentlyChanged,
token_count: 10,
};
let result = ContextResult {
elements: vec![ctx_elem],
total_tokens: 10,
max_tokens: 4000,
truncated: false,
};
let prompt = result.to_prompt();
assert!(prompt.contains("function: test.rs::main"));
assert!(prompt.contains("File: test.rs:1:5"));
assert!(!prompt.contains("Context truncated"));
}
#[test]
fn test_context_result_to_prompt_truncated() {
let elem = CodeElement {
qualified_name: "test.rs::foo".to_string(),
element_type: "function".to_string(),
name: "foo".to_string(),
file_path: "test.rs".to_string(),
line_start: 1,
line_end: 100,
language: "rust".to_string(),
parent_qualified: None,
metadata: serde_json::json!({}),
};
let ctx_elem = ContextElement {
element: elem,
priority: ContextPriority::Imported,
token_count: 100,
};
let result = ContextResult {
elements: vec![ctx_elem],
total_tokens: 100,
max_tokens: 50,
truncated: true,
};
let prompt = result.to_prompt();
assert!(prompt.contains("Context truncated"));
assert!(prompt.contains("100 tokens total"));
assert!(prompt.contains("max: 50"));
}
}