1pub mod confidence;
7pub mod crud;
8pub mod dedup;
9pub mod embed;
10pub mod error;
11pub mod extract;
12pub mod gc;
13pub mod ingest;
14pub mod llm;
15pub mod pipeline;
16pub mod pipeline_sync;
17pub mod query;
18pub mod search;
19pub mod store;
20pub mod traverse;
21pub mod types;
22pub mod util;
23pub mod utility;
24pub mod vigil_sync;
25
26use std::collections::HashMap;
27use std::path::{Path, PathBuf};
28
29use embed::{FastEmbedder, LazyEmbedder};
30use error::GraphError;
31use store::Db;
32pub use store::ServerConfig;
33#[allow(unused_imports)] use surrealdb::types::SurrealValue;
35use surrealdb::Surreal;
36use types::*;
37
38pub(crate) fn deserialize_take<T: serde::de::DeserializeOwned>(
41 response: &mut surrealdb::IndexedResults,
42 index: usize,
43) -> Result<Vec<T>, GraphError> {
44 let values: Vec<serde_json::Value> = response.take(index)?;
45 values
46 .into_iter()
47 .map(|v| serde_json::from_value(v).map_err(GraphError::from))
48 .collect()
49}
50
51pub(crate) fn deserialize_take_opt<T: serde::de::DeserializeOwned>(
52 response: &mut surrealdb::IndexedResults,
53 index: usize,
54) -> Result<Option<T>, GraphError> {
55 let values: Vec<T> = deserialize_take(response, index)?;
56 Ok(values.into_iter().next())
57}
58
59pub struct GraphMemory {
61 db: Surreal<Db>,
62 embedder: LazyEmbedder,
63 path: PathBuf,
64 scoring: crate::config::GraphScoringConfig,
65}
66
67impl GraphMemory {
68 pub async fn open(path: &Path) -> Result<Self, GraphError> {
76 let memory_dir = path.parent().unwrap_or(path);
77 let config = crate::config::load_from_dir(memory_dir);
78 let mode = config
79 .graph
80 .as_ref()
81 .map(|g| g.mode.clone())
82 .unwrap_or_else(|| "embedded".to_string());
83
84 match mode.as_str() {
85 "server" => Self::open_server(path).await,
86 _ => Self::open_embedded(path).await,
87 }
88 }
89
90 pub async fn open_embedded(path: &Path) -> Result<Self, GraphError> {
92 std::fs::create_dir_all(path)?;
93
94 let db = store::open(path).await?;
95 store::init_schema(&db).await?;
96
97 let models_dir = path.join("models");
98 std::fs::create_dir_all(&models_dir)?;
99 let embedder = LazyEmbedder::new(&models_dir);
100
101 let scoring = load_scoring_config(path);
102
103 Ok(Self {
104 db,
105 embedder,
106 path: path.to_path_buf(),
107 scoring,
108 })
109 }
110
111 pub async fn open_server(path: &Path) -> Result<Self, GraphError> {
115 let memory_dir = path.parent().unwrap_or(path);
116 let config = crate::config::load_from_dir(memory_dir);
117
118 let graph_section = config.graph.unwrap_or_default();
119 let password = if graph_section.password_file.is_empty() {
120 String::new()
121 } else {
122 let pw_path = if graph_section.password_file.starts_with('/') {
123 std::path::PathBuf::from(&graph_section.password_file)
124 } else {
125 let entity_root = memory_dir.parent().unwrap_or(memory_dir);
127 entity_root.join(&graph_section.password_file)
128 };
129 std::fs::read_to_string(&pw_path)
130 .map(|s| s.trim().to_string())
131 .map_err(|e| {
132 GraphError::Io(std::io::Error::new(
133 e.kind(),
134 format!(
135 "failed to read graph password file {}: {e}",
136 pw_path.display()
137 ),
138 ))
139 })?
140 };
141
142 let scoring = graph_section.scoring.clone();
143 let server_config = store::ServerConfig {
144 url: graph_section.url,
145 username: graph_section.username,
146 password,
147 namespace: graph_section.namespace,
148 database: graph_section.database,
149 };
150
151 let models_dir = path.join("models");
152 let mut gm = Self::connect(&server_config, &models_dir).await?;
153 gm.scoring = scoring;
154 Ok(gm)
155 }
156
157 pub async fn connect(
159 config: &store::ServerConfig,
160 models_dir: &Path,
161 ) -> Result<Self, GraphError> {
162 let db = store::connect(config).await?;
163 store::init_schema(&db).await?;
164
165 std::fs::create_dir_all(models_dir)?;
166 let embedder = LazyEmbedder::new(models_dir);
167
168 Ok(Self {
169 db,
170 embedder,
171 path: models_dir.to_path_buf(),
172 scoring: crate::config::GraphScoringConfig::default(),
173 })
174 }
175
176 pub fn path(&self) -> &Path {
178 &self.path
179 }
180
181 #[allow(dead_code)]
183 pub(crate) fn db(&self) -> &Surreal<Db> {
184 &self.db
185 }
186
187 #[allow(dead_code)]
189 pub(crate) fn embedder(&self) -> Result<&FastEmbedder, GraphError> {
190 self.embedder.get()
191 }
192
193 pub async fn add_entity(&self, entity: NewEntity) -> Result<Entity, GraphError> {
197 crud::add_entity(&self.db, self.embedder.get()?, entity).await
198 }
199
200 pub async fn get_entity(&self, name: &str) -> Result<Option<Entity>, GraphError> {
202 crud::get_entity_by_name(&self.db, name).await
203 }
204
205 pub async fn get_entity_by_id(&self, id: &str) -> Result<Option<Entity>, GraphError> {
207 crud::get_entity_by_id(&self.db, id).await
208 }
209
210 pub async fn update_entity(
212 &self,
213 id: &str,
214 updates: EntityUpdate,
215 ) -> Result<Entity, GraphError> {
216 crud::update_entity(&self.db, self.embedder.get()?, id, updates).await
217 }
218
219 pub async fn delete_entity(&self, id: &str) -> Result<(), GraphError> {
221 crud::delete_entity(&self.db, id).await
222 }
223
224 pub async fn list_entities(
226 &self,
227 entity_type: Option<&str>,
228 ) -> Result<Vec<Entity>, GraphError> {
229 crud::list_entities(&self.db, entity_type).await
230 }
231
232 pub async fn add_relationship(&self, rel: NewRelationship) -> Result<Relationship, GraphError> {
236 crud::add_relationship(&self.db, rel).await
237 }
238
239 pub async fn get_relationships(
241 &self,
242 entity_name: &str,
243 direction: Direction,
244 ) -> Result<Vec<Relationship>, GraphError> {
245 crud::get_relationships(&self.db, entity_name, direction).await
246 }
247
248 pub async fn supersede_relationship(
250 &self,
251 old_id: &str,
252 new: NewRelationship,
253 ) -> Result<Relationship, GraphError> {
254 crud::supersede_relationship(&self.db, old_id, new).await
255 }
256
257 pub async fn update_relationship_confidence(
259 &self,
260 rel_id: &str,
261 confidence: f64,
262 ) -> Result<(), GraphError> {
263 crud::update_relationship_confidence(&self.db, rel_id, confidence).await
264 }
265
266 pub async fn reinforce_relationship(
271 &self,
272 rel_id: &str,
273 new_confidence: f64,
274 ) -> Result<(), GraphError> {
275 crud::reinforce_relationship(&self.db, rel_id, new_confidence).await
276 }
277
278 pub async fn add_episode(&self, episode: NewEpisode) -> Result<Episode, GraphError> {
282 crud::add_episode(&self.db, self.embedder.get()?, episode).await
283 }
284
285 pub async fn get_episodes_by_session(
287 &self,
288 session_id: &str,
289 ) -> Result<Vec<Episode>, GraphError> {
290 crud::get_episodes_by_session(&self.db, session_id).await
291 }
292
293 pub async fn get_episode_by_log_number(
295 &self,
296 log_number: u32,
297 ) -> Result<Option<Episode>, GraphError> {
298 crud::get_episode_by_log_number(&self.db, log_number).await
299 }
300
301 pub async fn ingest_archive(
305 &self,
306 archive_text: &str,
307 session_id: &str,
308 log_number: Option<u32>,
309 llm: Option<&dyn llm::LlmProvider>,
310 ) -> Result<IngestionReport, GraphError> {
311 ingest::ingest_archive(self, archive_text, session_id, log_number, llm).await
312 }
313
314 pub async fn extract_from_archive(
316 &self,
317 archive_text: &str,
318 session_id: &str,
319 log_number: Option<u32>,
320 llm: &dyn llm::LlmProvider,
321 ) -> Result<IngestionReport, GraphError> {
322 ingest::extract_from_archive(self, archive_text, session_id, log_number, llm).await
323 }
324
325 pub async fn mark_extracted(&self, log_number: u32) -> Result<(), GraphError> {
327 crud::mark_episodes_extracted(&self.db, log_number).await
328 }
329
330 pub async fn unextracted_log_numbers(&self) -> Result<Vec<i64>, GraphError> {
332 crud::get_unextracted_log_numbers(&self.db).await
333 }
334
335 pub async fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>, GraphError> {
339 search::search(&self.db, self.embedder.get()?, &self.scoring, query, limit).await
340 }
341
342 pub async fn search_with_options(
344 &self,
345 query: &str,
346 options: &SearchOptions,
347 ) -> Result<Vec<ScoredEntity>, GraphError> {
348 search::search_with_options(
349 &self.db,
350 self.embedder.get()?,
351 &self.scoring,
352 query,
353 options,
354 )
355 .await
356 }
357
358 pub async fn search_episodes(
360 &self,
361 query: &str,
362 limit: usize,
363 ) -> Result<Vec<EpisodeSearchResult>, GraphError> {
364 search::search_episodes(&self.db, self.embedder.get()?, query, limit).await
365 }
366
367 pub async fn query(
371 &self,
372 query_text: &str,
373 options: &QueryOptions,
374 ) -> Result<QueryResult, GraphError> {
375 query::query(
376 &self.db,
377 self.embedder.get()?,
378 &self.scoring,
379 query_text,
380 options,
381 )
382 .await
383 }
384
385 pub async fn traverse(
389 &self,
390 entity_name: &str,
391 depth: u32,
392 ) -> Result<TraversalNode, GraphError> {
393 traverse::traverse(&self.db, entity_name, depth).await
394 }
395
396 pub async fn traverse_filtered(
398 &self,
399 entity_name: &str,
400 depth: u32,
401 type_filter: Option<&str>,
402 ) -> Result<TraversalNode, GraphError> {
403 traverse::traverse_filtered(&self.db, entity_name, depth, type_filter).await
404 }
405
406 pub async fn sync_pipeline(
410 &self,
411 docs: &PipelineDocuments,
412 ) -> Result<PipelineSyncReport, GraphError> {
413 pipeline_sync::sync_pipeline(self, docs).await
414 }
415
416 pub async fn pipeline_stats(
418 &self,
419 staleness_days: u32,
420 ) -> Result<PipelineGraphStats, GraphError> {
421 query::pipeline_stats(&self.db, staleness_days).await
422 }
423
424 pub async fn pipeline_entities(
426 &self,
427 stage: &str,
428 status: Option<&str>,
429 ) -> Result<Vec<EntityDetail>, GraphError> {
430 query::pipeline_entities(&self.db, stage, status).await
431 }
432
433 pub async fn pipeline_flow(
435 &self,
436 entity_name: &str,
437 ) -> Result<Vec<(EntityDetail, String, EntityDetail)>, GraphError> {
438 query::pipeline_flow(&self.db, entity_name).await
439 }
440
441 pub async fn sync_vigil_signals(
445 &self,
446 signals_path: &std::path::Path,
447 ) -> Result<VigilSyncReport, GraphError> {
448 vigil_sync::sync_vigil_signals(self, signals_path).await
449 }
450
451 pub async fn sync_outcomes(
453 &self,
454 outcomes_path: &std::path::Path,
455 ) -> Result<VigilSyncReport, GraphError> {
456 vigil_sync::sync_outcomes(self, outcomes_path).await
457 }
458
459 pub async fn sync_vigil(
461 &self,
462 signals_path: &std::path::Path,
463 outcomes_path: &std::path::Path,
464 ) -> Result<VigilSyncReport, GraphError> {
465 vigil_sync::sync_vigil(self, signals_path, outcomes_path).await
466 }
467
468 pub async fn record_outcome_feedback(
473 &self,
474 session_id: &str,
475 outcome: utility::OutcomeKind,
476 retrieved_entity_ids: &[String],
477 used_entity_ids: Option<&[String]>,
478 ) -> Result<utility::FeedbackReport, GraphError> {
479 utility::record_outcome_feedback(
480 &self.db,
481 session_id,
482 outcome,
483 retrieved_entity_ids,
484 used_entity_ids,
485 )
486 .await
487 }
488
489 pub async fn run_gc(&self, config: &gc::GcConfig) -> Result<gc::GcReport, GraphError> {
493 gc::run_gc(&self.db, config).await
494 }
495
496 pub async fn gc_stats(&self) -> Result<gc::GcStatsReport, GraphError> {
498 gc::stats_only(&self.db).await
499 }
500
501 pub async fn delete_relationship(&self, id: &str) -> Result<(), GraphError> {
503 crud::delete_relationship(&self.db, id).await
504 }
505
506 pub async fn stats(&self) -> Result<GraphStats, GraphError> {
510 let entity_count = db_count(&self.db, "entity").await?;
511 let relationship_count = db_count(&self.db, "relates_to").await?;
512 let episode_count = db_count(&self.db, "episode").await?;
513
514 let mut type_response = self
516 .db
517 .query("SELECT entity_type, count() AS count FROM entity GROUP BY entity_type")
518 .await?;
519
520 let type_rows: Vec<TypeCount> = type_response.take(0)?;
521 let entity_type_counts: HashMap<String, u64> = type_rows
522 .into_iter()
523 .map(|r| (r.entity_type, r.count))
524 .collect();
525
526 Ok(GraphStats {
527 entity_count,
528 relationship_count,
529 episode_count,
530 entity_type_counts,
531 })
532 }
533}
534
535fn load_scoring_config(graph_path: &Path) -> crate::config::GraphScoringConfig {
539 let memory_dir = graph_path.parent().unwrap_or(graph_path);
540 crate::config::load_from_dir(memory_dir)
541 .graph
542 .map(|g| g.scoring)
543 .unwrap_or_default()
544}
545
546async fn db_count(db: &Surreal<Db>, table: &str) -> Result<u64, GraphError> {
547 let query = format!("SELECT count() AS count FROM {table} GROUP ALL");
548 let mut response = db.query(&query).await?;
549 let rows: Vec<CountRow> = response.take(0)?;
550 Ok(rows.first().map(|r| r.count).unwrap_or(0))
551}
552
553#[derive(serde::Deserialize, surrealdb::types::SurrealValue)]
554struct CountRow {
555 count: u64,
556}
557
558#[derive(serde::Deserialize, surrealdb::types::SurrealValue)]
559struct TypeCount {
560 entity_type: String,
561 count: u64,
562}