claude_agent/context/
provider.rs

1//! Memory provider trait and implementations.
2
3use std::path::PathBuf;
4
5use async_trait::async_trait;
6
7use super::{ContextResult, MemoryContent};
8
9pub const MAX_IMPORT_DEPTH: usize = 5;
10
11#[async_trait]
12pub trait MemoryProvider: Send + Sync {
13    fn name(&self) -> &str;
14    async fn load(&self) -> ContextResult<MemoryContent>;
15    fn priority(&self) -> i32 {
16        0
17    }
18}
19
20#[derive(Debug, Clone, Default)]
21pub struct InMemoryProvider {
22    pub system_prompt: Option<String>,
23    pub claude_md: Vec<String>,
24    pub local_md: Vec<String>,
25    pub priority: i32,
26}
27
28impl InMemoryProvider {
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
34        self.system_prompt = Some(prompt.into());
35        self
36    }
37
38    pub fn with_claude_md(mut self, content: impl Into<String>) -> Self {
39        self.claude_md.push(content.into());
40        self
41    }
42
43    pub fn with_priority(mut self, priority: i32) -> Self {
44        self.priority = priority;
45        self
46    }
47}
48
49#[async_trait]
50impl MemoryProvider for InMemoryProvider {
51    fn name(&self) -> &str {
52        "in-memory"
53    }
54
55    async fn load(&self) -> ContextResult<MemoryContent> {
56        let mut content = MemoryContent::default();
57
58        if let Some(ref prompt) = self.system_prompt {
59            content.claude_md.push(prompt.clone());
60        }
61
62        content.claude_md.extend(self.claude_md.clone());
63        content.local_md.extend(self.local_md.clone());
64
65        Ok(content)
66    }
67
68    fn priority(&self) -> i32 {
69        self.priority
70    }
71}
72
73pub struct FileMemoryProvider {
74    path: PathBuf,
75    priority: i32,
76}
77
78impl FileMemoryProvider {
79    pub fn new(path: impl Into<PathBuf>) -> Self {
80        Self {
81            path: path.into(),
82            priority: 0,
83        }
84    }
85
86    pub fn with_priority(mut self, priority: i32) -> Self {
87        self.priority = priority;
88        self
89    }
90}
91
92#[async_trait]
93impl MemoryProvider for FileMemoryProvider {
94    fn name(&self) -> &str {
95        "file"
96    }
97
98    async fn load(&self) -> ContextResult<MemoryContent> {
99        let mut loader = super::MemoryLoader::new();
100        loader.load_all(&self.path).await
101    }
102
103    fn priority(&self) -> i32 {
104        self.priority
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[tokio::test]
113    async fn test_in_memory_provider() {
114        let provider = InMemoryProvider::new()
115            .with_system_prompt("You are a helpful assistant.")
116            .with_claude_md("# Project Rules");
117
118        let content = provider.load().await.unwrap();
119        assert_eq!(content.claude_md.len(), 2);
120    }
121}