claude_agent/context/
provider.rs1use 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}