comindly_core/
llm.rs

1//! Comindly LLM enrichment — per-symbol summaries, symbol-aware query generation, and a
2//! repo style guide, via **Rig** (provider-agnostic: OpenAI by default; any OpenAI-compatible
3//! endpoint via `COMINDLY_LLM_BASE_URL`; Bedrock/Vertex/others can be added behind cargo features).
4//!
5//! **Opt-in / data egress:** these functions send code (signatures, snippets) to the configured
6//! LLM provider. They run only when the caller explicitly enables enrichment
7//! (`comindly link --enrich`), never as part of plain indexing. Requires `OPENAI_API_KEY` (or a
8//! provider key + `COMINDLY_LLM_BASE_URL`).
9
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::Arc;
12
13use anyhow::{Context, Result};
14use rig::completion::CompletionRequest;
15use rig::providers::openai;
16
17/// Default model — cheap and fast, suitable for bulk symbol summaries.
18pub const DEFAULT_MODEL: &str = "gpt-6-luna";
19
20/// Where requests go without an endpoint setting: OpenAI.
21pub const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
22
23pub struct LlmClient {
24    // Chat Completions API (works with OpenAI *and* any OpenAI-compatible endpoint), not the
25    // Responses API — proxies/Ollama/vLLM implement `/chat/completions`, not `/responses`.
26    client: openai::OpenAI,
27    model: String,
28    // OpenAI's own API takes `max_completion_tokens` (required by newer models); compatible
29    // servers (Ollama, vLLM, LiteLLM) expect `max_tokens`.
30    openai_native: bool,
31    // Cumulative token usage across all calls (for live progress + a final total).
32    input_tokens: Arc<AtomicU64>,
33    output_tokens: Arc<AtomicU64>,
34}
35
36/// LLM-generated enrichment for one symbol.
37#[derive(Debug, Clone)]
38pub struct Enrichment {
39    pub summary: String,
40    pub queries: Vec<String>,
41}
42
43impl LlmClient {
44    /// A client for `model` at `base_url` (OpenAI when absent), with `key`.
45    pub fn new(base_url: Option<&str>, model: &str, key: Option<&str>) -> Result<Self> {
46        let key = key.map(str::to_string);
47        let client = match base_url {
48            Some(base) => openai::OpenAIConfig::new(key.unwrap_or_else(|| "sk-noauth".into()))
49                .with_base_url(base)
50                .client(),
51            None => openai::OpenAIConfig::new(key.context("no LLM key")?).client(),
52        };
53        Ok(Self {
54            client,
55            model: model.to_string(),
56            openai_native: base_url.is_none(),
57            input_tokens: Arc::new(AtomicU64::new(0)),
58            output_tokens: Arc::new(AtomicU64::new(0)),
59        })
60    }
61
62    /// Cumulative `(input, output)` tokens used so far, across all calls.
63    pub fn token_usage(&self) -> (u64, u64) {
64        (
65            self.input_tokens.load(Ordering::Relaxed),
66            self.output_tokens.load(Ordering::Relaxed),
67        )
68    }
69
70    async fn complete(&self, system: &str, user: &str, max_tokens: u32) -> Result<String> {
71        // Chat Completions (not the Responses API): what OpenAI-compatible servers implement.
72        let req = CompletionRequest::new(user.to_string()).preamble(system);
73        // Reasoning models spend part of the completion budget thinking, hence the headroom.
74        let req = if self.openai_native {
75            req.additional_params(serde_json::json!({
76                "max_completion_tokens": (max_tokens as u64 * 4).max(4000)
77            }))
78        } else {
79            req.max_tokens(max_tokens as u64)
80        };
81        let resp = self
82            .client
83            .chat(&self.model)
84            .call(req)
85            .await
86            .context("LLM completion")?;
87        self.input_tokens
88            .fetch_add(resp.usage.input_tokens.unwrap_or(0), Ordering::Relaxed);
89        self.output_tokens
90            .fetch_add(resp.usage.output_tokens.unwrap_or(0), Ordering::Relaxed);
91        Ok(resp.text().trim().to_string())
92    }
93
94    /// One-line summary of what a symbol does, plus a few natural-language queries it answers.
95    pub async fn enrich_symbol(
96        &self,
97        name: &str,
98        signature: &str,
99        context: &str,
100    ) -> Result<Enrichment> {
101        let system = "You document code for a search index. Reply in exactly this format:\n\
102             SUMMARY: <one concise sentence describing what this does>\n\
103             QUERIES: <3 natural-language questions a developer might ask that this code answers, separated by ` | `>\n\
104             No preamble, no code fences.";
105        let user = format!("Symbol: {name}\nSignature: {signature}\nContext: {context}");
106        let raw = self.complete(system, &user, 160).await?;
107        Ok(parse_enrichment(&raw))
108    }
109
110    /// Synthesize an evidence-based coding style guide for one repo from a measured evidence pack
111    /// (naming/size/idiom stats + enforced-config facts). The prompt forbids generic advice and
112    /// requires every rule to cite the evidence.
113    /// Synthesize ONE section of the style guide from its tailored evidence. Returns the section's
114    /// bullet rules (no heading — the caller adds it). Multiple sections are generated per repo.
115    pub async fn guide_section(
116        &self,
117        repo: &str,
118        title: &str,
119        guidance: &str,
120        evidence: &str,
121    ) -> Result<String> {
122        let system = "You are a senior engineer documenting ONE section of a codebase's house-style \
123            guide so an AI coding agent writes code INDISTINGUISHABLE from this repo's. You are given \
124            REAL evidence (signatures, code excerpts, dependency facts, measured stats) for this \
125            section only. Make the repo's tacit conventions explicit.\n\
126            Hard rules:\n\
127            - Output ONLY this section's rules as terse markdown bullets — NO top-level heading, no \
128            preamble, no closing note.\n\
129            - Cite a real FILE PATH (and symbol/signature) for every non-trivial claim; quote a short \
130            real snippet as the example. Where useful, show a correct vs incorrect example.\n\
131            - Prefer the repo's OWN wrapper/abstraction over a raw library when the evidence shows \
132            one; name it exactly.\n\
133            - NEVER output generic advice that would apply to any codebase ('use meaningful names', \
134            'write tests', 'handle errors gracefully'). Delete such lines.\n\
135            - Strength tags: **MUST**/**NEVER** for patterns across ≥3 files or enforced by config; \
136            'prefer' for weaker signals. Give a one-line reason where non-obvious.\n\
137            - If the evidence shows no consistent convention, reply with exactly: \
138            No consistent convention observed.";
139        let user = format!(
140            "Repository: {repo}\nSection: {title}\nWhat to cover: {guidance}\n\n# Evidence\n\n{evidence}"
141        );
142        self.complete(system, &user, 900).await
143    }
144
145    /// Narrate an execution flow: given an entry point and its ordered call trace, produce a
146    /// concise markdown walkthrough plus a few flow-oriented questions the flow answers.
147    pub async fn narrate_flow(
148        &self,
149        entry: &str,
150        signature: &str,
151        trace: &str,
152    ) -> Result<(String, Vec<String>)> {
153        let system = "You explain how a code flow works to a developer new to the codebase. \
154             Given an entry point and its call trace (each line: depth, relation, symbol, \
155             location), reply in exactly this format:\n\
156             NARRATION: <2-5 sentence walkthrough of what happens step by step; reference the \
157             actual symbol names>\n\
158             QUERIES: <3 questions a developer might ask about this flow, separated by ` | `>\n\
159             No preamble, no code fences.";
160        let user = format!("Entry point: {entry}\nSignature: {signature}\nCall trace:\n{trace}");
161        let raw = self.complete(system, &user, 320).await?;
162        Ok(parse_flow(&raw))
163    }
164}
165
166/// Parse the `NARRATION:` / `QUERIES:` reply format, tolerantly (narration may span lines).
167fn parse_flow(raw: &str) -> (String, Vec<String>) {
168    let (narr_part, q_part) = raw.split_once("QUERIES:").unwrap_or((raw, ""));
169    let narration = narr_part
170        .trim()
171        .trim_start_matches("NARRATION:")
172        .trim()
173        .to_string();
174    let queries = q_part
175        .split('|')
176        .map(|s| s.trim().to_string())
177        .filter(|s| !s.is_empty())
178        .collect();
179    (
180        if narration.is_empty() {
181            raw.trim().to_string()
182        } else {
183            narration
184        },
185        queries,
186    )
187}
188
189/// Parse the `SUMMARY:` / `QUERIES:` reply format, tolerantly.
190fn parse_enrichment(raw: &str) -> Enrichment {
191    let mut summary = String::new();
192    let mut queries = Vec::new();
193    for line in raw.lines() {
194        let l = line.trim();
195        if let Some(rest) = l.strip_prefix("SUMMARY:") {
196            summary = rest.trim().to_string();
197        } else if let Some(rest) = l.strip_prefix("QUERIES:") {
198            queries = rest
199                .split('|')
200                .map(|q| q.trim().to_string())
201                .filter(|q| !q.is_empty())
202                .collect();
203        }
204    }
205    if summary.is_empty() {
206        // Model didn't follow format — fall back to the first non-empty line.
207        summary = raw
208            .lines()
209            .map(str::trim)
210            .find(|l| !l.is_empty())
211            .unwrap_or("")
212            .to_string();
213    }
214    Enrichment { summary, queries }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn parses_enrichment_format() {
223        let raw = "SUMMARY: Opens a pooled Postgres connection.\nQUERIES: how to connect to postgres | get a db connection | open database pool";
224        let e = parse_enrichment(raw);
225        assert_eq!(e.summary, "Opens a pooled Postgres connection.");
226        assert_eq!(e.queries.len(), 3);
227        assert_eq!(e.queries[0], "how to connect to postgres");
228    }
229
230    // Real API call — run with a key: `cargo test -p comindly-llm -- --ignored`.
231    #[tokio::test]
232    #[ignore]
233    async fn live_enrich_symbol() {
234        let key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY");
235        let c = LlmClient::new(None, DEFAULT_MODEL, Some(&key)).unwrap();
236        let e = c
237            .enrich_symbol(
238                "AsyncTaskRunner",
239                "class AsyncTaskRunner(BaseTaskRunner)",
240                "acme/database/executors.py",
241            )
242            .await
243            .expect("enrich");
244        assert!(!e.summary.is_empty());
245        eprintln!("summary: {}", e.summary);
246        eprintln!("queries: {:?}", e.queries);
247    }
248}
comindly_core/llm.rs

comindly_core/
llm.rs
