tj_core/
llm.rs

1//! Pluggable LLM backend for the journal's optional AI operations
2//! (consolidation, dream backfill). One small trait, several adapters, picked by
3//! name so this public package can grow new providers without touching callers.
4//!
5//! Default is **`claude-p`** — the local Claude CLI on your subscription, so the
6//! out-of-the-box experience needs no API key. Override with `TJ_BACKEND` (env,
7//! global) or a per-command `--backend`:
8//!
9//! - `claude-p` (default) — local `claude -p`, Haiku, subscription auth.
10//! - `anthropic` — direct Anthropic API (`ANTHROPIC_API_KEY`).
11//! - `openai` — any OpenAI-compatible chat API (`OPENAI_API_KEY`,
12//!   `TJ_OPENAI_BASE_URL`, `TJ_OPENAI_MODEL`). Covers OpenAI, Codex, and other
13//!   compatible providers by pointing the base URL.
14//! - `ollama` — a local Ollama model (its OpenAI-compatible endpoint), **free**:
15//!   no key, no network beyond localhost. `TJ_OLLAMA_URL`, `TJ_OLLAMA_MODEL`.
16//!
17//! A backend that isn't usable (no key, no `claude` on PATH) yields `Ok(None)`
18//! from [`backend_from_env`] so the caller skips cleanly — we never fabricate
19//! output without a model.
20
21use anyhow::{anyhow, Context};
22use serde::{Deserialize, Serialize};
23use std::time::Duration;
24
25/// One AI call: a prompt in, the model's text reply out.
26pub trait LlmBackend: Send + Sync {
27    fn complete(&self, prompt: &str, max_tokens: u32) -> anyhow::Result<String>;
28    /// Stable label for logs / provenance.
29    fn name(&self) -> &'static str;
30}
31
32/// Resolve the backend from an explicit name (e.g. a `--backend` flag) or
33/// `TJ_BACKEND`, defaulting to `claude-p`. Returns:
34/// - `Ok(Some(_))` — a usable backend,
35/// - `Ok(None)` — the chosen backend is unavailable (no key / no `claude`); the
36///   caller should skip,
37/// - `Err(_)` — an unknown backend name (a typo worth surfacing).
38pub fn backend_from_env(explicit: Option<&str>) -> anyhow::Result<Option<Box<dyn LlmBackend>>> {
39    let name = explicit
40        .map(str::to_string)
41        .or_else(|| std::env::var("TJ_BACKEND").ok())
42        .filter(|s| !s.trim().is_empty())
43        .unwrap_or_else(|| "claude-p".to_string());
44
45    match name.trim() {
46        "claude-p" | "claude" | "agent-sdk" => {
47            if crate::classifier::agent_sdk::claude_on_path() {
48                Ok(Some(Box::new(ClaudeCliBackend::from_env())))
49            } else {
50                Ok(None)
51            }
52        }
53        "anthropic" | "api" => match std::env::var("ANTHROPIC_API_KEY") {
54            Ok(key) if !key.is_empty() => Ok(Some(Box::new(AnthropicBackend::new(key)))),
55            _ => Ok(None),
56        },
57        "openai" | "codex" => match std::env::var("OPENAI_API_KEY") {
58            Ok(key) if !key.is_empty() => Ok(Some(Box::new(OpenAiBackend::openai(key)))),
59            _ => Ok(None),
60        },
61        "ollama" => Ok(Some(Box::new(OpenAiBackend::ollama()))),
62        other => Err(anyhow!(
63            "unknown backend '{other}' (expected: claude-p, anthropic, openai, ollama)"
64        )),
65    }
66}
67
68// ---------------------------------------------------------------------------
69// claude -p (default) — local CLI, subscription auth, no API key.
70// ---------------------------------------------------------------------------
71
72pub struct ClaudeCliBackend {
73    model: String,
74}
75
76impl ClaudeCliBackend {
77    pub fn from_env() -> Self {
78        let model = std::env::var("TJ_CONSOLIDATE_MODEL")
79            .unwrap_or_else(|_| crate::classifier::agent_sdk::DEFAULT_MODEL.to_string());
80        Self { model }
81    }
82}
83
84impl LlmBackend for ClaudeCliBackend {
85    fn complete(&self, prompt: &str, _max_tokens: u32) -> anyhow::Result<String> {
86        crate::classifier::agent_sdk::run_claude_json(
87            &crate::classifier::agent_sdk::ClaudeBinaryStdinRunner,
88            &self.model,
89            prompt,
90        )
91    }
92    fn name(&self) -> &'static str {
93        "claude-p"
94    }
95}
96
97// ---------------------------------------------------------------------------
98// Anthropic direct API.
99// ---------------------------------------------------------------------------
100
101pub struct AnthropicBackend {
102    api_key: String,
103    model: String,
104    base_url: String,
105    timeout: Duration,
106}
107
108impl AnthropicBackend {
109    pub fn new(api_key: String) -> Self {
110        let model = std::env::var("TJ_CONSOLIDATE_MODEL")
111            .unwrap_or_else(|_| "claude-haiku-4-5-20251001".to_string());
112        let base_url = std::env::var("TJ_CONSOLIDATE_BASE_URL")
113            .unwrap_or_else(|_| "https://api.anthropic.com".to_string());
114        Self {
115            api_key,
116            model,
117            base_url,
118            timeout: Duration::from_secs(60),
119        }
120    }
121}
122
123#[derive(Serialize)]
124struct AnthropicReq<'a> {
125    model: &'a str,
126    max_tokens: u32,
127    messages: Vec<AnthropicMsg<'a>>,
128}
129#[derive(Serialize)]
130struct AnthropicMsg<'a> {
131    role: &'a str,
132    content: &'a str,
133}
134#[derive(Deserialize)]
135struct AnthropicResp {
136    content: Vec<AnthropicBlock>,
137}
138#[derive(Deserialize)]
139struct AnthropicBlock {
140    #[serde(rename = "type")]
141    kind: String,
142    #[serde(default)]
143    text: String,
144}
145
146impl LlmBackend for AnthropicBackend {
147    fn complete(&self, prompt: &str, max_tokens: u32) -> anyhow::Result<String> {
148        let body = AnthropicReq {
149            model: &self.model,
150            max_tokens,
151            messages: vec![AnthropicMsg {
152                role: "user",
153                content: prompt,
154            }],
155        };
156        let resp: AnthropicResp = ureq::post(&format!("{}/v1/messages", self.base_url))
157            .timeout(self.timeout)
158            .set("x-api-key", &self.api_key)
159            .set("anthropic-version", "2023-06-01")
160            .set("content-type", "application/json")
161            .send_json(serde_json::to_value(&body)?)
162            .context("Anthropic API request failed")?
163            .into_json()
164            .context("decode Anthropic response")?;
165        resp.content
166            .iter()
167            .find(|b| b.kind == "text")
168            .map(|b| b.text.clone())
169            .ok_or_else(|| anyhow!("no text content in Anthropic response"))
170    }
171    fn name(&self) -> &'static str {
172        "anthropic"
173    }
174}
175
176// ---------------------------------------------------------------------------
177// OpenAI-compatible — covers OpenAI, Codex, Ollama, and any compatible server.
178// ---------------------------------------------------------------------------
179
180pub struct OpenAiBackend {
181    api_key: Option<String>,
182    model: String,
183    base_url: String,
184    label: &'static str,
185    timeout: Duration,
186}
187
188impl OpenAiBackend {
189    pub fn openai(api_key: String) -> Self {
190        Self {
191            api_key: Some(api_key),
192            model: std::env::var("TJ_OPENAI_MODEL").unwrap_or_else(|_| "gpt-4o-mini".to_string()),
193            base_url: std::env::var("TJ_OPENAI_BASE_URL")
194                .unwrap_or_else(|_| "https://api.openai.com".to_string()),
195            label: "openai",
196            timeout: Duration::from_secs(60),
197        }
198    }
199
200    pub fn ollama() -> Self {
201        Self {
202            api_key: None, // local; no auth
203            model: std::env::var("TJ_OLLAMA_MODEL").unwrap_or_else(|_| "llama3.1".to_string()),
204            base_url: std::env::var("TJ_OLLAMA_URL")
205                .unwrap_or_else(|_| "http://localhost:11434".to_string()),
206            label: "ollama",
207            timeout: Duration::from_secs(120),
208        }
209    }
210}
211
212#[derive(Serialize)]
213struct OpenAiReq<'a> {
214    model: &'a str,
215    max_tokens: u32,
216    messages: Vec<AnthropicMsg<'a>>,
217}
218#[derive(Deserialize)]
219struct OpenAiResp {
220    choices: Vec<OpenAiChoice>,
221}
222#[derive(Deserialize)]
223struct OpenAiChoice {
224    message: OpenAiMsg,
225}
226#[derive(Deserialize)]
227struct OpenAiMsg {
228    #[serde(default)]
229    content: String,
230}
231
232impl LlmBackend for OpenAiBackend {
233    fn complete(&self, prompt: &str, max_tokens: u32) -> anyhow::Result<String> {
234        let body = OpenAiReq {
235            model: &self.model,
236            max_tokens,
237            messages: vec![AnthropicMsg {
238                role: "user",
239                content: prompt,
240            }],
241        };
242        let mut req = ureq::post(&format!("{}/v1/chat/completions", self.base_url))
243            .timeout(self.timeout)
244            .set("content-type", "application/json");
245        if let Some(key) = &self.api_key {
246            req = req.set("authorization", &format!("Bearer {key}"));
247        }
248        let resp: OpenAiResp = req
249            .send_json(serde_json::to_value(&body)?)
250            .with_context(|| format!("{} request failed", self.label))?
251            .into_json()
252            .context("decode OpenAI-compatible response")?;
253        resp.choices
254            .into_iter()
255            .next()
256            .map(|c| c.message.content)
257            .ok_or_else(|| anyhow!("no choices in {} response", self.label))
258    }
259    fn name(&self) -> &'static str {
260        self.label
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    struct EnvGuard(&'static str, Option<String>);
269    impl EnvGuard {
270        fn set(k: &'static str, v: &str) -> Self {
271            let prev = std::env::var(k).ok();
272            std::env::set_var(k, v);
273            Self(k, prev)
274        }
275        fn unset(k: &'static str) -> Self {
276            let prev = std::env::var(k).ok();
277            std::env::remove_var(k);
278            Self(k, prev)
279        }
280    }
281    impl Drop for EnvGuard {
282        fn drop(&mut self) {
283            match &self.1 {
284                Some(v) => std::env::set_var(self.0, v),
285                None => std::env::remove_var(self.0),
286            }
287        }
288    }
289
290    // Serialise env-touching tests (process-global env).
291    static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
292
293    #[test]
294    fn unknown_backend_errors() {
295        let _l = ENV_LOCK.lock().unwrap();
296        assert!(backend_from_env(Some("nonsense")).is_err());
297    }
298
299    #[test]
300    fn anthropic_unavailable_without_key_is_none() {
301        let _l = ENV_LOCK.lock().unwrap();
302        let _g = EnvGuard::unset("ANTHROPIC_API_KEY");
303        assert!(backend_from_env(Some("anthropic")).unwrap().is_none());
304    }
305
306    #[test]
307    fn anthropic_with_key_resolves() {
308        let _l = ENV_LOCK.lock().unwrap();
309        let _g = EnvGuard::set("ANTHROPIC_API_KEY", "k");
310        let b = backend_from_env(Some("anthropic")).unwrap().unwrap();
311        assert_eq!(b.name(), "anthropic");
312    }
313
314    #[test]
315    fn ollama_always_resolves_no_key() {
316        let _l = ENV_LOCK.lock().unwrap();
317        let b = backend_from_env(Some("ollama")).unwrap().unwrap();
318        assert_eq!(b.name(), "ollama");
319    }
320
321    #[test]
322    fn openai_calls_chat_completions_and_parses() {
323        let mut server = mockito::Server::new();
324        let m = server
325            .mock("POST", "/v1/chat/completions")
326            .with_status(200)
327            .with_header("content-type", "application/json")
328            .with_body(
329                serde_json::json!({
330                    "choices": [{"message": {"role": "assistant", "content": "hello from openai"}}]
331                })
332                .to_string(),
333            )
334            .create();
335        let b = OpenAiBackend {
336            api_key: Some("k".into()),
337            model: "gpt-4o-mini".into(),
338            base_url: server.url(),
339            label: "openai",
340            timeout: Duration::from_secs(5),
341        };
342        let out = b.complete("hi", 64).unwrap();
343        m.assert();
344        assert_eq!(out, "hello from openai");
345    }
346
347    #[test]
348    fn anthropic_calls_messages_and_parses() {
349        let mut server = mockito::Server::new();
350        let m = server
351            .mock("POST", "/v1/messages")
352            .with_status(200)
353            .with_header("content-type", "application/json")
354            .with_body(
355                serde_json::json!({
356                    "content": [{"type": "text", "text": "hello from anthropic"}]
357                })
358                .to_string(),
359            )
360            .create();
361        let b = AnthropicBackend {
362            api_key: "k".into(),
363            model: "claude-haiku-4-5-20251001".into(),
364            base_url: server.url(),
365            timeout: Duration::from_secs(5),
366        };
367        let out = b.complete("hi", 64).unwrap();
368        m.assert();
369        assert_eq!(out, "hello from anthropic");
370    }
371}
tj_core/llm.rs

tj_core/
llm.rs
