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 concise sentence describing what the code does.
48    pub summary: String,
49    /// Natural-language questions a developer might ask that this code answers.
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(
127        &self,
128        prompts: &Prompts,
129        name: &str,
130        signature: &str,
131        context: &str,
132    ) -> Result<Enrichment> {
133        let user = format!("Symbol: {name}\nSignature: {signature}\nContext: {context}");
134        let r: Enrichment = self
135            .complete(Task::Summary, &prompts.system(Task::Summary), &user, 160)
136            .await?;
137        Ok(Enrichment {
138            summary: r.summary.trim().to_string(),
139            queries: tidy(r.queries),
140        })
141    }
142
143    /// ONE section of a repo's style guide from its tailored evidence, as bullet rules without a
144    /// heading, or `None` when the evidence shows no consistent convention.
145    pub async fn guide_section(
146        &self,
147        prompts: &Prompts,
148        repo: &str,
149        title: &str,
150        guidance: &str,
151        evidence: &str,
152    ) -> Result<Option<String>> {
153        let user = format!(
154            "Repository: {repo}\nSection: {title}\nWhat to cover: {guidance}\n\n# Evidence\n\n{evidence}"
155        );
156        let r: GuideSection = self
157            .complete(Task::Guide, &prompts.system(Task::Guide), &user, 900)
158            .await?;
159        let rules = tidy(r.rules);
160        Ok((r.observed && !rules.is_empty()).then(|| {
161            rules
162                .iter()
163                .map(|r| format!("- {}", r.trim_start_matches("- ")))
164                .collect::<Vec<_>>()
165                .join("\n")
166        }))
167    }
168
169    /// Narrate an execution flow: given an entry point and its ordered call trace, produce a
170    /// concise markdown walkthrough plus a few flow-oriented questions the flow answers.
171    pub async fn narrate_flow(
172        &self,
173        prompts: &Prompts,
174        entry: &str,
175        signature: &str,
176        trace: &str,
177    ) -> Result<(String, Vec<String>)> {
178        let user = format!("Entry point: {entry}\nSignature: {signature}\nCall trace:\n{trace}");
179        let r: Walkthrough = self
180            .complete(Task::Flow, &prompts.system(Task::Flow), &user, 320)
181            .await?;
182        Ok((r.narration.trim().to_string(), tidy(r.queries)))
183    }
184}
185
186/// What a guide section call returns.
187#[derive(Deserialize, JsonSchema)]
188#[schemars(crate = "rig::schemars")]
189struct GuideSection {
190    /// False when the evidence shows no consistent convention.
191    observed: bool,
192    /// One rule per item, in markdown, without a leading dash.
193    rules: Vec<String>,
194}
195
196/// What a flow call returns.
197#[derive(Deserialize, JsonSchema)]
198#[schemars(crate = "rig::schemars")]
199struct Walkthrough {
200    /// A 2-5 sentence walkthrough naming the actual symbols.
201    narration: String,
202    /// Questions a developer might ask about this flow.
203    queries: Vec<String>,
204}
205
206/// The JSON schema of `task`'s reply, as text for the system message.
207pub fn reply_schema(task: Task) -> String {
208    let schema = match task {
209        Task::Summary => schemars::schema_for!(Enrichment),
210        Task::Guide => schemars::schema_for!(GuideSection),
211        Task::Flow => schemars::schema_for!(Walkthrough),
212    };
213    serde_json::to_string_pretty(&schema).unwrap_or_default()
214}
215
216fn tidy(items: Vec<String>) -> Vec<String> {
217    items
218        .into_iter()
219        .map(|q| q.trim().to_string())
220        .filter(|q| !q.is_empty())
221        .collect()
222}
223
224/// The JSON reply, also when a server that ignores the schema wraps it in prose or a code fence.
225fn parse<T: DeserializeOwned>(raw: &str) -> Result<T> {
226    let raw = raw.trim();
227    if let Ok(v) = serde_json::from_str(raw) {
228        return Ok(v);
229    }
230    let (Some(a), Some(b)) = (raw.find('{'), raw.rfind('}')) else {
231        anyhow::bail!("the LLM did not reply with JSON");
232    };
233    anyhow::ensure!(a < b, "the LLM did not reply with JSON");
234    serde_json::from_str(&raw[a..=b]).context("the LLM's JSON reply has the wrong shape")
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn reads_json_also_inside_prose() {
243        let e: Enrichment =
244            parse(r#"{"summary":"Opens a pooled connection.","queries":["how to connect"]}"#)
245                .unwrap();
246        assert_eq!(e.summary, "Opens a pooled connection.");
247        let fenced =
248            "Here it is:\n```json\n{\"narration\": \"Calls a, then b.\", \"queries\": []}\n```";
249        let w: Walkthrough = parse(fenced).unwrap();
250        assert_eq!(w.narration, "Calls a, then b.");
251        assert!(parse::<Enrichment>("SUMMARY: no json").is_err());
252        assert!(parse::<Enrichment>(r#"{"summary": 3}"#).is_err());
253    }
254
255    #[test]
256    fn the_schema_is_strict_enough_for_openai() {
257        let s = serde_json::to_value(schemars::schema_for!(GuideSection)).unwrap();
258        assert_eq!(s["required"], serde_json::json!(["observed", "rules"]));
259    }
260
261    // Real API call — run with a key: `cargo test -p comindly-llm -- --ignored`.
262    #[tokio::test]
263    #[ignore]
264    async fn live_enrich_symbol() {
265        let key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY");
266        let c = LlmClient::new(None, DEFAULT_MODEL, Some(&key)).unwrap();
267        let e = c
268            .enrich_symbol(
269                &Prompts::default(),
270                "AsyncTaskRunner",
271                "class AsyncTaskRunner(BaseTaskRunner)",
272                "acme/database/executors.py",
273            )
274            .await
275            .expect("enrich");
276        assert!(!e.summary.is_empty());
277        eprintln!("summary: {}", e.summary);
278        eprintln!("queries: {:?}", e.queries);
279    }
280
281    // A compatible server: `COMINDLY_TEST_LLM_BASE_URL` and `COMINDLY_TEST_LLM_MODEL`, e.g. Ollama.
282    #[tokio::test]
283    #[ignore]
284    async fn live_compatible_server() {
285        let base = std::env::var("COMINDLY_TEST_LLM_BASE_URL").expect("COMINDLY_TEST_LLM_BASE_URL");
286        let model = std::env::var("COMINDLY_TEST_LLM_MODEL").expect("COMINDLY_TEST_LLM_MODEL");
287        let c = LlmClient::new(Some(&base), &model, None).unwrap();
288        let p = Prompts::default();
289        let e = c
290            .enrich_symbol(
291                &p,
292                "open_pool",
293                "fn open_pool(cfg: &DbConfig) -> Pool",
294                "src/db/pool.rs",
295            )
296            .await
297            .expect("summary");
298        assert!(!e.summary.is_empty());
299        eprintln!("summary: {}\nqueries: {:?}", e.summary, e.queries);
300        let (n, _) = c
301            .narrate_flow(
302                &p,
303                "handle_order",
304                "fn handle_order(o: Order)",
305                "d1 calls acme:validate src/a.rs:3",
306            )
307            .await
308            .expect("flow");
309        eprintln!("flow: {n}");
310    }
311
312    #[tokio::test]
313    #[ignore]
314    async fn live_guide_and_flow() {
315        let key = std::env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY");
316        let c = LlmClient::new(None, DEFAULT_MODEL, Some(&key)).unwrap();
317        let p = Prompts::default();
318        let evidence = "snake_case functions 97% (1234/1268)\nsrc/db/pool.rs: pub fn open_pool(cfg: &DbConfig) -> Pool";
319        let g = c
320            .guide_section(&p, "acme", "Naming", "How things are named", evidence)
321            .await
322            .expect("guide");
323        eprintln!("guide: {g:?}");
324        let (n, q) = c
325            .narrate_flow(
326                &p,
327                "handle_order",
328                "fn handle_order(o: Order)",
329                "d1 calls acme:validate src/a.rs:3\nd2 calls acme:save src/b.rs:9",
330            )
331            .await
332            .expect("flow");
333        assert!(!n.is_empty());
334        eprintln!("flow: {n}\n{q:?}");
335    }
336}
comindly_core/llm.rs

comindly_core/
llm.rs
