rs_adk/tools/
preload_memory.rs

1//! Preload memory tool — preloads memory into the LLM context.
2//!
3//! Mirrors ADK-Python's `preload_memory_tool`. Automatically injects
4//! relevant memories into the LLM request context before generation.
5
6use crate::llm::LlmRequest;
7
8/// Tool that preloads relevant memories into the LLM request context.
9///
10/// Unlike [`crate::tools::LoadMemoryTool`] which is called by the model, this tool
11/// automatically injects memories into the system instruction or context
12/// before the model generates a response.
13#[derive(Debug, Clone, Default)]
14pub struct PreloadMemoryTool;
15
16impl PreloadMemoryTool {
17    /// Create a new preload memory tool.
18    pub fn new() -> Self {
19        Self
20    }
21
22    /// Inject memory context into the LLM request.
23    ///
24    /// Appends memory entries to the system instruction.
25    pub fn process_llm_request(&self, request: &mut LlmRequest, memories: &[String]) {
26        if memories.is_empty() {
27            return;
28        }
29
30        let memory_context = format!(
31            "\n\nRelevant memories from previous interactions:\n{}",
32            memories
33                .iter()
34                .map(|m| format!("- {m}"))
35                .collect::<Vec<_>>()
36                .join("\n")
37        );
38
39        if let Some(ref mut instruction) = request.system_instruction {
40            instruction.push_str(&memory_context);
41        } else {
42            request.system_instruction = Some(memory_context);
43        }
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn injects_memories_into_instruction() {
53        let tool = PreloadMemoryTool::new();
54        let mut request = LlmRequest::default();
55        request.system_instruction = Some("You are helpful.".into());
56
57        tool.process_llm_request(
58            &mut request,
59            &["User likes Rust".into(), "User is a developer".into()],
60        );
61
62        let instruction = request.system_instruction.unwrap();
63        assert!(instruction.contains("User likes Rust"));
64        assert!(instruction.contains("User is a developer"));
65    }
66
67    #[test]
68    fn empty_memories_noop() {
69        let tool = PreloadMemoryTool::new();
70        let mut request = LlmRequest::default();
71        request.system_instruction = Some("Original".into());
72
73        tool.process_llm_request(&mut request, &[]);
74        assert_eq!(request.system_instruction.unwrap(), "Original");
75    }
76
77    #[test]
78    fn creates_instruction_if_none() {
79        let tool = PreloadMemoryTool::new();
80        let mut request = LlmRequest::default();
81
82        tool.process_llm_request(&mut request, &["Memory 1".into()]);
83        assert!(request.system_instruction.is_some());
84        assert!(request.system_instruction.unwrap().contains("Memory 1"));
85    }
86}
rs_adk/tools/preload_memory.rs

rs_adk/tools/
preload_memory.rs
