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_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}