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;
16use rig::schemars::{self, JsonSchema};
17use serde::de::DeserializeOwned;
18use serde::Deserialize;
19
20use crate::prompts::{Prompts, Task};
21
22/// Default model — cheap and fast, suitable for bulk symbol summaries.
23pub const DEFAULT_MODEL: &str = "gpt-6-luna";
24
25/// Where requests go without an endpoint setting: OpenAI.
26pub const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
27
28pub struct LlmClient {
29    // Chat Completions API (works with OpenAI *and* any OpenAI-compatible endpoint), not the
30    // Responses API — proxies/Ollama/vLLM implement `/chat/completions`, not `/responses`.
31    client: openai::OpenAI,
32    model: String,
33    /// Tasks that run on another model than `model`.
34    models: std::collections::BTreeMap<Task, String>,
35    // OpenAI's own API takes `max_completion_tokens` (required by newer models); compatible
36    // servers (Ollama, vLLM, LiteLLM) expect `max_tokens`.
37    openai_native: bool,
38    // Cumulative token usage across all calls (for live progress + a final total).
39    input_tokens: Arc<AtomicU64>,
40    output_tokens: Arc<AtomicU64>,
41}
42
43/// LLM-generated enrichment for one symbol.
44#[derive(Debug, Clone, Deserialize, JsonSchema)]
45#[schemars(crate = "rig::schemars")]
46pub struct Enrichment {
47    /// One sentence, at most 30 words, on what the code does, what for, and its effects.
48    pub summary: String,
49    /// Three questions someone would type to find this code, in the domain's words.
50    pub queries: Vec<String>,
51}
52
53impl LlmClient {
54    /// A client for `model` at `base_url` (OpenAI when absent), with `key`.
55    pub fn new(base_url: Option<&str>, model: &str, key: Option<&str>) -> Result<Self> {
56        let key = key.map(str::to_string);
57        let client = match base_url {
58            Some(base) => openai::OpenAIConfig::new(key.unwrap_or_else(|| "sk-noauth".into()))
59                .with_base_url(base)
60                .client(),
61            None => openai::OpenAIConfig::new(key.context("no LLM key")?).client(),
62        };
63        Ok(Self {
64            client,
65            model: model.to_string(),
66            models: Default::default(),
67            openai_native: base_url.is_none(),
68            input_tokens: Arc::new(AtomicU64::new(0)),
69            output_tokens: Arc::new(AtomicU64::new(0)),
70        })
71    }
72
73    /// Cumulative `(input, output)` tokens used so far, across all calls.
74    pub fn token_usage(&self) -> (u64, u64) {
75        (
76            self.input_tokens.load(Ordering::Relaxed),
77            self.output_tokens.load(Ordering::Relaxed),
78        )
79    }
80
81    /// Run `task` on its own model instead of the default one.
82    pub fn with_models(mut self, models: std::collections::BTreeMap<Task, String>) -> Self {
83        self.models = models;
84        self
85    }
86
87    /// The model `task` runs on.
88    pub fn model(&self, task: Task) -> &str {
89        self.models.get(&task).unwrap_or(&self.model)
90    }
91
92    /// A reply in the shape of `T`, asked for as structured output (a JSON schema).
93    async fn complete<T: DeserializeOwned + JsonSchema>(
94        &self,
95        task: Task,
96        system: &str,
97        user: &str,
98        max_tokens: u32,
99    ) -> Result<T> {
100        // Chat Completions (not the Responses API): what OpenAI-compatible servers implement.
101        let req = CompletionRequest::new(user.to_string())
102            .preamble(system)
103            .output_schema(schemars::schema_for!(T));
104        // Reasoning models spend part of the completion budget thinking, hence the headroom.
105        let req = if self.openai_native {
106            req.additional_params(serde_json::json!({
107                "max_completion_tokens": (max_tokens as u64 * 4).max(4000)
108            }))
109        } else {
110            req.max_tokens(max_tokens as u64)
111        };
112        let resp = self
113            .client
114            .chat(self.model(task))
115            .call(req)
116            .await
117            .context("LLM completion")?;
118        self.input_tokens
119            .fetch_add(resp.usage.input_tokens.unwrap_or(0), Ordering::Relaxed);
120        self.output_tokens
121            .fetch_add(resp.usage.output_tokens.unwrap_or(0), Ordering::Relaxed);
122        parse(&resp.text())
123    }
124
125    /// One-line summary of what a symbol does, plus a few natural-language queries it answers.
126    pub async fn enrich_symbol(&self, prompts: &Prompts, brief: &str) -> Result<Enrichment> {
127        let r: Enrichment = self
128            .complete(Task::Summary, &prompts.system(Task::Summary), brief, 400)
129            .await?;
130        Ok(Enrichment {
131            summary: r.summary.trim().to_string(),
132            queries: tidy(r.queries),
133        })
134    }
135
136    /// ONE section of a repo's style guide from its tailored evidence, as bullet rules without a
137    /// heading, or `None` when the evidence shows no consistent convention.
138    pub async fn guide_section(
139        &self,
140        prompts: &Prompts,
141        repo: &str,
142        title: &str,
143        guidance: &str,
144        evidence: &str,
145    ) -> Result<Option<String>> {
146        let user = format!(
147            "Repository: {repo}\nSection: {title}\nWhat to cover: {guidance}\n\n# Evidence\n\n{evidence}"
148        );
149        let r: GuideSection = self
150            .complete(Task::Guide, &prompts.system(Task::Guide), &user, 900)
151            .await?;
152        let rules = tidy(r.rules);
153        Ok((r.observed && !rules.is_empty()).then(|| {
154            rules
155                .iter()
156                .map(|r| format!("- {}", r.trim_start_matches("- ")))
157                .collect::<Vec<_>>()
158                .join("\n")
159        }))
160    }
161
162    /// Narrate an execution flow: given an entry point and its ordered call trace, produce a
163    /// concise markdown walkthrough plus a few flow-oriented questions the flow answers.
164    pub async fn narrate_flow(
165        &self,
166        prompts: &Prompts,
167        entry: &str,
168        signature: &str,
169        trace: &str,
170    ) -> Result<(String, Vec<String>)> {
171        let user = format!("Entry point: {entry}\nSignature: {signature}\nCall trace:\n{trace}");
172        let r: Walkthrough = self
173            .complete(Task::Flow, &prompts.system(Task::Flow), &user, 700)
174            .await?;
175        Ok((r.narration.trim().to_string(), tidy(r.queries)))
176    }
177}
178
179/// What a guide section call returns.
180#[derive(Deserialize, JsonSchema)]
181#[schemars(crate = "rig::schemars")]
182struct GuideSection {
183    /// False when the evidence shows no consistent convention.
184    observed: bool,
185    /// One rule per item, in markdown, without a leading dash.
186    rules: Vec<String>,
187}
188
189/// What a flow call returns.
190#[derive(Deserialize, JsonSchema)]
191#[schemars(crate = "rig::schemars")]
192struct Walkthrough {
193    /// The walkthrough in 3 to 6 short sentences, at most 120 words, symbols in backticks.
194    narration: String,
195    /// Three questions this flow answers, in the domain's words.
196    queries: Vec<String>,
197}
198
199/// The JSON schema of `task`'s reply, as text for the system message.
200pub fn reply_schema(task: Task) -> String {
201    let schema = match task {
202        Task::Summary => schemars::schema_for!(Enrichment),
203        Task::Guide => schemars::schema_for!(GuideSection),
204        Task::Flow => schemars::schema_for!(Walkthrough),
205    };
206    serde_json::to_string_pretty(&schema).unwrap_or_default()
207}
208
209fn tidy(items: Vec<String>) -> Vec<String> {
210    items
211        .into_iter()
212        .map(|q| q.trim().to_string())
213        .filter(|q| !q.is_empty())
214        .collect()
215}
216
217/// The JSON reply, also when a server that ignores the schema wraps it in prose or a code fence.
218fn parse<T: DeserializeOwned>(raw: &str) -> Result<T> {
219    let raw = raw.trim();
220    if let Ok(v) = serde_json::from_str(raw) {
221        return Ok(v);
222    }
223    let (Some(a), Some(b)) = (raw.find('{'), raw.rfind('}')) else {
224        anyhow::bail!("the LLM did not reply with JSON");
225    };
226    anyhow::ensure!(a < b, "the LLM did not reply with JSON");
227    serde_json::from_str(&raw[a..=b]).context("the LLM's JSON reply has the wrong shape")
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn reads_json_also_inside_prose() {
236        let e: Enrichment =
237            parse(r#"{"summary":"Opens a pooled connection.","queries":["how to connect"]}"#)
238                .unwrap();
239        assert_eq!(e.summary, "Opens a pooled connection.");
240        let fenced =
241            "Here it is:\n```json\n{\"narration\": \"Calls a, then b.\", \"queries\": []}\n```";
242        let w: Walkthrough = parse(fenced).unwrap();
243        assert_eq!(w.narration, "Calls a, then b.");
244        assert!(parse::<Enrichment>("SUMMARY: no json").is_err());
245        assert!(parse::<Enrichment>(r#"{"summary": 3}"#).is_err());
246    }
247
248    #[test]
249    fn the_schema_is_strict_enough_for_openai() {
250        let s = serde_json::to_value(schemars::schema_for!(GuideSection)).unwrap();
251        assert_eq!(s["required"], serde_json::json!(["observed", "rules"]));
252    }
253
254    // Real API call — run with a key: `cargo test -p comindly-llm -- --ignored`.
255    #[tokio::test]
256    #[ignore]
257    async fn live_enrich_symbol() {
258        let key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY");
259        let c = LlmClient::new(None, DEFAULT_MODEL, Some(&key)).unwrap();
260        let e = c
261            .enrich_symbol(
262                &Prompts::default(),
263                "Symbol: AsyncTaskRunner\nKind: Class in repository acme, acme/database/executors.py:12\nSignature: class AsyncTaskRunner(BaseTaskRunner)",
264            )
265            .await
266            .expect("enrich");
267        assert!(!e.summary.is_empty());
268        eprintln!("summary: {}", e.summary);
269        eprintln!("queries: {:?}", e.queries);
270    }
271
272    // A compatible server: `COMINDLY_TEST_LLM_BASE_URL` and `COMINDLY_TEST_LLM_MODEL`, e.g. Ollama.
273    #[tokio::test]
274    #[ignore]
275    async fn live_compatible_server() {
276        let base = std::env::var("COMINDLY_TEST_LLM_BASE_URL").expect("COMINDLY_TEST_LLM_BASE_URL");
277        let model = std::env::var("COMINDLY_TEST_LLM_MODEL").expect("COMINDLY_TEST_LLM_MODEL");
278        let c = LlmClient::new(Some(&base), &model, None).unwrap();
279        let p = Prompts::default();
280        let e = c
281            .enrich_symbol(
282                &p,
283                "Symbol: open_pool\nKind: Function in repository acme, src/db/pool.rs:3\nSignature: fn open_pool(cfg: &DbConfig) -> Pool",
284            )
285            .await
286            .expect("summary");
287        assert!(!e.summary.is_empty());
288        eprintln!("summary: {}\nqueries: {:?}", e.summary, e.queries);
289        let (n, _) = c
290            .narrate_flow(
291                &p,
292                "handle_order",
293                "fn handle_order(o: Order)",
294                "d1 calls acme:validate src/a.rs:3",
295            )
296            .await
297            .expect("flow");
298        eprintln!("flow: {n}");
299    }
300
301    #[tokio::test]
302    #[ignore]
303    async fn live_guide_and_flow() {
304        let key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY");
305        let c = LlmClient::new(None, DEFAULT_MODEL, Some(&key)).unwrap();
306        let p = Prompts::default();
307        let evidence = "snake_case functions 97% (1234/1268)\nsrc/db/pool.rs: pub fn open_pool(cfg: &DbConfig) -> Pool";
308        let g = c
309            .guide_section(&p, "acme", "Naming", "How things are named", evidence)
310            .await
311            .expect("guide");
312        eprintln!("guide: {g:?}");
313        let (n, q) = c
314            .narrate_flow(
315                &p,
316                "handle_order",
317                "fn handle_order(o: Order)",
318                "d1 calls acme:validate src/a.rs:3\nd2 calls acme:save src/b.rs:9",
319            )
320            .await
321            .expect("flow");
322        assert!(!n.is_empty());
323        eprintln!("flow: {n}\n{q:?}");
324    }
325}
comindly_core/llm.rs

comindly_core/
llm.rs
