1use crate::connectors::types::{ObjectDescriptor, ProtoData};
2use crate::ids::{EntityId, ObjectId, ProviderId, RelationshipId};
3use crate::knowledge::{Entity, Relationship};
4use crate::status::UnderstandingStatus;
5use serde::{Deserialize, Serialize};
6use std::collections::{BTreeMap, HashMap, HashSet};
7
8#[derive(Clone, Debug)]
9pub struct StoredObject {
10 pub descriptor: ObjectDescriptor,
11 pub proto: ProtoData,
12 pub status: UnderstandingStatus,
13 pub provider: Option<ProviderId>,
14 pub classification_reason: Option<String>,
15 pub entity_ids: Vec<EntityId>,
16 pub diagnostics: Vec<String>,
17}
18
19#[derive(Clone, Debug, Default, Serialize, Deserialize)]
20pub struct InventoryCounts {
21 pub source_objects: u64,
22 pub modules: u64,
23 pub types: u64,
24 pub functions: u64,
25 pub documents: u64,
26 pub datasets: u64,
27 pub relationships: u64,
28 pub understood: u64,
29 pub partial: u64,
30 pub unknown: u64,
31 pub failed: u64,
32}
33
34#[derive(Debug, Default)]
35pub struct KnowledgeStore {
36 objects: HashMap<ObjectId, StoredObject>,
37 entities: HashMap<EntityId, Entity>,
38 relationships: HashMap<RelationshipId, Relationship>,
39 by_name: HashMap<String, HashSet<EntityId>>,
40 by_kind: HashMap<String, HashSet<EntityId>>,
41 by_path: HashMap<String, ObjectId>,
42 reverse_rels: HashMap<EntityId, HashSet<RelationshipId>>,
43 forward_rels: HashMap<EntityId, HashSet<RelationshipId>>,
44 root: Option<std::path::PathBuf>,
45 git_branch: Option<String>,
46 enabled_providers: Vec<String>,
47}
48
49impl KnowledgeStore {
50 pub fn new() -> Self {
51 Self::default()
52 }
53
54 pub fn set_root(&mut self, root: std::path::PathBuf) {
55 self.root = Some(root);
56 }
57
58 pub fn root(&self) -> Option<&std::path::PathBuf> {
59 self.root.as_ref()
60 }
61
62 pub fn set_git_branch(&mut self, branch: Option<String>) {
63 self.git_branch = branch;
64 }
65
66 pub fn git_branch(&self) -> Option<&str> {
67 self.git_branch.as_deref()
68 }
69
70 pub fn set_enabled_providers(&mut self, providers: Vec<String>) {
71 self.enabled_providers = providers;
72 }
73
74 pub fn enabled_providers(&self) -> &[String] {
75 &self.enabled_providers
76 }
77
78 pub fn upsert_object(&mut self, object: StoredObject) {
79 let path_key = object
80 .descriptor
81 .relative_path
82 .to_string_lossy()
83 .to_string();
84 self.by_path.insert(path_key, object.descriptor.id.clone());
85 self.objects.insert(object.descriptor.id.clone(), object);
86 }
87
88 pub fn add_entity(&mut self, entity: Entity) {
89 let id = entity.id.clone();
90 self.by_name
91 .entry(entity.name.to_lowercase())
92 .or_default()
93 .insert(id.clone());
94 self.by_kind
95 .entry(entity.kind.clone())
96 .or_default()
97 .insert(id.clone());
98 if let Some(obj) = self.objects.get_mut(&entity.source_object) {
99 if !obj.entity_ids.contains(&id) {
100 obj.entity_ids.push(id.clone());
101 }
102 }
103 self.entities.insert(id, entity);
104 }
105
106 pub fn add_relationship(&mut self, rel: Relationship) {
107 let id = rel.id.clone();
108 self.forward_rels
109 .entry(rel.from.clone())
110 .or_default()
111 .insert(id.clone());
112 self.reverse_rels
113 .entry(rel.to.clone())
114 .or_default()
115 .insert(id.clone());
116 self.relationships.insert(id, rel);
117 }
118
119 pub fn get_object(&self, id: &ObjectId) -> Option<&StoredObject> {
120 self.objects.get(id)
121 }
122
123 pub fn get_object_by_path(&self, path: &str) -> Option<&StoredObject> {
124 self.by_path.get(path).and_then(|id| self.objects.get(id))
125 }
126
127 pub fn objects(&self) -> impl Iterator<Item = &StoredObject> {
128 self.objects.values()
129 }
130
131 pub fn objects_by_status(&self, status: UnderstandingStatus) -> Vec<&StoredObject> {
132 self.objects
133 .values()
134 .filter(|o| o.status == status)
135 .collect()
136 }
137
138 pub fn get_entity(&self, id: &EntityId) -> Option<&Entity> {
139 self.entities.get(id)
140 }
141
142 pub fn entities(&self) -> impl Iterator<Item = &Entity> {
143 self.entities.values()
144 }
145
146 pub fn relationships(&self) -> impl Iterator<Item = &Relationship> {
147 self.relationships.values()
148 }
149
150 pub fn find_entities_by_name(&self, name: &str) -> Vec<&Entity> {
151 let key = name.to_lowercase();
152 self.by_name
153 .get(&key)
154 .into_iter()
155 .flatten()
156 .filter_map(|id| self.entities.get(id))
157 .collect()
158 }
159
160 pub fn find(&self, text: &str) -> FindResults {
161 let q = text.to_lowercase();
162 let mut entities = Vec::new();
163 for e in self.entities.values() {
164 if e.name.to_lowercase().contains(&q)
165 || e.kind.to_lowercase().contains(&q)
166 || e.attributes.values().any(|v| v.to_lowercase().contains(&q))
167 {
168 entities.push(e.clone());
169 }
170 }
171 let mut objects = Vec::new();
172 for o in self.objects.values() {
173 let path = o.descriptor.relative_path.to_string_lossy().to_lowercase();
174 if path.contains(&q)
175 || o.descriptor.media_type.to_lowercase().contains(&q)
176 || o.classification_reason
177 .as_ref()
178 .map(|r| r.to_lowercase().contains(&q))
179 .unwrap_or(false)
180 {
181 objects.push(o.clone());
182 }
183 }
184 entities.sort_by(|a, b| a.name.cmp(&b.name));
185 objects.sort_by(|a, b| a.descriptor.relative_path.cmp(&b.descriptor.relative_path));
186 FindResults { entities, objects }
187 }
188
189 pub fn neighborhood(&self, entity_id: &EntityId, depth: usize) -> GraphNeighborhood {
190 let mut nodes = BTreeMap::new();
191 let mut edges = Vec::new();
192 let mut frontier = vec![entity_id.clone()];
193 let mut seen = HashSet::new();
194 seen.insert(entity_id.clone());
195
196 if let Some(e) = self.entities.get(entity_id) {
197 nodes.insert(entity_id.clone(), e.clone());
198 }
199
200 for _ in 0..depth {
201 let mut next = Vec::new();
202 for id in &frontier {
203 let rel_ids: Vec<_> = self
204 .forward_rels
205 .get(id)
206 .into_iter()
207 .flatten()
208 .chain(self.reverse_rels.get(id).into_iter().flatten())
209 .cloned()
210 .collect();
211 for rid in rel_ids {
212 if let Some(rel) = self.relationships.get(&rid) {
213 edges.push(rel.clone());
214 for neighbor in [&rel.from, &rel.to] {
215 if seen.insert(neighbor.clone()) {
216 if let Some(e) = self.entities.get(neighbor) {
217 nodes.insert(neighbor.clone(), e.clone());
218 next.push(neighbor.clone());
219 }
220 }
221 }
222 }
223 }
224 }
225 frontier = next;
226 if frontier.is_empty() {
227 break;
228 }
229 }
230
231 GraphNeighborhood { nodes, edges }
232 }
233
234 pub fn inventory(&self) -> InventoryCounts {
235 let mut counts = InventoryCounts {
236 source_objects: self.objects.len() as u64,
237 relationships: self.relationships.len() as u64,
238 ..Default::default()
239 };
240 for o in self.objects.values() {
241 match o.status {
242 UnderstandingStatus::Understood => counts.understood += 1,
243 UnderstandingStatus::PartiallyUnderstood => counts.partial += 1,
244 UnderstandingStatus::Unknown => counts.unknown += 1,
245 UnderstandingStatus::Failed => counts.failed += 1,
246 }
247 }
248 for e in self.entities.values() {
249 match e.kind.as_str() {
250 "module" | "namespace" | "package" => counts.modules += 1,
251 "class" | "struct" | "enum" | "trait" | "protocol" | "type" => counts.types += 1,
252 "function" | "method" => counts.functions += 1,
253 "document" | "section" | "heading" => counts.documents += 1,
254 "dataset" | "column" => counts.datasets += 1,
255 _ => {}
256 }
257 }
258 counts
259 }
260}
261
262#[derive(Clone, Debug, Default)]
263pub struct FindResults {
264 pub entities: Vec<Entity>,
265 pub objects: Vec<StoredObject>,
266}
267
268#[derive(Clone, Debug, Default)]
269pub struct GraphNeighborhood {
270 pub nodes: BTreeMap<EntityId, Entity>,
271 pub edges: Vec<Relationship>,
272}