use crate::kg_write::{CachePolicy, KgWriteError};
use crate::{ActivityFilter, ActivitySource, DaemonEvent};
use std::sync::Arc;
use trusty_common::memory_core::dream::{DreamConfig, Dreamer, PersistedDreamStats};
use trusty_common::memory_core::palace::PalaceId;
use trusty_common::memory_core::store::kg::{ExpandDirection, Triple};
use trusty_common::memory_core::{PalaceHandle, PalaceRegistry};
use super::core::KG_GRAPH_MAX_TRIPLES;
use super::helpers::{list_palaces_blocking, refresh_gaps_cache};
use super::types::{
DreamStatusPayload, KgAssertBody, KgGraphPayload, KgNeighborsPayload, KgNodeView,
KgSeedPayload, ServiceError, ServiceResult,
};
use super::MemoryService;
async fn adjacency_counts(handle: &Arc<PalaceHandle>) -> (u64, u64, u64) {
let handle = Arc::clone(handle);
tokio::task::spawn_blocking(move || {
(
handle.kg.node_count() as u64,
handle.kg.edge_count() as u64,
handle.kg.community_count() as u64,
)
})
.await
.unwrap_or_else(|e| {
tracing::warn!("kg adjacency counts task failed: {e}");
(0, 0, 0)
})
}
pub(crate) const DEFAULT_KG_LIST_LIMIT: usize = 50;
pub(crate) const MAX_KG_LIST_LIMIT: usize = 200;
pub(crate) const MAINTENANCE_LEASE_NOT_HELD: &str =
"this process does not hold this data root's maintenance lease (another \
process holds it, or maintenance.lock could not be opened); see the pid in \
maintenance.lock (#8733)";
impl MemoryService {
pub async fn kg_query(&self, id: &str, subject: &str) -> ServiceResult<Vec<Triple>> {
let handle = self.open_handle(id)?;
handle
.kg
.query_active(subject)
.await
.map_err(|e| ServiceError::internal(format!("kg query: {e:#}")))
}
pub async fn kg_assert(&self, id: &str, body: KgAssertBody) -> ServiceResult<()> {
let handle = self.open_handle(id)?;
let triple = Triple {
subject: body.subject,
predicate: body.predicate,
object: body.object,
valid_from: chrono::Utc::now(),
valid_to: None,
confidence: body.confidence.unwrap_or(1.0),
provenance: body.provenance,
};
crate::kg_write::assert_triple(&self.state, &handle, triple, CachePolicy::Inline)
.await
.map(|_| ())
.map_err(|e| match e {
KgWriteError::Admission(inner) => ServiceError::bad_request(format!("{inner:#}")),
other => ServiceError::internal(format!("{other}")),
})
}
pub async fn kg_retract_triple(
&self,
id: &str,
subject: &str,
predicate: &str,
object: &str,
) -> ServiceResult<usize> {
let handle = self.open_handle(id)?;
let closed = handle
.kg
.retract_triple(subject, predicate, object)
.await
.map_err(|e| ServiceError::internal(format!("kg retract_triple: {e:#}")))?;
if closed > 0 && crate::prompt_facts::is_hot_predicate(predicate) {
if let Err(e) = crate::prompt_facts::rebuild_prompt_cache(&self.state).await {
tracing::warn!("rebuild_prompt_cache after kg_retract_triple failed: {e:#}");
}
}
Ok(closed)
}
pub async fn kg_list_subjects(&self, id: &str, limit: usize) -> ServiceResult<Vec<String>> {
let handle = self.open_handle(id)?;
handle
.kg
.list_subjects(limit)
.map_err(|e| ServiceError::internal(format!("kg list_subjects: {e:#}")))
}
pub async fn kg_list_subjects_with_counts(
&self,
id: &str,
limit: usize,
) -> ServiceResult<Vec<(String, u64)>> {
let handle = self.open_handle(id)?;
handle
.kg
.list_subjects_with_counts(limit)
.map_err(|e| ServiceError::internal(format!("kg list_subjects_with_counts: {e:#}")))
}
pub async fn kg_list_all(
&self,
id: &str,
limit: usize,
offset: usize,
) -> ServiceResult<Vec<Triple>> {
let handle = self.open_handle(id)?;
handle
.kg
.list_active(limit, offset)
.await
.map_err(|e| ServiceError::internal(format!("kg list_active: {e:#}")))
}
pub async fn kg_count(&self, id: &str) -> ServiceResult<usize> {
let handle = self.open_handle(id)?;
handle.kg.count_active_triples().map_err(|e| {
ServiceError::internal(format!("kg count_active_triples for palace {id}: {e:#}"))
})
}
pub async fn kg_graph(&self, id: &str) -> ServiceResult<KgGraphPayload> {
self.kg_graph_with_cap(id, KG_GRAPH_MAX_TRIPLES).await
}
pub async fn kg_graph_with_cap(
&self,
id: &str,
max_triples: usize,
) -> ServiceResult<KgGraphPayload> {
let handle = self.open_handle(id)?;
let triples = handle
.kg
.list_active(max_triples, 0)
.await
.map_err(|e| ServiceError::internal(format!("kg list_active: {e:#}")))?;
let active_triple_count = handle.kg.count_active_triples().map_err(|e| {
ServiceError::internal(format!("kg count_active_triples for palace {id}: {e:#}"))
})? as u64;
let returned_triple_count = triples.len() as u64;
let (node_count, edge_count, community_count) = adjacency_counts(&handle).await;
Ok(KgGraphPayload {
triples,
node_count,
edge_count,
community_count,
returned_triple_count,
active_triple_count,
truncated: returned_triple_count < active_triple_count,
})
}
pub async fn kg_graph_seed(&self, id: &str, limit: usize) -> ServiceResult<KgSeedPayload> {
let handle = self.open_handle(id)?;
let (nodes, triples) = handle
.kg
.top_degree_subgraph(limit)
.map_err(|e| ServiceError::internal(format!("kg top_degree_subgraph: {e:#}")))?;
let (node_count, edge_count, community_count) = adjacency_counts(&handle).await;
let returned_node_count = nodes.len() as u64;
Ok(KgSeedPayload {
nodes: nodes.into_iter().map(KgNodeView::from).collect(),
returned_triple_count: triples.len() as u64,
triples,
node_count,
edge_count,
community_count,
returned_node_count,
limit: limit as u64,
truncated: returned_node_count < node_count,
})
}
pub async fn kg_neighbors(
&self,
id: &str,
node: &str,
direction: ExpandDirection,
max_hops: usize,
) -> ServiceResult<KgNeighborsPayload> {
let handle = self.open_handle(id)?;
let (nodes, triples) = handle
.kg
.expand_neighbors(node, direction, max_hops)
.map_err(|e| ServiceError::internal(format!("kg expand_neighbors: {e:#}")))?;
Ok(KgNeighborsPayload {
origin: node.to_string(),
returned_node_count: nodes.len() as u64,
returned_triple_count: triples.len() as u64,
nodes: nodes.into_iter().map(KgNodeView::from).collect(),
triples,
direction: match direction {
ExpandDirection::In => "in",
ExpandDirection::Out => "out",
ExpandDirection::Both => "both",
}
.to_string(),
max_hops: max_hops as u64,
})
}
pub async fn dream_status_aggregate(&self) -> DreamStatusPayload {
let palaces = PalaceRegistry::list_palaces(&self.state.data_root).unwrap_or_default();
let mut out = DreamStatusPayload::default();
let mut latest: Option<chrono::DateTime<chrono::Utc>> = None;
for p in palaces {
let data_dir = self.state.data_root.join(p.id.as_str());
let snap = match PersistedDreamStats::load(&data_dir) {
Ok(Some(s)) => s,
_ => continue,
};
out.merged = out.merged.saturating_add(snap.stats.merged);
out.pruned = out.pruned.saturating_add(snap.stats.pruned);
out.compacted = out.compacted.saturating_add(snap.stats.compacted);
out.closets_updated = out
.closets_updated
.saturating_add(snap.stats.closets_updated);
out.duration_ms = out.duration_ms.saturating_add(snap.stats.duration_ms);
latest = match latest {
Some(t) if t >= snap.last_run_at => Some(t),
_ => Some(snap.last_run_at),
};
}
out.last_run_at = latest;
out
}
pub async fn dream_status_for_palace(&self, id: &str) -> ServiceResult<DreamStatusPayload> {
let data_dir = self.state.data_root.join(id);
if !data_dir.exists() {
return Err(ServiceError::not_found(format!("palace not found: {id}")));
}
match PersistedDreamStats::load(&data_dir) {
Ok(Some(s)) => Ok(s.into()),
Ok(None) => Ok(DreamStatusPayload::default()),
Err(e) => Err(ServiceError::internal(format!("read dream stats: {e:#}"))),
}
}
pub async fn dream_run(&self) -> ServiceResult<DreamStatusPayload> {
if !self.state.registry.may_run_maintenance() {
return Err(ServiceError::conflict(MAINTENANCE_LEASE_NOT_HELD));
}
let palaces = list_palaces_blocking(&self.state)
.await
.map_err(|e| ServiceError::internal(format!("{e:#}")))?;
let dreamer = Dreamer::new(DreamConfig::default());
let mut out = DreamStatusPayload::default();
for p in palaces {
let registry = std::sync::Arc::clone(&self.state.registry);
let root = self.state.data_root.clone();
let pid = p.id.clone();
let opened =
tokio::task::spawn_blocking(move || registry.open_palace(&root, &pid)).await;
let handle = match opened {
Ok(Ok(h)) => h,
Ok(Err(e)) => {
tracing::warn!(palace = %p.id, "dream_run: open failed: {e:#}");
continue;
}
Err(e) => {
tracing::warn!(palace = %p.id, "dream_run: join open failed: {e}");
continue;
}
};
match dreamer.dream_cycle(&handle).await {
Ok(stats) => {
out.merged = out.merged.saturating_add(stats.merged);
out.pruned = out.pruned.saturating_add(stats.pruned);
out.compacted = out.compacted.saturating_add(stats.compacted);
out.closets_updated = out.closets_updated.saturating_add(stats.closets_updated);
out.duration_ms = out.duration_ms.saturating_add(stats.duration_ms);
}
Err(e) => tracing::warn!(palace = %p.id, "dream_run: cycle failed: {e:#}"),
}
refresh_gaps_cache(&self.state, &handle).await;
}
out.last_run_at = Some(chrono::Utc::now());
self.state.emit(DaemonEvent::DreamCompleted {
palace_id: None,
merged: out.merged,
pruned: out.pruned,
compacted: out.compacted,
closets_updated: out.closets_updated,
duration_ms: out.duration_ms,
source: ActivitySource::Http,
});
self.state.emit(self.aggregate_status_event());
Ok(out)
}
pub async fn list_activity(
&self,
filter: ActivityFilter,
limit: usize,
offset: usize,
) -> ServiceResult<(Vec<crate::ActivityEntry>, u64)> {
let entries = self
.state
.activity_log
.list(&filter, limit, offset)
.map_err(|e| ServiceError::internal(format!("activity list: {e:#}")))?;
let total = self
.state
.activity_log
.count()
.map_err(|e| ServiceError::internal(format!("activity count: {e:#}")))?;
Ok((entries, total))
}
pub fn open_handle(&self, id: &str) -> ServiceResult<Arc<PalaceHandle>> {
self.state
.registry
.open_palace(&self.state.data_root, &PalaceId::new(id))
.map_err(|e| {
if PalaceRegistry::open_error_is_absent(&e) {
ServiceError::not_found(format!("palace not found: {id} ({e:#})"))
} else {
ServiceError::internal(format!("palace could not be loaded: {id} ({e:#})"))
}
})
}
}