pub mod confidence;
pub mod crud;
pub mod dedup;
pub mod embed;
pub mod error;
pub mod extract;
pub mod gc;
pub mod ingest;
pub mod llm;
pub mod pipeline;
pub mod pipeline_sync;
pub mod query;
pub mod search;
pub mod store;
pub mod traverse;
pub mod types;
pub mod util;
pub mod utility;
pub mod vigil_sync;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
pub use confidence::{Provenance, ProvenanceWeights};
use embed::{FastEmbedder, LazyEmbedder};
use error::GraphError;
pub use ingest::{IngestContext, ProvenancePolicy};
use store::Db;
pub use store::ServerConfig;
#[allow(unused_imports)] use surrealdb::types::SurrealValue;
use surrealdb::Surreal;
use types::*;
pub(crate) fn deserialize_take<T: serde::de::DeserializeOwned>(
response: &mut surrealdb::IndexedResults,
index: usize,
) -> Result<Vec<T>, GraphError> {
let values: Vec<serde_json::Value> = response.take(index)?;
values
.into_iter()
.map(|v| serde_json::from_value(v).map_err(GraphError::from))
.collect()
}
pub(crate) fn deserialize_take_opt<T: serde::de::DeserializeOwned>(
response: &mut surrealdb::IndexedResults,
index: usize,
) -> Result<Option<T>, GraphError> {
let values: Vec<T> = deserialize_take(response, index)?;
Ok(values.into_iter().next())
}
pub struct GraphMemory {
db: Surreal<Db>,
embedder: LazyEmbedder,
path: PathBuf,
scoring: crate::config::GraphScoringConfig,
provenance: confidence::ProvenanceWeights,
}
impl GraphMemory {
pub async fn open(path: &Path) -> Result<Self, GraphError> {
let memory_dir = path.parent().unwrap_or(path);
let config = crate::config::load_from_dir(memory_dir);
let mode = config
.graph
.as_ref()
.map(|g| g.mode.clone())
.unwrap_or_else(|| "embedded".to_string());
match mode.as_str() {
"server" => Self::open_server(path).await,
_ => Self::open_embedded(path).await,
}
}
pub async fn open_embedded(path: &Path) -> Result<Self, GraphError> {
std::fs::create_dir_all(path)?;
let db = store::open(path).await?;
store::init_schema(&db).await?;
let models_dir = path.join("models");
std::fs::create_dir_all(&models_dir)?;
let embedder = LazyEmbedder::new(&models_dir);
let graph_config = load_graph_section(path);
Ok(Self {
db,
embedder,
path: path.to_path_buf(),
scoring: graph_config.scoring,
provenance: graph_config.provenance,
})
}
pub async fn open_server(path: &Path) -> Result<Self, GraphError> {
let memory_dir = path.parent().unwrap_or(path);
let config = crate::config::load_from_dir(memory_dir);
let graph_section = config.graph.unwrap_or_default();
let password = if graph_section.password_file.is_empty() {
String::new()
} else {
let pw_path = if graph_section.password_file.starts_with('/') {
std::path::PathBuf::from(&graph_section.password_file)
} else {
let entity_root = memory_dir.parent().unwrap_or(memory_dir);
entity_root.join(&graph_section.password_file)
};
std::fs::read_to_string(&pw_path)
.map(|s| s.trim().to_string())
.map_err(|e| {
GraphError::Io(std::io::Error::new(
e.kind(),
format!(
"failed to read graph password file {}: {e}",
pw_path.display()
),
))
})?
};
let scoring = graph_section.scoring.clone();
let provenance = graph_section.provenance;
let server_config = store::ServerConfig {
url: graph_section.url,
username: graph_section.username,
password,
namespace: graph_section.namespace,
database: graph_section.database,
};
let models_dir = path.join("models");
let mut gm = Self::connect(&server_config, &models_dir).await?;
gm.scoring = scoring;
gm.provenance = provenance;
Ok(gm)
}
pub async fn connect(
config: &store::ServerConfig,
models_dir: &Path,
) -> Result<Self, GraphError> {
let db = store::connect(config).await?;
store::init_schema(&db).await?;
std::fs::create_dir_all(models_dir)?;
let embedder = LazyEmbedder::new(models_dir);
Ok(Self {
db,
embedder,
path: models_dir.to_path_buf(),
scoring: crate::config::GraphScoringConfig::default(),
provenance: confidence::ProvenanceWeights::default(),
})
}
pub fn path(&self) -> &Path {
&self.path
}
#[must_use]
pub fn provenance_weights(&self) -> &confidence::ProvenanceWeights {
&self.provenance
}
#[allow(dead_code)]
pub(crate) fn db(&self) -> &Surreal<Db> {
&self.db
}
#[allow(dead_code)]
pub(crate) fn embedder(&self) -> Result<&FastEmbedder, GraphError> {
self.embedder.get()
}
pub async fn add_entity(&self, entity: NewEntity) -> Result<Entity, GraphError> {
crud::add_entity(&self.db, self.embedder.get()?, entity).await
}
pub async fn get_entity(&self, name: &str) -> Result<Option<Entity>, GraphError> {
crud::get_entity_by_name(&self.db, name).await
}
pub async fn get_entity_by_id(&self, id: &str) -> Result<Option<Entity>, GraphError> {
crud::get_entity_by_id(&self.db, id).await
}
pub async fn update_entity(
&self,
id: &str,
updates: EntityUpdate,
) -> Result<Entity, GraphError> {
crud::update_entity(&self.db, self.embedder.get()?, id, updates).await
}
pub async fn delete_entity(&self, id: &str) -> Result<(), GraphError> {
crud::delete_entity(&self.db, id).await
}
pub async fn list_entities(
&self,
entity_type: Option<&str>,
) -> Result<Vec<Entity>, GraphError> {
crud::list_entities(&self.db, entity_type).await
}
pub async fn add_relationship(&self, rel: NewRelationship) -> Result<Relationship, GraphError> {
crud::add_relationship(&self.db, rel).await
}
pub async fn get_relationships(
&self,
entity_name: &str,
direction: Direction,
) -> Result<Vec<Relationship>, GraphError> {
crud::get_relationships(&self.db, entity_name, direction).await
}
pub async fn supersede_relationship(
&self,
old_id: &str,
new: NewRelationship,
) -> Result<Relationship, GraphError> {
crud::supersede_relationship(&self.db, old_id, new).await
}
pub async fn update_relationship_confidence(
&self,
rel_id: &str,
confidence: f64,
) -> Result<(), GraphError> {
crud::update_relationship_confidence(&self.db, rel_id, confidence).await
}
pub async fn reinforce_relationship(
&self,
rel_id: &str,
evidence: confidence::EdgeEvidence,
) -> Result<(), GraphError> {
crud::reinforce_relationship(&self.db, rel_id, evidence).await
}
pub async fn add_episode(&self, episode: NewEpisode) -> Result<Episode, GraphError> {
crud::add_episode(&self.db, self.embedder.get()?, episode).await
}
pub async fn add_episode_from(
&self,
episode: NewEpisode,
provenance: Provenance,
) -> Result<Episode, GraphError> {
crud::add_episode_from(&self.db, self.embedder.get()?, episode, provenance).await
}
pub async fn get_episodes_by_session(
&self,
session_id: &str,
) -> Result<Vec<Episode>, GraphError> {
crud::get_episodes_by_session(&self.db, session_id).await
}
pub async fn get_episode_by_log_number(
&self,
log_number: u32,
) -> Result<Option<Episode>, GraphError> {
crud::get_episode_by_log_number(&self.db, log_number).await
}
pub async fn ingest_archive(
&self,
archive_text: &str,
context: &IngestContext,
llm: Option<&dyn llm::LlmProvider>,
) -> Result<IngestionReport, GraphError> {
ingest::ingest_archive(self, archive_text, context, llm).await
}
pub async fn extract_from_archive(
&self,
archive_text: &str,
context: &IngestContext,
llm: &dyn llm::LlmProvider,
) -> Result<IngestionReport, GraphError> {
ingest::extract_from_archive(self, archive_text, context, llm).await
}
pub async fn mark_extracted(&self, log_number: u32) -> Result<(), GraphError> {
crud::mark_episodes_extracted(&self.db, log_number).await
}
pub async fn unextracted_log_numbers(&self) -> Result<Vec<i64>, GraphError> {
crud::get_unextracted_log_numbers(&self.db).await
}
pub async fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>, GraphError> {
search::search(&self.db, self.embedder.get()?, &self.scoring, query, limit).await
}
pub async fn search_with_options(
&self,
query: &str,
options: &SearchOptions,
) -> Result<Vec<ScoredEntity>, GraphError> {
search::search_with_options(
&self.db,
self.embedder.get()?,
&self.scoring,
query,
options,
)
.await
}
pub async fn search_episodes(
&self,
query: &str,
limit: usize,
) -> Result<Vec<EpisodeSearchResult>, GraphError> {
search::search_episodes(&self.db, self.embedder.get()?, query, limit).await
}
pub async fn query(
&self,
query_text: &str,
options: &QueryOptions,
) -> Result<QueryResult, GraphError> {
query::query(
&self.db,
self.embedder.get()?,
&self.scoring,
query_text,
options,
)
.await
}
pub async fn traverse(
&self,
entity_name: &str,
depth: u32,
) -> Result<TraversalNode, GraphError> {
traverse::traverse(&self.db, entity_name, depth).await
}
pub async fn traverse_filtered(
&self,
entity_name: &str,
depth: u32,
type_filter: Option<&str>,
) -> Result<TraversalNode, GraphError> {
traverse::traverse_filtered(&self.db, entity_name, depth, type_filter).await
}
pub async fn sync_pipeline(
&self,
docs: &PipelineDocuments,
) -> Result<PipelineSyncReport, GraphError> {
pipeline_sync::sync_pipeline(self, docs).await
}
pub async fn pipeline_stats(
&self,
staleness_days: u32,
) -> Result<PipelineGraphStats, GraphError> {
query::pipeline_stats(&self.db, staleness_days).await
}
pub async fn pipeline_entities(
&self,
stage: &str,
status: Option<&str>,
) -> Result<Vec<EntityDetail>, GraphError> {
query::pipeline_entities(&self.db, stage, status).await
}
pub async fn pipeline_flow(
&self,
entity_name: &str,
) -> Result<Vec<(EntityDetail, String, EntityDetail)>, GraphError> {
query::pipeline_flow(&self.db, entity_name).await
}
pub async fn sync_vigil_signals(
&self,
signals_path: &std::path::Path,
) -> Result<VigilSyncReport, GraphError> {
vigil_sync::sync_vigil_signals(self, signals_path).await
}
pub async fn sync_outcomes(
&self,
outcomes_path: &std::path::Path,
) -> Result<VigilSyncReport, GraphError> {
vigil_sync::sync_outcomes(self, outcomes_path).await
}
pub async fn sync_vigil(
&self,
signals_path: &std::path::Path,
outcomes_path: &std::path::Path,
) -> Result<VigilSyncReport, GraphError> {
vigil_sync::sync_vigil(self, signals_path, outcomes_path).await
}
pub async fn record_outcome_feedback(
&self,
session_id: &str,
outcome: utility::OutcomeKind,
retrieved_entity_ids: &[String],
used_entity_ids: Option<&[String]>,
) -> Result<utility::FeedbackReport, GraphError> {
utility::record_outcome_feedback(
&self.db,
session_id,
outcome,
retrieved_entity_ids,
used_entity_ids,
)
.await
}
pub async fn record_session_outcome(
&self,
session_id: &str,
outcome: utility::OutcomeKind,
) -> Result<utility::FeedbackReport, GraphError> {
let session = utility::session_entities(&self.db, session_id).await?;
if session.is_empty() {
return Ok(utility::FeedbackReport::default());
}
utility::record_outcome_feedback(
&self.db,
session_id,
outcome,
&session.retrieved,
Some(&session.used),
)
.await
}
pub async fn record_session_use(
&self,
session_id: &str,
entity_ids: &[String],
) -> Result<u32, GraphError> {
utility::record_session_use(&self.db, session_id, entity_ids).await
}
pub async fn run_gc(&self, config: &gc::GcConfig) -> Result<gc::GcReport, GraphError> {
gc::run_gc(&self.db, config).await
}
pub async fn gc_stats(&self) -> Result<gc::GcStatsReport, GraphError> {
gc::stats_only(&self.db).await
}
pub async fn delete_relationship(&self, id: &str) -> Result<(), GraphError> {
crud::delete_relationship(&self.db, id).await
}
pub async fn stats(&self) -> Result<GraphStats, GraphError> {
let entity_count = db_count(&self.db, "entity").await?;
let relationship_count = db_count(&self.db, "relates_to").await?;
let episode_count = db_count(&self.db, "episode").await?;
let mut type_response = self
.db
.query("SELECT entity_type, count() AS count FROM entity GROUP BY entity_type")
.await?;
let type_rows: Vec<TypeCount> = type_response.take(0)?;
let entity_type_counts: HashMap<String, u64> = type_rows
.into_iter()
.map(|r| (r.entity_type, r.count))
.collect();
Ok(GraphStats {
entity_count,
relationship_count,
episode_count,
entity_type_counts,
})
}
}
fn load_graph_section(graph_path: &Path) -> crate::config::GraphSection {
let memory_dir = graph_path.parent().unwrap_or(graph_path);
crate::config::load_from_dir(memory_dir)
.graph
.unwrap_or_default()
}
async fn db_count(db: &Surreal<Db>, table: &str) -> Result<u64, GraphError> {
let query = format!("SELECT count() AS count FROM {table} GROUP ALL");
let mut response = db.query(&query).await?;
let rows: Vec<CountRow> = response.take(0)?;
Ok(rows.first().map(|r| r.count).unwrap_or(0))
}
#[derive(serde::Deserialize, surrealdb::types::SurrealValue)]
struct CountRow {
count: u64,
}
#[derive(serde::Deserialize, surrealdb::types::SurrealValue)]
struct TypeCount {
entity_type: String,
count: u64,
}