comindly_core/
brief.rs

1//! What the LLM is shown about a symbol or a flow: its code, docstring, place and neighbours.
2
3use std::collections::HashMap;
4use std::path::PathBuf;
5
6use crate::model::{Edge, EdgeKind, Symbol, SymbolKind};
7
8const CODE_LINES: usize = 60;
9const CODE_CHARS: usize = 3000;
10const CALLS: usize = 8;
11const CALLERS: usize = 6;
12
13/// The indexed checkouts, read once per file.
14pub struct Sources {
15    roots: HashMap<String, PathBuf>,
16    files: HashMap<(String, String), Option<String>>,
17}
18
19impl Sources {
20    pub fn new(roots: &[(String, PathBuf)]) -> Sources {
21        Sources {
22            roots: roots.iter().cloned().collect(),
23            files: HashMap::new(),
24        }
25    }
26
27    /// Use `text` as the contents of `path` in `repo`, e.g. read from the index rather than a checkout.
28    pub fn insert(&mut self, repo: &str, path: &str, text: Option<String>) {
29        self.files.insert((repo.into(), path.into()), text);
30    }
31
32    /// The symbol's own lines, cut to a size an LLM call can carry.
33    pub fn code(&mut self, s: &Symbol) -> Option<String> {
34        let key = (s.repo.0.clone(), s.file_path.clone());
35        let text = self
36            .files
37            .entry(key)
38            .or_insert_with(|| {
39                let root = self.roots.get(&s.repo.0)?;
40                std::fs::read_to_string(root.join(&s.file_path)).ok()
41            })
42            .as_deref()?;
43        let (start, end) = (s.range.start.line as usize, s.range.end.line as usize);
44        let lines: Vec<&str> = text
45            .lines()
46            .skip(start.saturating_sub(1))
47            .take(end.saturating_sub(start) + 1)
48            .collect();
49        if lines.iter().all(|l| l.trim().is_empty()) {
50            return None;
51        }
52        let mut code = lines[..lines.len().min(CODE_LINES)].join("\n");
53        if code.len() > CODE_CHARS {
54            let cut = (0..=CODE_CHARS)
55                .rev()
56                .find(|&i| code.is_char_boundary(i))
57                .unwrap_or(0);
58            code.truncate(cut);
59        }
60        if lines.len() > CODE_LINES || code.len() == CODE_CHARS {
61            code.push_str("\n… (cut)");
62        }
63        Some(code)
64    }
65}
66
67/// Each symbol's edges in both directions, for naming its neighbours.
68pub struct Neighbours<'a> {
69    by_id: HashMap<String, &'a Symbol>,
70    out: HashMap<String, Vec<(&'a EdgeKind, String)>>,
71    into: HashMap<String, Vec<(&'a EdgeKind, String)>>,
72}
73
74impl<'a> Neighbours<'a> {
75    pub fn new(symbols: &'a [Symbol], edges: &'a [Edge]) -> Neighbours<'a> {
76        let mut n = Neighbours {
77            by_id: symbols.iter().map(|s| (s.id.render(), s)).collect(),
78            out: HashMap::new(),
79            into: HashMap::new(),
80        };
81        for e in edges {
82            let (src, dst) = (e.src.render(), e.dst.render());
83            n.out
84                .entry(src.clone())
85                .or_default()
86                .push((&e.kind, dst.clone()));
87            n.into.entry(dst).or_default().push((&e.kind, src));
88        }
89        n
90    }
91
92    fn name(&self, id: &str, with_repo: bool) -> Option<String> {
93        let s = self.by_id.get(id)?;
94        Some(if with_repo {
95            format!("{} ({})", s.name, s.repo.0)
96        } else {
97            s.name.clone()
98        })
99    }
100
101    fn list(
102        &self,
103        edges: Option<&Vec<(&EdgeKind, String)>>,
104        kind: &EdgeKind,
105        max: usize,
106        with_repo: bool,
107    ) -> Option<String> {
108        let mut names: Vec<String> = edges?
109            .iter()
110            .filter(|(k, _)| *k == kind)
111            .filter_map(|(_, id)| self.name(id, with_repo))
112            .collect();
113        names.sort();
114        names.dedup();
115        if names.is_empty() {
116            return None;
117        }
118        let more = names.len().saturating_sub(max);
119        names.truncate(max);
120        Some(match more {
121            0 => names.join(", "),
122            n => format!("{} and {n} more", names.join(", ")),
123        })
124    }
125}
126
127/// Everything the summary prompt sees about `s`.
128pub fn symbol(s: &Symbol, n: &Neighbours, sources: &mut Sources) -> String {
129    let id = s.id.render();
130    let mut out = vec![
131        format!("Symbol: {}", s.name),
132        format!(
133            "Kind: {:?} in repository {}, {}:{}",
134            s.kind, s.repo.0, s.file_path, s.range.start.line
135        ),
136    ];
137    let owner = n.into.get(&id).and_then(|v| {
138        v.iter()
139            .filter(|(k, _)| **k == EdgeKind::Contains)
140            .filter_map(|(_, p)| n.by_id.get(p))
141            .find(|p| !matches!(p.kind, SymbolKind::File | SymbolKind::Module))
142    });
143    if let Some(o) = owner {
144        out.push(format!("Belongs to: {:?} {}", o.kind, o.name));
145    }
146    if let Some(sig) = s.signature.as_deref().filter(|x| !x.trim().is_empty()) {
147        out.push(format!("Signature: {sig}"));
148    }
149    if let Some(doc) = s.docstring.as_deref().filter(|d| !d.trim().is_empty()) {
150        out.push(format!("Docstring: {}", doc.trim()));
151    }
152    match sources.code(s) {
153        Some(code) => out.push(format!("Code:\n```\n{code}\n```")),
154        None => out.push("Code: not available".into()),
155    }
156    let outgoing = n.out.get(&id);
157    if let Some(calls) = n.list(outgoing, &EdgeKind::Calls, CALLS, false) {
158        out.push(format!("Calls: {calls}"));
159    }
160    if let Some(callers) = n.list(n.into.get(&id), &EdgeKind::Calls, CALLERS, true) {
161        out.push(format!("Called by: {callers}"));
162    }
163    let contracts: Vec<String> = [
164        EdgeKind::Serves,
165        EdgeKind::Consumes,
166        EdgeKind::Publishes,
167        EdgeKind::Subscribes,
168        EdgeKind::Reads,
169        EdgeKind::Writes,
170    ]
171    .iter()
172    .filter_map(|k| Some(format!("{k} {}", n.list(outgoing, k, CALLS, false)?)))
173    .collect();
174    if !contracts.is_empty() {
175        out.push(format!("Contracts: {}", contracts.join("; ")));
176    }
177    out.join("\n")
178}
179
180/// A flow's trace followed by its entry point's code, when there is some.
181pub fn flow(trace: String, entry: &Symbol, sources: &mut Sources) -> String {
182    match sources.code(entry) {
183        Some(code) => format!(
184            "{trace}\n\nEntry point code ({}):\n```\n{code}\n```",
185            entry.repo.0
186        ),
187        None => trace,
188    }
189}
190
191/// The trace a walkthrough is written from: each step with its one-line summary when known.
192pub fn trace(steps: &[crate::graph::FlowStep], summaries: &HashMap<String, String>) -> String {
193    steps
194        .iter()
195        .map(|h| {
196            let line = format!(
197                "d{} {} {}:{} {}",
198                h.depth, h.via, h.node.repo, h.node.name, h.node.location
199            );
200            match summaries.get(&h.node.id) {
201                Some(s) => format!("{line} — {s}"),
202                None => line,
203            }
204        })
205        .collect::<Vec<_>>()
206        .join("\n")
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use crate::model::{GlobalSymbolId, Language, Position, Range, RepoId};
213
214    fn sym(name: &str, kind: SymbolKind, lines: (u32, u32)) -> Symbol {
215        Symbol {
216            id: GlobalSymbolId {
217                scheme: "t".into(),
218                package_manager: ".".into(),
219                package: "shop".into(),
220                version: ".".into(),
221                descriptor: name.into(),
222            },
223            name: name.into(),
224            kind,
225            language: Language::Python,
226            repo: RepoId("shop".into()),
227            file_path: "shop/api.py".into(),
228            range: Range {
229                start: Position {
230                    line: lines.0,
231                    column: 0,
232                },
233                end: Position {
234                    line: lines.1,
235                    column: 0,
236                },
237            },
238            signature: Some(format!("def {name}()")),
239            docstring: None,
240            fingerprint: None,
241        }
242    }
243
244    fn edge(a: &Symbol, b: &Symbol, kind: EdgeKind) -> Edge {
245        Edge {
246            src: a.id.clone(),
247            dst: b.id.clone(),
248            kind,
249            confidence: 1.0,
250            cross_repo: false,
251        }
252    }
253
254    #[test]
255    fn a_symbol_brief_has_its_code_and_neighbours() {
256        let dir = std::env::temp_dir().join(format!("comindly-brief-{}", std::process::id()));
257        std::fs::create_dir_all(dir.join("shop")).unwrap();
258        std::fs::write(
259            dir.join("shop/api.py"),
260            "import db\n\ndef create():\n    db.save()\n\ndef handler():\n    create()\n",
261        )
262        .unwrap();
263        let create = sym("create", SymbolKind::Function, (3, 4));
264        let save = sym("save", SymbolKind::Function, (1, 1));
265        let handler = sym("handler", SymbolKind::Function, (6, 7));
266        let mut route = sym("http:POST /orders", SymbolKind::Contract, (1, 1));
267        route.signature = None;
268        let symbols = vec![create.clone(), save.clone(), handler.clone(), route.clone()];
269        let edges = vec![
270            edge(&create, &save, EdgeKind::Calls),
271            edge(&handler, &create, EdgeKind::Calls),
272            edge(&create, &route, EdgeKind::Serves),
273        ];
274        let n = Neighbours::new(&symbols, &edges);
275        let mut sources = Sources::new(&[("shop".into(), dir.clone())]);
276        let brief = symbol(&create, &n, &mut sources);
277        std::fs::remove_dir_all(&dir).ok();
278        assert!(brief.contains("def create():\n    db.save()"), "{brief}");
279        assert!(
280            !brief.contains("def handler"),
281            "only its own lines: {brief}"
282        );
283        assert!(brief.contains("Calls: save"), "{brief}");
284        assert!(brief.contains("Called by: handler (shop)"), "{brief}");
285        assert!(
286            brief.contains("Contracts: serves http:POST /orders"),
287            "{brief}"
288        );
289    }
290}
comindly_core/brief.rs

comindly_core/
brief.rs
