pub mod anchored_summary;
pub mod node;
pub mod pattern_promotion;
pub mod query;
pub mod snapshot;
pub mod store;
pub mod trajectory_table;
mod trajectory_persist;
pub use anchored_summary::{anchored_summary_id, ANCHORED_SUMMARY_TAG};
pub use trajectory_persist::{
persist_trajectory_coarse_tools, persist_trajectory_for_episode, trajectory_env_enabled,
};
pub use node::{
AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, FailureNode, MemoryCategory,
PersonaLayer, PersonaNode, PersonaSource, ProceduralNode, ProcedureType, RuntimeStateNode,
SemanticNode, Sentiment, StrengthEvent, TrajectoryNode,
};
pub use query::{
count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
recall_task_scoped_episodes, walk_from, GraphQuery,
};
pub use snapshot::{
AgentGraphSnapshot, DanglingEdgeDetail, GraphValidationReport, SnapshotEdge,
SNAPSHOT_SCHEMA_VERSION,
};
pub use store::{GraphStore, GraphValidationError, SnapshotImportError, SqliteGraphStore};
pub use trajectory_table::TrajectoryDetailRecord;
use uuid::Uuid;
pub struct GraphMemory {
store: SqliteGraphStore,
}
impl GraphMemory {
pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
let store = SqliteGraphStore::open(db_path)?;
Ok(Self { store })
}
pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
let store = SqliteGraphStore::from_connection(conn)?;
Ok(Self { store })
}
pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
Self { store }
}
pub fn write_episode(
&self,
tool_calls: Vec<String>,
delegation_to: Option<String>,
trace_event: Option<serde_json::Value>,
) -> Result<Uuid, String> {
let turn_id = Uuid::new_v4();
let timestamp = chrono::Utc::now().timestamp();
let node =
AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
let node_id = node.id;
self.store.write_node(&node)?;
Ok(node_id)
}
pub fn write_fact(
&self,
fact: String,
confidence: f32,
source_turn_id: Uuid,
) -> Result<Uuid, String> {
let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
let node_id = node.id;
self.store.write_node(&node)?;
Ok(node_id)
}
pub fn store_pattern(
&self,
pattern_name: String,
compiled_graph: Vec<u8>,
) -> Result<Uuid, String> {
let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
let node_id = node.id;
self.store.write_node(&node)?;
Ok(node_id)
}
pub fn write_procedural(
&self,
pattern_name: &str,
tool_sequence: Vec<String>,
confidence: f32,
) -> Result<Uuid, String> {
let mut node = AinlMemoryNode::new_procedural_tools(
pattern_name.to_string(),
tool_sequence,
confidence,
);
if let AinlNodeType::Procedural { ref mut procedural } = node.node_type {
procedural.pattern_observation_count =
procedural
.pattern_observation_count
.max(crate::pattern_promotion::DEFAULT_MIN_OBSERVATIONS);
let floor = crate::pattern_promotion::DEFAULT_FITNESS_FLOOR;
if let Some(f) = procedural.fitness {
procedural.fitness = Some(f.max(floor));
} else {
procedural.fitness = Some(floor);
}
procedural.prompt_eligible = true;
}
let node_id = node.id;
self.store.write_node(&node)?;
Ok(node_id)
}
pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
self.store.insert_graph_edge(source, target, rel)
}
pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
let since = chrono::Utc::now().timestamp() - seconds_ago;
self.store.query_episodes_since(since, 100)
}
pub fn recall_by_type(
&self,
kind: AinlNodeKind,
seconds_ago: i64,
) -> Result<Vec<AinlMemoryNode>, String> {
let since = chrono::Utc::now().timestamp() - seconds_ago;
self.store
.query_nodes_by_type_since(kind.as_str(), since, 500)
}
pub fn find_procedural_by_tool_sequence(
&self,
agent_id: &str,
tool_sequence: &[String],
) -> Result<Option<AinlMemoryNode>, String> {
let norm: Vec<String> = tool_sequence.iter().map(|s| s.trim().to_string()).collect();
if norm.is_empty() {
return Ok(None);
}
let nodes = self.recall_by_type(AinlNodeKind::Procedural, 60 * 60 * 24 * 365 * 5)?;
for n in nodes {
if n.agent_id != agent_id {
continue;
}
let AinlNodeType::Procedural { ref procedural } = n.node_type else {
continue;
};
if procedural.tool_sequence.len() != norm.len() {
continue;
}
let same = procedural
.tool_sequence
.iter()
.zip(norm.iter())
.all(|(a, b)| a.trim() == b.trim());
if same {
return Ok(Some(n));
}
}
Ok(None)
}
pub fn write_persona(
&self,
trait_name: &str,
strength: f32,
learned_from: Vec<Uuid>,
) -> Result<Uuid, String> {
let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
let node_id = node.id;
self.store.write_node(&node)?;
Ok(node_id)
}
pub fn store(&self) -> &dyn GraphStore {
&self.store
}
pub fn sqlite_store(&self) -> &SqliteGraphStore {
&self.store
}
pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
self.store.validate_graph(agent_id)
}
pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
self.store.export_graph(agent_id)
}
pub fn import_graph(
&mut self,
snapshot: &AgentGraphSnapshot,
allow_dangling_edges: bool,
) -> Result<(), String> {
self.store.import_graph(snapshot, allow_dangling_edges)
}
pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
self.store.agent_subgraph_edges(agent_id)
}
pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
self.store.write_node_with_edges(node)
}
pub fn insert_graph_edge_checked(
&self,
from_id: Uuid,
to_id: Uuid,
label: &str,
) -> Result<(), String> {
self.store.insert_graph_edge_checked(from_id, to_id, label)
}
pub fn read_runtime_state(&self, agent_id: &str) -> Result<Option<RuntimeStateNode>, String> {
self.store.read_runtime_state(agent_id)
}
pub fn write_runtime_state(&self, state: &RuntimeStateNode) -> Result<(), String> {
self.store.write_runtime_state(state)
}
pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
self.store.write_node(node)
}
pub fn insert_trajectory_detail(&self, row: &TrajectoryDetailRecord) -> Result<(), String> {
self.store.insert_trajectory_detail(row)
}
pub fn list_trajectories_for_agent(
&self,
agent_id: &str,
limit: usize,
since_timestamp: Option<i64>,
) -> Result<Vec<TrajectoryDetailRecord>, String> {
self.store
.list_trajectories_for_agent(agent_id, limit, since_timestamp)
}
pub fn count_trajectory_details_before(
&self,
agent_id: &str,
before_recorded_at: i64,
) -> Result<usize, String> {
self.store
.count_trajectory_details_before(agent_id, before_recorded_at)
}
pub fn prune_trajectory_details_before(
&self,
agent_id: &str,
before_recorded_at: i64,
) -> Result<usize, String> {
self.store
.delete_trajectory_details_before(agent_id, before_recorded_at)
}
pub fn search_failures_for_agent(
&self,
agent_id: &str,
query: &str,
limit: usize,
) -> Result<Vec<AinlMemoryNode>, String> {
self.store
.search_failures_fts_for_agent(agent_id, query, limit)
}
pub fn search_all_nodes_fts(
&self,
agent_id: &str,
query: &str,
project_id: Option<&str>,
limit: usize,
) -> Result<Vec<AinlMemoryNode>, String> {
self.store
.search_all_nodes_fts_for_agent(agent_id, query, project_id, limit)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_graph_memory_api() {
let temp_dir = std::env::temp_dir();
let db_path = temp_dir.join("ainl_lib_test.db");
let _ = std::fs::remove_file(&db_path);
let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
let episode_id = memory
.write_episode(
vec!["file_read".to_string(), "agent_delegate".to_string()],
Some("agent-B".to_string()),
None,
)
.expect("Failed to write episode");
assert_ne!(episode_id, Uuid::nil());
let fact_id = memory
.write_fact(
"User prefers concise responses".to_string(),
0.85,
episode_id,
)
.expect("Failed to write fact");
assert_ne!(fact_id, Uuid::nil());
let recent = memory.recall_recent(60).expect("Failed to recall");
assert_eq!(recent.len(), 1);
if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
assert_eq!(episodic.tool_calls.len(), 2);
} else {
panic!("Wrong node type");
}
}
#[test]
fn test_store_pattern() {
let temp_dir = std::env::temp_dir();
let db_path = temp_dir.join("ainl_lib_test_pattern.db");
let _ = std::fs::remove_file(&db_path);
let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
let pattern_id = memory
.store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
.expect("Failed to store pattern");
assert_ne!(pattern_id, Uuid::nil());
let patterns = find_patterns(memory.store(), "research").expect("Query failed");
assert_eq!(patterns.len(), 1);
}
#[test]
fn failure_write_and_fts_search_roundtrip() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_failure_fts_smoke.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent_id = "agent-smoke-fts";
let mut node = AinlMemoryNode::new_loop_guard_failure(
"block",
Some("shell_exec"),
"repeated identical tool invocation blocked by loop guard",
Some("session-xyz"),
);
node.agent_id = agent_id.to_string();
let nid = node.id;
memory.write_node(&node).expect("write failure node");
let hits = memory
.search_failures_for_agent(agent_id, "loop", 10)
.expect("search loop");
assert_eq!(hits.len(), 1, "expected one FTS hit for token 'loop'");
assert_eq!(hits[0].id, nid);
assert!(
matches!(&hits[0].node_type, AinlNodeType::Failure { .. }),
"expected Failure node type"
);
let hits2 = memory
.search_failures_for_agent(agent_id, "shell_exec", 10)
.expect("search tool name");
assert_eq!(hits2.len(), 1);
assert_eq!(hits2[0].id, nid);
let empty = memory
.search_failures_for_agent(agent_id, " ", 10)
.expect("whitespace-only query");
assert!(empty.is_empty());
let wrong_agent = memory
.search_failures_for_agent("other-agent", "loop", 10)
.expect("wrong agent id");
assert!(wrong_agent.is_empty());
}
#[test]
fn all_nodes_fts_write_and_search_roundtrip() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_all_nodes_fts.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent_id = "agent-fts-all";
let mut node = AinlMemoryNode::new_fact("unique-fts-violet-cat-42".into(), 0.8, Uuid::new_v4());
node.agent_id = agent_id.to_string();
let nid = node.id;
memory.write_node(&node).expect("write fact");
let hits = memory
.search_all_nodes_fts(agent_id, "violet", None, 10)
.expect("search");
assert_eq!(hits.len(), 1, "expected one all-nodes FTS hit");
assert_eq!(hits[0].id, nid);
}
#[test]
fn tool_execution_failure_write_and_fts_search_roundtrip() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_tool_failure_fts.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent_id = "agent-tool-ft";
let mut node = AinlMemoryNode::new_tool_execution_failure(
"file_read",
"ENOENT: no such file or directory",
Some("sess-tool-1"),
);
node.agent_id = agent_id.to_string();
let nid = node.id;
memory.write_node(&node).expect("write tool failure node");
let hits = memory
.search_failures_for_agent(agent_id, "ENOENT", 10)
.expect("search ENOENT");
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].id, nid);
let src_hits = memory
.search_failures_for_agent(agent_id, "tool_runner", 10)
.expect("search source");
assert_eq!(src_hits.len(), 1);
assert_eq!(src_hits[0].id, nid);
}
#[test]
fn agent_loop_precheck_failure_write_and_fts_search_roundtrip() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_precheck_failure_fts.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent_id = "agent-precheck-ft";
let mut node = AinlMemoryNode::new_agent_loop_precheck_failure(
"param_validation",
"file_write",
"missing required field: path",
Some("sess-pv-1"),
);
node.agent_id = agent_id.to_string();
let nid = node.id;
memory.write_node(&node).expect("write precheck failure");
let hits = memory
.search_failures_for_agent(agent_id, "param_validation", 10)
.expect("search kind");
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].id, nid);
let hits2 = memory
.search_failures_for_agent(agent_id, "agent_loop", 10)
.expect("search agent_loop prefix");
assert_eq!(hits2.len(), 1);
}
#[test]
fn ainl_runtime_graph_validation_failure_write_and_fts_search_roundtrip() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_graph_validation_failure_fts.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent_id = "agent-graph-val-ft";
let mut node = AinlMemoryNode::new_ainl_runtime_graph_validation_failure(
"graph validation failed before turn: dangling edges …",
Some("sess-gv-1"),
);
node.agent_id = agent_id.to_string();
let nid = node.id;
memory.write_node(&node).expect("write graph validation failure");
let hits = memory
.search_failures_for_agent(agent_id, "graph_validation", 10)
.expect("search source label");
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].id, nid);
let hits2 = memory
.search_failures_for_agent(agent_id, "dangling", 10)
.expect("search message body");
assert_eq!(hits2.len(), 1);
}
#[test]
fn trajectory_detail_prune_before_drops_only_old_rows() {
use ainl_contracts::{TrajectoryOutcome, TrajectoryStep};
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("ainl_traj_prune.db");
let memory = GraphMemory::new(&db_path).expect("graph memory");
let agent = "agent-traj-prune";
let ep_old = memory
.write_episode(vec![], None, None)
.expect("episode for old traj");
let ep_new = memory
.write_episode(vec![], None, None)
.expect("episode for new traj");
let mk_step = |sid: &str| TrajectoryStep {
step_id: sid.to_string(),
timestamp_ms: 0,
adapter: "a".into(),
operation: "o".into(),
inputs_preview: None,
outputs_preview: None,
duration_ms: 1,
success: true,
error: None,
vitals: None,
freshness_at_step: None,
frame_vars: None,
tool_telemetry: None,
};
let r_old = TrajectoryDetailRecord {
id: Uuid::new_v4(),
episode_id: ep_old,
graph_trajectory_node_id: None,
agent_id: agent.to_string(),
session_id: "s-old".into(),
project_id: None,
recorded_at: 100,
outcome: TrajectoryOutcome::Success,
ainl_source_hash: None,
duration_ms: 1,
steps: vec![mk_step("1")],
frame_vars: None,
fitness_delta: None,
};
let r_new = TrajectoryDetailRecord {
id: Uuid::new_v4(),
episode_id: ep_new,
graph_trajectory_node_id: None,
agent_id: agent.to_string(),
session_id: "s-new".into(),
project_id: None,
recorded_at: 200,
outcome: TrajectoryOutcome::Success,
ainl_source_hash: None,
duration_ms: 1,
steps: vec![mk_step("2")],
frame_vars: None,
fitness_delta: None,
};
memory.insert_trajectory_detail(&r_old).expect("insert old");
memory.insert_trajectory_detail(&r_new).expect("insert new");
let before = memory
.list_trajectories_for_agent(agent, 10, None)
.expect("list");
assert_eq!(before.len(), 2);
let removed = memory
.prune_trajectory_details_before(agent, 200)
.expect("prune");
assert_eq!(removed, 1);
let after = memory
.list_trajectories_for_agent(agent, 10, None)
.expect("list after");
assert_eq!(after.len(), 1);
assert_eq!(after[0].recorded_at, 200);
}
}