gemini_memory_rs/retrieval/
extractor.rs

1//! The out-of-band retrieval-plan extractor seam.
2//!
3//! Plan extraction is a small structured-output model call that runs *after*
4//! the final transcript, while the model is already speaking. It never sits on
5//! the response path, so its failure mode is "slightly worse retrieval", not
6//! "slower conversation".
7
8use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use std::sync::Arc;
11use std::time::Duration;
12
13use super::deterministic::DeterministicPlanner;
14use super::plan::RetrievalPlan;
15use crate::core::{MemoryError, TurnId};
16
17/// The system instruction for the plan extractor.
18///
19/// The prohibitions matter as much as the task: a model asked to plan retrieval
20/// will otherwise answer the user, invent memories, or treat its own previous
21/// output as something the user said.
22pub const RETRIEVAL_PLAN_INSTRUCTION: &str = "\
23You produce a retrieval plan for a personal memory system. Given the user's \
24most recent utterance and a little surrounding conversation, decide which of \
25the user's existing stored memories might be relevant to answering them.
26
27Return only the structured plan. Specifically:
28- Do NOT answer the user's question.
29- Do NOT propose anything to remember; that is a separate task.
30- Do NOT treat the assistant's own statements as facts about the user.
31- Do NOT infer sensitive attributes (health, religion, politics, sexuality) \
32  that the user did not state.
33- Set requires_memory to false for generic factual, visual or world-knowledge \
34  questions that do not depend on this user's history.
35- Prefer few, specific search terms over many broad ones.
36- When the user is not speaking English, put BOTH their words and the English \
37  equivalents in topics and lexical_queries. Stored facts are written in \
38  English and retrieval is literal word matching, so a question about \
39  'khaana' finds nothing unless the plan also carries 'food' and 'diet'. \
40  English on both sides is what makes the two meet; do not rely on the stored \
41  fact having guessed which Hindi or Tamil word would be used months later.";
42
43/// What the extractor is given.
44#[derive(Debug, Clone)]
45pub struct RetrievalExtractionContext {
46    /// The finalized user utterance.
47    pub transcript: String,
48    /// Up to four preceding user turns, oldest first.
49    pub recent_user_turns: Vec<String>,
50    /// Up to two preceding assistant turns, for reference resolution only.
51    pub recent_assistant_turns: Vec<String>,
52    /// Entities already known in this session.
53    pub known_entities: Vec<String>,
54    /// The rule-based plan, as a starting point.
55    pub deterministic: RetrievalPlan,
56    /// The turn being planned for.
57    pub turn_id: TurnId,
58    /// The generation the request was issued at.
59    pub generation: u64,
60    /// Evaluation time.
61    pub now: DateTime<Utc>,
62}
63
64impl RetrievalExtractionContext {
65    /// Render the context as a prompt body.
66    pub fn to_prompt(&self) -> String {
67        let mut out = String::new();
68        if !self.recent_user_turns.is_empty() {
69            out.push_str("Earlier user turns:\n");
70            for turn in &self.recent_user_turns {
71                out.push_str("- ");
72                out.push_str(turn);
73                out.push('\n');
74            }
75        }
76        if !self.recent_assistant_turns.is_empty() {
77            out.push_str("\nAssistant turns (for reference resolution only):\n");
78            for turn in &self.recent_assistant_turns {
79                out.push_str("- ");
80                out.push_str(turn);
81                out.push('\n');
82            }
83        }
84        if !self.known_entities.is_empty() {
85            out.push_str("\nKnown entities: ");
86            out.push_str(&self.known_entities.join(", "));
87            out.push('\n');
88        }
89        out.push_str("\nCurrent user utterance:\n");
90        out.push_str(&self.transcript);
91        out.push('\n');
92        out
93    }
94}
95
96/// Produces a retrieval plan from conversation context.
97#[async_trait]
98pub trait RetrievalPlanExtractor: Send + Sync {
99    /// Extract a plan.
100    async fn extract(
101        &self,
102        context: RetrievalExtractionContext,
103    ) -> Result<RetrievalPlan, MemoryError>;
104}
105
106/// The rule-based planner, exposed as an extractor.
107///
108/// This is the default: it needs no model, it is deterministic, and it is what
109/// every other implementation degrades to.
110pub struct DeterministicPlanExtractor {
111    planner: Arc<DeterministicPlanner>,
112}
113
114impl DeterministicPlanExtractor {
115    /// Wrap a planner.
116    pub fn new(planner: Arc<DeterministicPlanner>) -> Self {
117        Self { planner }
118    }
119}
120
121#[async_trait]
122impl RetrievalPlanExtractor for DeterministicPlanExtractor {
123    async fn extract(
124        &self,
125        context: RetrievalExtractionContext,
126    ) -> Result<RetrievalPlan, MemoryError> {
127        Ok(self.planner.plan(
128            &context.transcript,
129            context.turn_id,
130            context.generation,
131            context.now,
132        ))
133    }
134}
135
136/// Runs an extractor under a deadline and falls back to the deterministic plan.
137///
138/// The deterministic plan is already in the context, so a failed or slow model
139/// call costs nothing beyond the deadline itself.
140pub struct BoundedPlanExtractor {
141    inner: Arc<dyn RetrievalPlanExtractor>,
142    timeout: Duration,
143}
144
145impl BoundedPlanExtractor {
146    /// Bound `inner` to `timeout`.
147    pub fn new(inner: Arc<dyn RetrievalPlanExtractor>, timeout: Duration) -> Self {
148        Self { inner, timeout }
149    }
150}
151
152#[async_trait]
153impl RetrievalPlanExtractor for BoundedPlanExtractor {
154    async fn extract(
155        &self,
156        context: RetrievalExtractionContext,
157    ) -> Result<RetrievalPlan, MemoryError> {
158        let fallback = context.deterministic.clone();
159        match tokio::time::timeout(self.timeout, self.inner.extract(context)).await {
160            Ok(Ok(plan)) => Ok(plan.normalized()),
161            Ok(Err(_)) | Err(_) => Ok(fallback),
162        }
163    }
164}
165
166/// The JSON Schema a structured-output model call should be constrained to.
167pub fn retrieval_plan_schema() -> serde_json::Value {
168    let schema = schemars::schema_for!(RetrievalPlan);
169    serde_json::to_value(schema).unwrap_or(serde_json::Value::Null)
170}
171
172/// Build an extraction context from a transcript and a rule-based plan.
173pub fn context_for(
174    planner: &DeterministicPlanner,
175    transcript: &str,
176    turn_id: TurnId,
177    generation: u64,
178    now: DateTime<Utc>,
179) -> RetrievalExtractionContext {
180    RetrievalExtractionContext {
181        transcript: transcript.to_string(),
182        recent_user_turns: Vec::new(),
183        recent_assistant_turns: Vec::new(),
184        known_entities: Vec::new(),
185        deterministic: planner.plan(transcript, turn_id, generation, now),
186        turn_id,
187        generation,
188        now,
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::retrieval::deterministic::KnownEntities;
196    use crate::retrieval::plan::RetrievalIntent;
197
198    fn planner() -> Arc<DeterministicPlanner> {
199        let mut known = KnownEntities::new();
200        known.insert("Rhea", "rhea");
201        Arc::new(DeterministicPlanner::with_entities(known))
202    }
203
204    fn context(text: &str) -> RetrievalExtractionContext {
205        context_for(&planner(), text, TurnId(2), 2, Utc::now())
206    }
207
208    struct AlwaysFails;
209
210    #[async_trait]
211    impl RetrievalPlanExtractor for AlwaysFails {
212        async fn extract(
213            &self,
214            _context: RetrievalExtractionContext,
215        ) -> Result<RetrievalPlan, MemoryError> {
216            Err(MemoryError::Extraction("model unavailable".into()))
217        }
218    }
219
220    struct Hangs;
221
222    #[async_trait]
223    impl RetrievalPlanExtractor for Hangs {
224        async fn extract(
225            &self,
226            _context: RetrievalExtractionContext,
227        ) -> Result<RetrievalPlan, MemoryError> {
228            tokio::time::sleep(Duration::from_secs(30)).await;
229            unreachable!("the bound should fire first")
230        }
231    }
232
233    #[tokio::test]
234    async fn the_deterministic_extractor_produces_a_usable_plan() {
235        let plan = DeterministicPlanExtractor::new(planner())
236            .extract(context("what does Rhea like to eat"))
237            .await
238            .unwrap();
239        assert!(plan.requires_memory);
240        assert!(!plan.lexical_queries.is_empty());
241    }
242
243    #[tokio::test]
244    async fn a_failing_model_extractor_degrades_to_the_rule_based_plan() {
245        let bounded = BoundedPlanExtractor::new(Arc::new(AlwaysFails), Duration::from_millis(50));
246        let plan = bounded
247            .extract(context("what does Rhea like to eat"))
248            .await
249            .unwrap();
250        assert!(plan.requires_memory);
251        assert_eq!(plan.intent, RetrievalIntent::RelationshipReference);
252    }
253
254    #[tokio::test]
255    async fn a_hanging_model_extractor_is_abandoned_at_the_deadline() {
256        let bounded = BoundedPlanExtractor::new(Arc::new(Hangs), Duration::from_millis(20));
257        let started = std::time::Instant::now();
258        let plan = bounded
259            .extract(context("what does Rhea like to eat"))
260            .await
261            .unwrap();
262        assert!(started.elapsed() < Duration::from_secs(1));
263        assert!(plan.requires_memory);
264    }
265
266    #[test]
267    fn the_prompt_separates_user_turns_from_assistant_turns() {
268        let mut ctx = context("what should we cook");
269        ctx.recent_user_turns = vec!["I am pescatarian".into()];
270        ctx.recent_assistant_turns = vec!["You mentioned salmon".into()];
271        let prompt = ctx.to_prompt();
272        assert!(prompt.contains("Earlier user turns:"));
273        assert!(prompt.contains("reference resolution only"));
274        assert!(prompt.ends_with("what should we cook\n"));
275    }
276
277    #[test]
278    fn the_instruction_forbids_answering_and_storing() {
279        assert!(RETRIEVAL_PLAN_INSTRUCTION.contains("Do NOT answer"));
280        assert!(RETRIEVAL_PLAN_INSTRUCTION.contains("separate task"));
281    }
282
283    #[test]
284    fn the_plan_schema_is_a_json_object_schema() {
285        let schema = retrieval_plan_schema();
286        assert_eq!(schema["type"], "object");
287        assert!(schema["properties"]["requires_memory"].is_object());
288    }
289}
gemini_memory_rs/retrieval/extractor.rs

gemini_memory_rs/retrieval/
extractor.rs
