use crate::error::{Result, SconeError};
#[derive(Debug, Clone, PartialEq)]
pub struct ExtractedFact {
pub subject: String,
pub predicate: String,
pub object: String,
pub confidence: f32,
}
pub const ANSWER_SYSTEM_V1: &str = "Answer from the provided memory context. \
Cite nothing you cannot find there; say so when the context lacks the answer.";
pub const ANSWER_SYSTEM_V2: &str = "You answer questions from retrieved \
personal memory. Reply with ONLY the specific fact or detail asked for - a \
short phrase, no preamble, no explanation. Prefer the exact wording found in \
the context. If the question refers to time ('first', 'last', 'in May'), use \
the timestamps and ordering in the context to pick the right instance. If the \
context does not contain the answer, reply exactly: unknown";
pub trait LlmProvider: Send {
fn id(&self) -> &str;
fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>>;
fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String>;
fn answer(&self, question: &str, context: &str) -> Result<String> {
self.answer_with_system(ANSWER_SYSTEM_V1, question, context)
}
}
pub struct FakeLlm {
facts: Vec<ExtractedFact>,
fail: Option<String>,
answer: Option<String>,
calls: std::cell::RefCell<Vec<String>>,
}
impl FakeLlm {
pub fn new(facts: Vec<ExtractedFact>) -> Self {
Self {
facts,
fail: None,
answer: None,
calls: std::cell::RefCell::new(Vec::new()),
}
}
pub fn with_answer(mut self, answer: &str) -> Self {
self.answer = Some(answer.to_owned());
self
}
pub fn failing(message: &str) -> Self {
Self {
facts: Vec::new(),
fail: Some(message.to_owned()),
answer: None,
calls: std::cell::RefCell::new(Vec::new()),
}
}
pub fn calls(&self) -> Vec<String> {
self.calls.borrow().clone()
}
}
impl LlmProvider for FakeLlm {
fn id(&self) -> &str {
"fake"
}
fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
self.calls.borrow_mut().push(text.to_owned());
match &self.fail {
Some(msg) => Err(SconeError::Llm(msg.clone())),
None => Ok(self.facts.clone()),
}
}
fn answer_with_system(&self, _system: &str, question: &str, context: &str) -> Result<String> {
match (&self.fail, &self.answer) {
(Some(msg), _) => Err(SconeError::Llm(msg.clone())),
(None, Some(programmed)) => Ok(programmed.clone()),
(None, None) => Ok(format!(
"answer to {question} given {} bytes",
context.len()
)),
}
}
}
const EXTRACTION_PROMPT: &str = "Extract durable factual statements from the text as a STRICT \
JSON array. Each element: {\"subject\": string, \"predicate\": string, \"object\": string, \
\"confidence\": number 0..1}. Subjects are entities (people, projects, tools, places). \
Predicates are short verb phrases. Only facts stated or strongly implied; no speculation. \
Reply with the JSON array ONLY — no prose, no code fences.";
fn parse_extraction(content: &str) -> Result<Vec<ExtractedFact>> {
let trimmed = content
.trim()
.trim_start_matches("```json")
.trim_start_matches("```")
.trim_end_matches("```")
.trim();
let value: serde_json::Value = serde_json::from_str(trimmed).map_err(|e| {
SconeError::Llm(format!(
"model did not return JSON: {e}; got: {}",
content.chars().take(120).collect::<String>()
))
})?;
let array = value
.as_array()
.ok_or_else(|| SconeError::Llm("model did not return JSON array".into()))?;
array
.iter()
.map(|f| {
Ok(ExtractedFact {
subject: f["subject"]
.as_str()
.ok_or_else(|| SconeError::Llm("fact missing subject".into()))?
.to_owned(),
predicate: f["predicate"]
.as_str()
.ok_or_else(|| SconeError::Llm("fact missing predicate".into()))?
.to_owned(),
object: f["object"]
.as_str()
.ok_or_else(|| SconeError::Llm("fact missing object".into()))?
.to_owned(),
confidence: f["confidence"].as_f64().unwrap_or(0.5) as f32,
})
})
.collect()
}
pub const DEFAULT_LLM_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(180);
fn http_json(
req: ureq::RequestBuilder<ureq::typestate::WithBody>,
timeout: std::time::Duration,
body: serde_json::Value,
) -> Result<serde_json::Value> {
let mut res = req
.config()
.timeout_global(Some(timeout))
.build()
.send_json(body)
.map_err(|e| SconeError::Llm(format!("http: {e}")))?;
res.body_mut()
.read_json()
.map_err(|e| SconeError::Llm(format!("http body: {e}")))
}
pub struct OpenAiCompatible {
base_url: String,
model: String,
api_key: Option<String>,
timeout: std::time::Duration,
}
impl OpenAiCompatible {
pub fn new(base_url: &str, model: &str, api_key: Option<String>) -> Self {
Self {
base_url: base_url.trim_end_matches('/').to_owned(),
model: model.to_owned(),
api_key,
timeout: DEFAULT_LLM_TIMEOUT,
}
}
pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
fn chat(&self, system: &str, user: &str) -> Result<String> {
let mut req = ureq::post(format!("{}/chat/completions", self.base_url));
if let Some(key) = &self.api_key {
req = req.header("authorization", format!("Bearer {key}"));
}
let body = serde_json::json!({
"model": self.model,
"messages": [
{"role": "system", "content": system},
{"role": "user", "content": user},
],
});
let value = http_json(req, self.timeout, body)?;
value["choices"][0]["message"]["content"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| SconeError::Llm("no content in chat response".into()))
}
}
impl LlmProvider for OpenAiCompatible {
fn id(&self) -> &str {
&self.model
}
fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
parse_extraction(&self.chat(EXTRACTION_PROMPT, text)?)
}
fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String> {
self.chat(
system,
&format!("Context:\n{context}\n\nQuestion: {question}"),
)
}
}
pub struct AnthropicProvider {
base_url: String,
model: String,
api_key: String,
timeout: std::time::Duration,
}
impl AnthropicProvider {
pub fn new(base_url: &str, model: &str, api_key: &str) -> Self {
Self {
base_url: base_url.trim_end_matches('/').to_owned(),
model: model.to_owned(),
api_key: api_key.to_owned(),
timeout: DEFAULT_LLM_TIMEOUT,
}
}
pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
fn message(&self, system: &str, user: &str) -> Result<String> {
let req = ureq::post(format!("{}/v1/messages", self.base_url))
.header("x-api-key", self.api_key.as_str())
.header("anthropic-version", "2023-06-01");
let body = serde_json::json!({
"model": self.model,
"max_tokens": 1024,
"system": system,
"messages": [{"role": "user", "content": user}],
});
let value = http_json(req, self.timeout, body)?;
value["content"][0]["text"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| SconeError::Llm("no text in messages response".into()))
}
}
impl LlmProvider for AnthropicProvider {
fn id(&self) -> &str {
&self.model
}
fn extract_facts(&self, text: &str) -> Result<Vec<ExtractedFact>> {
parse_extraction(&self.message(EXTRACTION_PROMPT, text)?)
}
fn answer_with_system(&self, system: &str, question: &str, context: &str) -> Result<String> {
self.message(
system,
&format!("Context:\n{context}\n\nQuestion: {question}"),
)
}
}