oxi/store/
extracting_backend.rs

1//! LLM fact-extraction wrapper around any [`oxi_agent::tools::MemoryBackend`].
2//!
3//! When `settings.memory_llm_extract = true`, every `put(content, …)`
4//! is sent through a structured extraction prompt and the model's
5//! JSON response is fanned out into one `inner.put(fact_text, …)`
6//! per atomic fact.
7//!
8//! Gap-2 design: the LLM call must run inside the async runtime.
9//! The LLM extractor is implemented directly as an async function
10//! that awaits the provider stream. The sync [`FactExtractor`]
11//! trait is retained only for the **heuristic** path (pure CPU, no
12//! I/O).
13//!
14//! Failure policy: when the extractor errors, returns no facts,
15//! or returns malformed JSON, the wrapper falls back to storing
16//! `content` verbatim so `memory_retain` never silently drops user
17//! input.
18#![allow(clippy::manual_pattern_char_comparison)]
19#![allow(missing_docs)]
20#![allow(dead_code)]
21
22use std::pin::Pin;
23use std::sync::Arc;
24
25use futures::StreamExt;
26use oxi_agent::tools::{MemoryBackend, MemoryItem, ToolError};
27use oxi_ai::{Context, ProviderEvent};
28use serde::Deserialize;
29
30use super::settings::Settings;
31
32/// One atomic fact returned by the LLM extractor.
33#[derive(Debug, Clone, Deserialize)]
34pub struct ExtractedFact {
35    /// Atomic fact text. Must be non-empty after trim.
36    pub content: String,
37    /// Free-form importance in `[0, 1]`. `None` → backend default.
38    #[serde(default)]
39    pub importance: Option<f64>,
40    /// Free-form kind (`"fact"` / `"preference"` / `"context"` /
41    /// `"summary"`). `None` → kind passed into `put`.
42    #[serde(default)]
43    pub kind: Option<String>,
44}
45
46/// JSON shape returned by the extraction prompt.
47#[derive(Debug, Default, Deserialize)]
48struct ExtractionPayload {
49    #[serde(default)]
50    pub facts: Vec<ExtractedFact>,
51}
52
53/// Sync, no-IO extractor (heuristic). Safe to call anywhere.
54pub trait FactExtractor: Send + Sync {
55    /// Return `None` to defer to the inner backend with the raw payload.
56    fn extract(&self, content: &str, kind: &str) -> Option<Vec<ExtractedFact>>;
57}
58
59/// Default heuristic — splits on sentence boundaries.
60pub struct HeuristicFactExtractor;
61
62impl FactExtractor for HeuristicFactExtractor {
63    fn extract(&self, content: &str, _kind: &str) -> Option<Vec<ExtractedFact>> {
64        let facts: Vec<ExtractedFact> = content
65            .split(|c: char| c == '.' || c == '!' || c == '?' || c == '\n')
66            .map(str::trim)
67            .filter(|s| s.len() >= 10 && !s.ends_with('?'))
68            .map(|s| ExtractedFact {
69                content: s.to_string(),
70                importance: None,
71                kind: None,
72            })
73            .collect();
74        if facts.is_empty() { None } else { Some(facts) }
75    }
76}
77
78/// omp `stage_one_system.md`, truncated to the parts that drive
79/// our extraction contract (we don't need `rollout_slug` here — we
80/// only store a flat list of facts).
81pub const STAGE_ONE_SYSTEM_PROMPT: &str = "\
82You extract atomic durable facts from free-form text.\n\
83Return STRICT JSON ONLY (no markdown, no commentary).\n\
84Output contract: {\"facts\": [{\"content\": \"string\", \"importance\": number, \"kind\": \"string\"}]}.\n\
85Rules:\n\
86- Each fact is one self-contained claim (a constraint, decision, workflow, or resolved pitfall).\n\
87- Skip transient chatter, greetings, and pure questions.\n\
88- importance in [0, 1] (1 = critical, never forget).\n\
89- kind is one of: \"fact\", \"preference\", \"context\", \"summary\".\n\
90- When no durable signal exists, return {\"facts\": []}.";
91
92/// omp `stage_one_input.md` user template (`{{content}}`).
93pub const STAGE_ONE_USER_TEMPLATE: &str = "\
94Text to extract from:\n\
95{{content}}\n\n\
96Return JSON now.";
97
98/// Async LLM fact extractor. Owns a `Provider` and a `Model`.
99pub struct LlmExtractor {
100    provider: Arc<dyn oxi_ai::Provider>,
101    model: oxi_ai::Model,
102}
103
104impl std::fmt::Debug for LlmExtractor {
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        f.debug_struct("LlmExtractor")
107            .field("model", &self.model.id)
108            .finish_non_exhaustive()
109    }
110}
111
112impl LlmExtractor {
113    /// Build an extractor backed by any `oxi_ai::Provider`.
114    pub fn new(provider: Arc<dyn oxi_ai::Provider>, model: oxi_ai::Model) -> Self {
115        Self { provider, model }
116    }
117}
118
119/// Async extraction helper shared by the wrap path.
120async fn run_extraction(
121    provider: &Arc<dyn oxi_ai::Provider>,
122    model: &oxi_ai::Model,
123    content: &str,
124) -> Option<Vec<ExtractedFact>> {
125    let user = STAGE_ONE_USER_TEMPLATE.replace("{{content}}", content);
126    let mut ctx = Context::default();
127    ctx.set_system_prompt(STAGE_ONE_SYSTEM_PROMPT.to_string());
128    ctx.add_message(oxi_ai::Message::user(user));
129    let mut stream = provider.stream(model, &ctx, None).await.ok()?;
130    let mut buf = String::new();
131    while let Some(ev) = stream.next().await {
132        match ev {
133            ProviderEvent::TextDelta { delta, .. } => buf.push_str(&delta),
134            ProviderEvent::Done { .. } => break,
135            ProviderEvent::Error { .. } => return None,
136            _ => {}
137        }
138    }
139    parse_facts(&buf)
140}
141
142/// Parse the strict-JSON `{"facts": […]}` payload that the model
143/// returns. Strips ```json fences defensively.
144fn parse_facts(buf: &str) -> Option<Vec<ExtractedFact>> {
145    let trimmed = buf
146        .trim()
147        .trim_start_matches("```json")
148        .trim_start_matches("```")
149        .trim_end_matches("```")
150        .trim();
151    let payload: ExtractionPayload = serde_json::from_str(trimmed).ok()?;
152    let facts: Vec<ExtractedFact> = payload
153        .facts
154        .into_iter()
155        .filter(|f| !f.content.trim().is_empty())
156        .collect();
157    if facts.is_empty() { None } else { Some(facts) }
158}
159
160/// Wrapper around any [`MemoryBackend`].
161pub struct ExtractingMemoryBackend {
162    inner: Arc<dyn MemoryBackend>,
163    heuristic: Arc<dyn FactExtractor>,
164    llm: Option<LlmExtractor>,
165}
166
167impl std::fmt::Debug for ExtractingMemoryBackend {
168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169        f.debug_struct("ExtractingMemoryBackend")
170            .field("inner", &"<dyn MemoryBackend>")
171            .field("heuristic", &"<dyn FactExtractor>")
172            .field("llm", &self.llm.as_ref().map(|l| l.model.id.clone()))
173            .finish()
174    }
175}
176
177impl ExtractingMemoryBackend {
178    /// Build a wrapper that runs extraction before delegating to the
179    /// inner backend.
180    pub fn new(
181        inner: Arc<dyn MemoryBackend>,
182        heuristic: Arc<dyn FactExtractor>,
183        llm: Option<LlmExtractor>,
184    ) -> Self {
185        Self {
186            inner,
187            heuristic,
188            llm,
189        }
190    }
191}
192
193impl MemoryBackend for ExtractingMemoryBackend {
194    fn put<'a>(
195        &'a self,
196        content: &'a str,
197        kind: &'a str,
198        subject: &'a str,
199    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
200        let inner = Arc::clone(&self.inner);
201        let heuristic = Arc::clone(&self.heuristic);
202        let llm = self.llm.as_ref().map(|l| LlmExtractorHandle {
203            provider: Arc::clone(&l.provider),
204            model: l.model.clone(),
205        });
206
207        Box::pin(async move {
208            const MAX_FACTS_PER_PUT: usize = 16;
209
210            let facts = if let Some(handle) = llm {
211                handle.extract(content).await
212            } else {
213                heuristic.extract(content, kind)
214            };
215
216            match facts {
217                Some(fs) if !fs.is_empty() => {
218                    let mut last: Result<String, ToolError> = Ok(content.to_string());
219                    for fact in fs.into_iter().take(MAX_FACTS_PER_PUT) {
220                        let fkind = fact.kind.as_deref().unwrap_or(kind);
221                        match inner.put(&fact.content, fkind, subject).await {
222                            Ok(id) => last = Ok(id),
223                            Err(e) => {
224                                last = Err(e);
225                                break;
226                            }
227                        }
228                    }
229                    last
230                }
231                _ => inner.put(content, kind, subject).await,
232            }
233        })
234    }
235
236    fn search<'a>(
237        &'a self,
238        query: &'a str,
239        k: usize,
240    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
241        self.inner.search(query, k)
242    }
243
244    fn list<'a>(
245        &'a self,
246        subject: &'a str,
247    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
248        self.inner.list(subject)
249    }
250
251    fn delete<'a>(
252        &'a self,
253        id: &'a str,
254    ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
255        self.inner.delete(id)
256    }
257
258    fn memory_info(&self) -> Option<String> {
259        self.inner.memory_info()
260    }
261}
262
263/// Borrow-friendly clone of an [`LlmExtractor`] so it can move into
264/// the `async move { }` block on `put`.
265struct LlmExtractorHandle {
266    provider: Arc<dyn oxi_ai::Provider>,
267    model: oxi_ai::Model,
268}
269
270impl LlmExtractorHandle {
271    async fn extract(&self, content: &str) -> Option<Vec<ExtractedFact>> {
272        run_extraction(&self.provider, &self.model, content).await
273    }
274}
275
276/// Resolve the LLM extractor from the wired `Oxi` engine. Returns
277/// `None` when the user hasn't configured
278/// `memory_llm_extract_model` or the pattern doesn't resolve.
279pub fn try_make_llm_extractor(oxi: &oxi_sdk::Oxi, settings: &Settings) -> Option<LlmExtractor> {
280    let pat = settings.memory_llm_extract_model.trim();
281    if pat.is_empty() {
282        return None;
283    }
284    let model = oxi.resolve_model(pat).ok()?;
285    let provider = oxi.providers().get(model.provider.as_str())?;
286    Some(LlmExtractor::new(provider, model))
287}
288
289/// Wire the wrapper on top of an existing memory backend.
290pub fn wrap_with_extractor(
291    inner: Arc<dyn MemoryBackend>,
292    settings: &Settings,
293    oxi: Option<&oxi_sdk::Oxi>,
294) -> Arc<dyn MemoryBackend> {
295    if !settings.memory_llm_extract {
296        return inner;
297    }
298    let heuristic: Arc<dyn FactExtractor> = Arc::new(HeuristicFactExtractor);
299    let llm = oxi.and_then(|o| try_make_llm_extractor(o, settings));
300    Arc::new(ExtractingMemoryBackend::new(inner, heuristic, llm))
301}
oxi/store/extracting_backend.rs

oxi/store/
extracting_backend.rs
