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_local_md(mut self, content: impl Into<String>) -> Self {
44        self.local_md.push(content.into());
45        self
46    }
47
48    pub fn with_priority(mut self, priority: i32) -> Self {
49        self.priority = priority;
50        self
51    }
52}
53
54#[async_trait]
55impl MemoryProvider for InMemoryProvider {
56    fn name(&self) -> &str {
57        "in-memory"
58    }
59
60    async fn load(&self) -> ContextResult<MemoryContent> {
61        let mut content = MemoryContent::default();
62
63        if let Some(ref prompt) = self.system_prompt {
64            content.claude_md.push(prompt.clone());
65        }
66
67        content.claude_md.extend(self.claude_md.clone());
68        content.local_md.extend(self.local_md.clone());
69
70        Ok(content)
71    }
72
73    fn priority(&self) -> i32 {
74        self.priority
75    }
76}
77
78pub struct FileMemoryProvider {
79    path: PathBuf,
80    priority: i32,
81}
82
83impl FileMemoryProvider {
84    pub fn new(path: impl Into<PathBuf>) -> Self {
85        Self {
86            path: path.into(),
87            priority: 0,
88        }
89    }
90
91    pub fn with_priority(mut self, priority: i32) -> Self {
92        self.priority = priority;
93        self
94    }
95}
96
97#[async_trait]
98impl MemoryProvider for FileMemoryProvider {
99    fn name(&self) -> &str {
100        "file"
101    }
102
103    async fn load(&self) -> ContextResult<MemoryContent> {
104        let mut loader = super::MemoryLoader::new();
105        loader.load(&self.path).await
106    }
107
108    fn priority(&self) -> i32 {
109        self.priority
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[tokio::test]
118    async fn test_in_memory_provider() {
119        let provider = InMemoryProvider::new()
120            .with_system_prompt("You are a helpful assistant.")
121            .with_claude_md("# Project Rules");
122
123        let content = provider.load().await.unwrap();
124        assert_eq!(content.claude_md.len(), 2);
125    }
126}