1use crate::error::{Result, SconeError};
7
8#[derive(Debug, Clone, PartialEq)]
10pub struct ExtractedFact {
11 pub subject: String,
12 pub predicate: String,
13 pub object: String,
14 pub confidence: f32,
15}
16
17pub const ANSWER_SYSTEM_V1: &str = "Answer from the provided memory context. \
19Cite nothing you cannot find there; say so when the context lacks the answer.";
20
21pub const ANSWER_SYSTEM_V2: &str = "You answer questions from retrieved \
24personal memory. Reply with ONLY the specific fact or detail asked for - a \
25short phrase, no preamble, no explanation. Prefer the exact wording found in \
26the context. If the question refers to time ('first', 'last', 'in May'), use \
27the timestamps and ordering in the context to pick the right instance. If the \
28context does not contain the answer, reply exactly: unknown";
29
30pub const ANSWER_SYSTEM_V3: &str = "You answer questions from retrieved \
39personal memory. Work in two steps, in one reply. Step 1: copy the 2 to 4 \
40context lines that bear on the question, each on its own line starting \
41EVIDENCE:, keeping their [timestamps]. Step 2: end with one line starting \
42ANSWER: followed by the specific fact or detail asked for, short, in the \
43context's own wording. When the question involves time ('first', 'last', \
44'before', 'in May'), order the evidence timestamps and pick accordingly; \
45when facts conflict, the latest timestamp wins. If the evidence does not \
46contain the answer, end with exactly: ANSWER: unknown";
47
48pub const TWO_PASS_EXTRACT_SYSTEM: &str = "From the retrieved memory \
56context, copy every line that could bear on the question, verbatim with \
57its [timestamp], one per line. No commentary, no answer. If nothing \
58bears on it, reply exactly: NO EVIDENCE";
59
60pub trait LlmProvider: Send {
61 fn id(&self) -> &str;
62 fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>>;
63 fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String>;
65 fn answer(&self, question: &str, context: &str) -> Result<String> {
67 self.answer_with_system(ANSWER_SYSTEM_V1, question, context)
68 }
69}
70
71pub struct FakeLlm {
74 facts: Vec<ExtractedFact>,
75 fail: Option<String>,
76 answer: Option<String>,
77 calls: std::cell::RefCell<Vec<String>>,
78}
79
80impl FakeLlm {
81 pub fn new(facts: Vec<ExtractedFact>) -> Self {
82 Self {
83 facts,
84 fail: None,
85 answer: None,
86 calls: std::cell::RefCell::new(Vec::new()),
87 }
88 }
89
90 pub fn with_answer(mut self, answer: &str) -> Self {
92 self.answer = Some(answer.to_owned());
93 self
94 }
95
96 pub fn failing(message: &str) -> Self {
97 Self {
98 facts: Vec::new(),
99 fail: Some(message.to_owned()),
100 answer: None,
101 calls: std::cell::RefCell::new(Vec::new()),
102 }
103 }
104
105 pub fn calls(&self) -> Vec<String> {
106 self.calls.borrow().clone()
107 }
108}
109
110impl LlmProvider for FakeLlm {
111 fn id(&self) -> &str {
112 "fake"
113 }
114
115 fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
116 self.calls.borrow_mut().push(text.to_owned());
117 match &self.fail {
118 Some(msg) => Err(SconeError::Llm(msg.clone())),
119 None => Ok(self.facts.clone()),
120 }
121 }
122
123 fn answer_with_system(&self, _system: &str, question: &str, context: &str) -> Result<String> {
124 match (&self.fail, &self.answer) {
125 (Some(msg), _) => Err(SconeError::Llm(msg.clone())),
126 (None, Some(programmed)) => Ok(programmed.clone()),
127 (None, None) => Ok(format!(
128 "answer to {question} given {} bytes",
129 context.len()
130 )),
131 }
132 }
133}
134
135const EXTRACTION_PROMPT: &str = "Extract durable factual statements from the text as a STRICT \
136JSON array. Each element: {\"subject\": string, \"predicate\": string, \"object\": string, \
137\"confidence\": number 0..1}. Subjects are entities (people, projects, tools, places). \
138Predicates are short verb phrases. Only facts stated or strongly implied; no speculation. \
139Reply with the JSON array ONLY — no prose, no code fences.";
140
141fn parse_extraction(content: &str) -> Result<Vec<ExtractedFact>> {
142 let trimmed = content
143 .trim()
144 .trim_start_matches("```json")
145 .trim_start_matches("```")
146 .trim_end_matches("```")
147 .trim();
148 let value: serde_json::Value = serde_json::from_str(trimmed).map_err(|e| {
149 SconeError::Llm(format!(
150 "model did not return JSON: {e}; got: {}",
151 content.chars().take(120).collect::<String>()
152 ))
153 })?;
154 let array = value
155 .as_array()
156 .ok_or_else(|| SconeError::Llm("model did not return JSON array".into()))?;
157 array
158 .iter()
159 .map(|f| {
160 Ok(ExtractedFact {
161 subject: f["subject"]
162 .as_str()
163 .ok_or_else(|| SconeError::Llm("fact missing subject".into()))?
164 .to_owned(),
165 predicate: f["predicate"]
166 .as_str()
167 .ok_or_else(|| SconeError::Llm("fact missing predicate".into()))?
168 .to_owned(),
169 object: f["object"]
170 .as_str()
171 .ok_or_else(|| SconeError::Llm("fact missing object".into()))?
172 .to_owned(),
173 confidence: f["confidence"].as_f64().unwrap_or(0.5) as f32,
174 })
175 })
176 .collect()
177}
178
179pub const DEFAULT_LLM_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(180);
183
184fn http_json(
185 req: ureq::RequestBuilder<ureq::typestate::WithBody>,
186 timeout: std::time::Duration,
187 body: serde_json::Value,
188) -> Result<serde_json::Value> {
189 let mut res = req
190 .config()
191 .timeout_global(Some(timeout))
192 .build()
193 .send_json(body)
194 .map_err(|e| SconeError::Llm(format!("http: {e}")))?;
195 res.body_mut()
196 .read_json()
197 .map_err(|e| SconeError::Llm(format!("http body: {e}")))
198}
199
200pub struct OpenAiCompatible {
202 base_url: String,
203 model: String,
204 api_key: Option<String>,
205 timeout: std::time::Duration,
206 think: Option<bool>,
207}
208
209impl OpenAiCompatible {
210 pub fn new(base_url: &str, model: &str, api_key: Option<String>) -> Self {
211 Self {
212 base_url: base_url.trim_end_matches('/').to_owned(),
213 model: model.to_owned(),
214 api_key,
215 timeout: DEFAULT_LLM_TIMEOUT,
216 think: None,
217 }
218 }
219
220 pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
221 self.timeout = timeout;
222 self
223 }
224
225 pub fn with_think(mut self, think: bool) -> Self {
229 self.think = Some(think);
230 self
231 }
232
233 fn chat(&self, system: &str, user: &str) -> Result<String> {
234 let mut req = ureq::post(format!("{}/chat/completions", self.base_url));
235 if let Some(key) = &self.api_key {
236 req = req.header("authorization", format!("Bearer {key}"));
237 }
238 let mut body = serde_json::json!({
239 "model": self.model,
240 "messages": [
241 {"role": "system", "content": system},
242 {"role": "user", "content": user},
243 ],
244 });
245 if let Some(think) = self.think {
246 body["think"] = serde_json::Value::Bool(think);
247 }
248 let value = http_json(req, self.timeout, body)?;
249 value["choices"][0]["message"]["content"]
250 .as_str()
251 .map(str::to_owned)
252 .ok_or_else(|| SconeError::Llm("no content in chat response".into()))
253 }
254}
255
256impl LlmProvider for OpenAiCompatible {
257 fn id(&self) -> &str {
258 &self.model
259 }
260
261 fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
262 parse_extraction(&self.chat(EXTRACTION_PROMPT, text)?)
263 }
264
265 fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String> {
266 self.chat(
267 system,
268 &format!("Context:\n{context}\n\nQuestion: {question}"),
269 )
270 }
271}
272
273pub struct AnthropicProvider {
275 base_url: String,
276 model: String,
277 api_key: String,
278 timeout: std::time::Duration,
279}
280
281impl AnthropicProvider {
282 pub fn new(base_url: &str, model: &str, api_key: &str) -> Self {
283 Self {
284 base_url: base_url.trim_end_matches('/').to_owned(),
285 model: model.to_owned(),
286 api_key: api_key.to_owned(),
287 timeout: DEFAULT_LLM_TIMEOUT,
288 }
289 }
290
291 pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
292 self.timeout = timeout;
293 self
294 }
295
296 fn message(&self, system: &str, user: &str) -> Result<String> {
297 let req = ureq::post(format!("{}/v1/messages", self.base_url))
298 .header("x-api-key", self.api_key.as_str())
299 .header("anthropic-version", "2023-06-01");
300 let body = serde_json::json!({
301 "model": self.model,
302 "max_tokens": 1024,
303 "system": system,
304 "messages": [{"role": "user", "content": user}],
305 });
306 let value = http_json(req, self.timeout, body)?;
307 value["content"][0]["text"]
308 .as_str()
309 .map(str::to_owned)
310 .ok_or_else(|| SconeError::Llm("no text in messages response".into()))
311 }
312}
313
314impl LlmProvider for AnthropicProvider {
315 fn id(&self) -> &str {
316 &self.model
317 }
318
319 fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
320 parse_extraction(&self.message(EXTRACTION_PROMPT, text)?)
321 }
322
323 fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String> {
324 self.message(
325 system,
326 &format!("Context:\n{context}\n\nQuestion: {question}"),
327 )
328 }
329}