ricecoder_domain_agents/
knowledge_base.rs1use crate::error::{DomainAgentError, Result};
4use crate::models::{Domain, KnowledgeBase, KnowledgeEntry};
5use ricecoder_storage::PathResolver;
6use tracing::{debug, info};
7
8pub struct KnowledgeBaseManager {
10 knowledge_bases: std::collections::HashMap<String, KnowledgeBase>,
11}
12
13impl KnowledgeBaseManager {
14 pub fn new() -> Self {
16 Self {
17 knowledge_bases: std::collections::HashMap::new(),
18 }
19 }
20}
21
22impl Default for KnowledgeBaseManager {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27
28impl KnowledgeBaseManager {
29
30 pub async fn load_knowledge_base(&mut self, domain: &str) -> Result<()> {
32 debug!("Loading knowledge base for domain: {}", domain);
33
34 let kb_path = self.get_knowledge_base_path(domain);
35
36 if kb_path.exists() {
37 let content = tokio::fs::read_to_string(&kb_path).await?;
38 let kb: KnowledgeBase = serde_yaml::from_str(&content)?;
39 self.knowledge_bases.insert(domain.to_string(), kb);
40 info!("Loaded knowledge base for domain: {}", domain);
41 } else {
42 debug!("Knowledge base not found for domain: {}, creating empty", domain);
43 let kb = KnowledgeBase::new(domain, "1.0.0");
44 self.knowledge_bases.insert(domain.to_string(), kb);
45 }
46
47 Ok(())
48 }
49
50 pub fn get_knowledge_base(&self, domain: &str) -> Result<&KnowledgeBase> {
52 self.knowledge_bases
53 .get(domain)
54 .ok_or_else(|| DomainAgentError::KnowledgeNotAvailable(domain.to_string()))
55 }
56
57 pub fn get_knowledge_base_mut(&mut self, domain: &str) -> Result<&mut KnowledgeBase> {
59 self.knowledge_bases
60 .get_mut(domain)
61 .ok_or_else(|| DomainAgentError::KnowledgeNotAvailable(domain.to_string()))
62 }
63
64 pub fn add_knowledge_entry(&mut self, domain: &str, entry: KnowledgeEntry) -> Result<()> {
66 let kb = self.get_knowledge_base_mut(domain)?;
67 kb.add_entry(entry);
68 Ok(())
69 }
70
71 pub fn search_by_tag(&self, domain: &str, tag: &str) -> Result<Vec<&KnowledgeEntry>> {
73 let kb = self.get_knowledge_base(domain)?;
74 Ok(kb.search_by_tag(tag))
75 }
76
77 pub fn search_by_category(&self, domain: &str, category: &str) -> Result<Vec<&KnowledgeEntry>> {
79 let kb = self.get_knowledge_base(domain)?;
80 Ok(kb.search_by_category(category))
81 }
82
83 pub async fn save_knowledge_base(&self, domain: &str) -> Result<()> {
85 let kb = self.get_knowledge_base(domain)?;
86 let kb_path = self.get_knowledge_base_path(domain);
87
88 if let Some(parent) = kb_path.parent() {
90 tokio::fs::create_dir_all(parent).await?;
91 }
92
93 let content = serde_yaml::to_string(kb)?;
94 tokio::fs::write(&kb_path, content).await?;
95 info!("Saved knowledge base for domain: {}", domain);
96
97 Ok(())
98 }
99
100 fn get_knowledge_base_path(&self, domain: &str) -> std::path::PathBuf {
102 PathResolver::resolve_project_path()
103 .join("knowledge_bases")
104 .join(format!("{}.yaml", domain))
105 }
106
107 pub async fn load_all_knowledge_bases(&mut self) -> Result<()> {
109 debug!("Loading all knowledge bases");
110
111 for domain in Domain::all() {
112 self.load_knowledge_base(domain.as_str()).await?;
113 }
114
115 info!("Loaded all knowledge bases");
116 Ok(())
117 }
118
119 pub fn get_loaded_domains(&self) -> Vec<&str> {
121 self.knowledge_bases.keys().map(|s| s.as_str()).collect()
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 #[test]
130 fn test_knowledge_base_creation() {
131 let kb = KnowledgeBase::new("frontend", "1.0.0");
132 assert_eq!(kb.domain, "frontend");
133 assert_eq!(kb.version, "1.0.0");
134 assert!(kb.entries.is_empty());
135 }
136
137 #[test]
138 fn test_add_knowledge_entry() {
139 let mut kb = KnowledgeBase::new("frontend", "1.0.0");
140 let entry = KnowledgeEntry {
141 id: "react-hooks".to_string(),
142 category: "patterns".to_string(),
143 title: "React Hooks".to_string(),
144 description: "Best practices for using React Hooks".to_string(),
145 tags: vec!["react".to_string(), "hooks".to_string()],
146 example: Some("const [count, setCount] = useState(0);".to_string()),
147 references: vec!["https://react.dev/reference/react/hooks".to_string()],
148 };
149
150 kb.add_entry(entry);
151 assert_eq!(kb.entries.len(), 1);
152 }
153
154 #[test]
155 fn test_search_by_tag() {
156 let mut kb = KnowledgeBase::new("frontend", "1.0.0");
157 let entry = KnowledgeEntry {
158 id: "react-hooks".to_string(),
159 category: "patterns".to_string(),
160 title: "React Hooks".to_string(),
161 description: "Best practices for using React Hooks".to_string(),
162 tags: vec!["react".to_string(), "hooks".to_string()],
163 example: None,
164 references: vec![],
165 };
166
167 kb.add_entry(entry);
168 let results = kb.search_by_tag("react");
169 assert_eq!(results.len(), 1);
170 }
171
172 #[test]
173 fn test_search_by_category() {
174 let mut kb = KnowledgeBase::new("frontend", "1.0.0");
175 let entry = KnowledgeEntry {
176 id: "react-hooks".to_string(),
177 category: "patterns".to_string(),
178 title: "React Hooks".to_string(),
179 description: "Best practices for using React Hooks".to_string(),
180 tags: vec![],
181 example: None,
182 references: vec![],
183 };
184
185 kb.add_entry(entry);
186 let results = kb.search_by_category("patterns");
187 assert_eq!(results.len(), 1);
188 }
189}