Skip to main content

codemem_engine/
enrichment_text.rs

1use crate::index::{CodeChunk, ResolvedEdge, Symbol};
2use crate::scoring;
3use crate::CodememEngine;
4use codemem_core::{GraphBackend, MemoryType};
5
6impl CodememEngine {
7    // ── Contextual Enrichment ────────────────────────────────────────────────
8
9    /// Build contextual text for a memory node.
10    ///
11    /// NOTE: Acquires the graph lock on each call. For batch operations,
12    /// consider passing a pre-acquired guard or caching results.
13    pub fn enrich_memory_text(
14        &self,
15        content: &str,
16        memory_type: MemoryType,
17        tags: &[String],
18        namespace: Option<&str>,
19        node_id: Option<&str>,
20    ) -> String {
21        let mut ctx = String::new();
22        ctx.push_str(&format!("[{}]", memory_type));
23
24        if let Some(ns) = namespace {
25            ctx.push_str(&format!(" [namespace:{}]", ns));
26        }
27
28        if !tags.is_empty() {
29            ctx.push_str(&format!(" [tags:{}]", tags.join(",")));
30        }
31
32        if let Some(nid) = node_id {
33            let graph = match self.lock_graph() {
34                Ok(g) => g,
35                Err(_) => return format!("{ctx}\n{content}"),
36            };
37            if let Ok(edges) = graph.get_edges(nid) {
38                let mut rels: Vec<String> = Vec::new();
39                for edge in edges.iter().take(8) {
40                    let other = if edge.src == nid {
41                        &edge.dst
42                    } else {
43                        &edge.src
44                    };
45                    let label = graph
46                        .get_node(other)
47                        .ok()
48                        .flatten()
49                        .map(|n| n.label.clone())
50                        .unwrap_or_else(|| other.to_string());
51                    let dir = if edge.src == nid { "->" } else { "<-" };
52                    rels.push(format!("{dir} {} ({})", label, edge.relationship));
53                }
54                if !rels.is_empty() {
55                    ctx.push_str(&format!("\nRelated: {}", rels.join("; ")));
56                }
57            }
58        }
59
60        format!("{ctx}\n{content}")
61    }
62
63    /// Build contextual text for a code symbol.
64    pub fn enrich_symbol_text(&self, sym: &Symbol, edges: &[ResolvedEdge]) -> String {
65        let mut ctx = String::new();
66        ctx.push_str(&format!("[{} {}]", sym.visibility, sym.kind));
67        ctx.push_str(&format!(" File: {}", sym.file_path));
68
69        if let Some(ref parent) = sym.parent {
70            ctx.push_str(&format!(" Parent: {}", parent));
71        }
72
73        let related: Vec<String> = edges
74            .iter()
75            .filter(|e| {
76                e.source_qualified_name == sym.qualified_name
77                    || e.target_qualified_name == sym.qualified_name
78            })
79            .take(8)
80            .map(|e| {
81                if e.source_qualified_name == sym.qualified_name {
82                    format!("-> {} ({})", e.target_qualified_name, e.relationship)
83                } else {
84                    format!("<- {} ({})", e.source_qualified_name, e.relationship)
85                }
86            })
87            .collect();
88        if !related.is_empty() {
89            ctx.push_str(&format!("\nRelated: {}", related.join("; ")));
90        }
91
92        let mut body = format!("{}: {}", sym.qualified_name, sym.signature);
93        if let Some(ref doc) = sym.doc_comment {
94            body.push('\n');
95            body.push_str(doc);
96        }
97
98        format!("{ctx}\n{body}")
99    }
100
101    /// Build contextual text for a code chunk before embedding.
102    pub fn enrich_chunk_text(&self, chunk: &CodeChunk) -> String {
103        let mut ctx = String::new();
104        ctx.push_str(&format!("[chunk:{}]", chunk.node_kind));
105        ctx.push_str(&format!(" File: {}", chunk.file_path));
106        ctx.push_str(&format!(" Lines: {}-{}", chunk.line_start, chunk.line_end));
107        if let Some(ref parent) = chunk.parent_symbol {
108            ctx.push_str(&format!(" Parent: {}", parent));
109        }
110
111        let body = scoring::truncate_content(&chunk.text, 4000);
112
113        format!("{ctx}\n{body}")
114    }
115}