Skip to main content

ainl_memory/
node.rs

1//! AINL graph node types - the vocabulary of agent memory.
2//!
3//! Four core memory types: Episode (episodic), Semantic, Procedural, Persona.
4//! Designed to be standalone (zero ArmaraOS deps) yet compatible with
5//! OrchestrationTraceEvent serialization.
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use uuid::Uuid;
10
11/// Coarse node kind for store queries (matches `node_type` column values).
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum AinlNodeKind {
14    Episode,
15    Semantic,
16    Procedural,
17    Persona,
18}
19
20impl AinlNodeKind {
21    pub fn as_str(&self) -> &'static str {
22        match self {
23            Self::Episode => "episode",
24            Self::Semantic => "semantic",
25            Self::Procedural => "procedural",
26            Self::Persona => "persona",
27        }
28    }
29}
30
31/// Memory category aligned with the four memory families (episodic ↔ `Episode` nodes).
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[serde(rename_all = "snake_case")]
34pub enum MemoryCategory {
35    Persona,
36    Semantic,
37    Episodic,
38    Procedural,
39    /// Agent-scoped runtime session counters / cache hints (persisted by `ainl-runtime`).
40    RuntimeState,
41}
42
43impl MemoryCategory {
44    pub fn from_node_type(node_type: &AinlNodeType) -> Self {
45        match node_type {
46            AinlNodeType::Episode { .. } => MemoryCategory::Episodic,
47            AinlNodeType::Semantic { .. } => MemoryCategory::Semantic,
48            AinlNodeType::Procedural { .. } => MemoryCategory::Procedural,
49            AinlNodeType::Persona { .. } => MemoryCategory::Persona,
50            AinlNodeType::RuntimeState { .. } => MemoryCategory::RuntimeState,
51        }
52    }
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
56#[serde(rename_all = "snake_case")]
57pub enum PersonaLayer {
58    #[default]
59    Base,
60    Delta,
61    Injection,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
65#[serde(rename_all = "snake_case")]
66pub enum PersonaSource {
67    SystemDefault,
68    #[default]
69    UserConfigured,
70    Evolved,
71    Feedback,
72    Injection,
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
76#[serde(rename_all = "snake_case")]
77pub enum Sentiment {
78    Positive,
79    Neutral,
80    Negative,
81    Mixed,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
85#[serde(rename_all = "snake_case")]
86pub enum ProcedureType {
87    #[default]
88    ToolSequence,
89    ResponsePattern,
90    WorkflowStep,
91    BehavioralRule,
92}
93
94/// One strength adjustment on a persona trait (evolution / provenance).
95#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
96pub struct StrengthEvent {
97    pub delta: f32,
98    pub reason: String,
99    pub episode_id: String,
100    pub timestamp: u64,
101}
102
103fn default_importance_score() -> f32 {
104    0.5
105}
106
107fn default_semantic_confidence() -> f32 {
108    0.7
109}
110
111fn default_decay_eligible() -> bool {
112    true
113}
114
115fn default_success_rate() -> f32 {
116    0.5
117}
118
119fn default_strength_floor() -> f32 {
120    0.0
121}
122
123/// Canonical persona payload (flattened under `AinlNodeType::Persona` in JSON).
124#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
125pub struct PersonaNode {
126    pub trait_name: String,
127    pub strength: f32,
128    #[serde(default)]
129    pub learned_from: Vec<Uuid>,
130    #[serde(default)]
131    pub layer: PersonaLayer,
132    #[serde(default)]
133    pub source: PersonaSource,
134    #[serde(default = "default_strength_floor")]
135    pub strength_floor: f32,
136    #[serde(default)]
137    pub locked: bool,
138    #[serde(default)]
139    pub relevance_score: f32,
140    #[serde(default)]
141    pub provenance_episode_ids: Vec<String>,
142    #[serde(default)]
143    pub evolution_log: Vec<StrengthEvent>,
144    /// Optional axis-evolution bundle (`ainl-persona`); omitted in JSON → empty map.
145    #[serde(default)]
146    pub axis_scores: HashMap<String, f32>,
147    #[serde(default)]
148    pub evolution_cycle: u32,
149    /// ISO-8601 timestamp of last persona evolution pass.
150    #[serde(default)]
151    pub last_evolved: String,
152    /// Redundant copy of owning agent id (mirrors `AinlMemoryNode.agent_id` for payload consumers).
153    #[serde(default)]
154    pub agent_id: String,
155    /// Soft labels: axes above the high-spectrum threshold, not discrete classes.
156    #[serde(default)]
157    pub dominant_axes: Vec<String>,
158}
159
160/// Semantic / factual memory payload.
161#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
162pub struct SemanticNode {
163    pub fact: String,
164    #[serde(default = "default_semantic_confidence")]
165    pub confidence: f32,
166    pub source_turn_id: Uuid,
167    #[serde(default)]
168    pub topic_cluster: Option<String>,
169    #[serde(default)]
170    pub source_episode_id: String,
171    #[serde(default)]
172    pub contradiction_ids: Vec<String>,
173    #[serde(default)]
174    pub last_referenced_at: u64,
175    /// How many times this node has been retrieved from the store.
176    /// Managed by the recall path only — never written by extractors.
177    #[serde(default)]
178    pub reference_count: u32,
179    #[serde(default = "default_decay_eligible")]
180    pub decay_eligible: bool,
181    /// Optional tag hints for analytics / persona (`ainl-persona`); omitted → empty.
182    #[serde(default)]
183    pub tags: Vec<String>,
184    /// How many times this exact fact has recurred across separate extraction events.
185    /// Written by `graph_extractor` when the same fact is observed again.
186    ///
187    /// Do **not** use `reference_count` as a substitute: that field tracks retrieval frequency,
188    /// not extraction recurrence. They measure different things. `graph_extractor` (Prompt 2)
189    /// must write `recurrence_count` directly; persona / domain extractors gate on this field only.
190    #[serde(default)]
191    pub recurrence_count: u32,
192    /// `reference_count` snapshot from the last graph-extractor pass (JSON key `_last_ref_snapshot`).
193    #[serde(rename = "_last_ref_snapshot", default)]
194    pub last_ref_snapshot: u32,
195}
196
197/// Episodic memory payload (one turn / moment).
198#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
199pub struct EpisodicNode {
200    pub turn_id: Uuid,
201    pub timestamp: i64,
202    #[serde(default)]
203    pub tool_calls: Vec<String>,
204    #[serde(default)]
205    pub delegation_to: Option<String>,
206    #[serde(skip_serializing_if = "Option::is_none")]
207    pub trace_event: Option<serde_json::Value>,
208    #[serde(default)]
209    pub turn_index: u32,
210    #[serde(default)]
211    pub user_message_tokens: u32,
212    #[serde(default)]
213    pub assistant_response_tokens: u32,
214    /// Preferred list of tools for analytics; mirrors `tool_calls` when not set explicitly.
215    #[serde(default)]
216    pub tools_invoked: Vec<String>,
217    /// Persona signal names emitted this turn (`Vec`, never `Option`). Omitted JSON → `[]`.
218    /// Serialized even when empty (no `skip_serializing_if`). Backfill: `read_node` → patch → `write_node`.
219    #[serde(default)]
220    pub persona_signals_emitted: Vec<String>,
221    #[serde(default)]
222    pub sentiment: Option<Sentiment>,
223    #[serde(default)]
224    pub flagged: bool,
225    #[serde(default)]
226    pub conversation_id: String,
227    #[serde(default)]
228    pub follows_episode_id: Option<String>,
229    /// Optional raw user message for offline extractors (`ainl-graph-extractor`); omitted unless set.
230    #[serde(default, skip_serializing_if = "Option::is_none")]
231    pub user_message: Option<String>,
232    /// Optional assistant reply text for offline extractors; omitted unless set.
233    #[serde(default, skip_serializing_if = "Option::is_none")]
234    pub assistant_response: Option<String>,
235}
236
237impl EpisodicNode {
238    /// Effective tool list: `tools_invoked` if non-empty, else `tool_calls`.
239    pub fn effective_tools(&self) -> &[String] {
240        if !self.tools_invoked.is_empty() {
241            &self.tools_invoked
242        } else {
243            &self.tool_calls
244        }
245    }
246}
247
248/// Procedural memory payload.
249#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
250pub struct ProceduralNode {
251    pub pattern_name: String,
252    #[serde(default)]
253    pub compiled_graph: Vec<u8>,
254    #[serde(default, skip_serializing_if = "Vec::is_empty")]
255    pub tool_sequence: Vec<String>,
256    #[serde(default, skip_serializing_if = "Option::is_none")]
257    pub confidence: Option<f32>,
258    #[serde(default)]
259    pub procedure_type: ProcedureType,
260    #[serde(default)]
261    pub trigger_conditions: Vec<String>,
262    #[serde(default)]
263    pub success_count: u32,
264    #[serde(default)]
265    pub failure_count: u32,
266    #[serde(default = "default_success_rate")]
267    pub success_rate: f32,
268    #[serde(default)]
269    pub last_invoked_at: u64,
270    #[serde(default)]
271    pub reinforcement_episode_ids: Vec<String>,
272    #[serde(default)]
273    pub suppression_episode_ids: Vec<String>,
274    /// Graph-patch / refinement generation (`ainl-persona`); omitted JSON → 0 (skip persona extract until bumped).
275    #[serde(default)]
276    pub patch_version: u32,
277    /// Optional fitness score in \[0,1\]; when absent, consumers may fall back to `success_rate`.
278    #[serde(default)]
279    pub fitness: Option<f32>,
280    /// Declared read dependencies for the procedure (metadata-only hints).
281    #[serde(default)]
282    pub declared_reads: Vec<String>,
283    /// When true, excluded from [`crate::GraphQuery::active_patches`] and skipped by patch dispatch.
284    #[serde(default)]
285    pub retired: bool,
286    /// IR label for graph-patch identity (empty → runtimes may fall back to [`Self::pattern_name`]).
287    #[serde(default)]
288    pub label: String,
289    /// Optional orchestration / turn correlation id (same namespace as episodic `trace_event.trace_id`).
290    #[serde(default, skip_serializing_if = "Option::is_none")]
291    pub trace_id: Option<String>,
292}
293
294impl ProceduralNode {
295    pub fn recompute_success_rate(&mut self) {
296        let total = self.success_count.saturating_add(self.failure_count);
297        self.success_rate = if total == 0 {
298            0.5
299        } else {
300            self.success_count as f32 / total as f32
301        };
302    }
303}
304
305/// Persisted session counters and persona prompt cache for one agent (`ainl-runtime` ↔ SQLite).
306#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
307pub struct RuntimeStateNode {
308    pub agent_id: String,
309    pub turn_count: u32,
310    pub last_extraction_turn: u32,
311    pub last_persona_prompt: Option<String>,
312    pub updated_at: String,
313}
314
315/// Core AINL node types - the vocabulary of agent memory.
316#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
317#[serde(tag = "type", rename_all = "snake_case")]
318pub enum AinlNodeType {
319    /// Episodic memory: what happened during an agent turn
320    Episode {
321        #[serde(flatten)]
322        episodic: EpisodicNode,
323    },
324
325    /// Semantic memory: facts learned, with confidence
326    Semantic {
327        #[serde(flatten)]
328        semantic: SemanticNode,
329    },
330
331    /// Procedural memory: reusable compiled workflow patterns
332    Procedural {
333        #[serde(flatten)]
334        procedural: ProceduralNode,
335    },
336
337    /// Persona memory: traits learned over time
338    Persona {
339        #[serde(flatten)]
340        persona: PersonaNode,
341    },
342
343    /// Runtime session state (turn counters, extraction cadence, persona cache snapshot).
344    RuntimeState {
345        runtime_state: RuntimeStateNode,
346    },
347}
348
349/// A node in the AINL memory graph
350#[derive(Serialize, Debug, Clone, PartialEq)]
351pub struct AinlMemoryNode {
352    pub id: Uuid,
353    pub memory_category: MemoryCategory,
354    pub importance_score: f32,
355    pub agent_id: String,
356    pub node_type: AinlNodeType,
357    pub edges: Vec<AinlEdge>,
358}
359
360#[derive(Deserialize)]
361struct AinlMemoryNodeWire {
362    id: Uuid,
363    #[serde(default)]
364    memory_category: Option<MemoryCategory>,
365    #[serde(default)]
366    importance_score: Option<f32>,
367    #[serde(default)]
368    agent_id: Option<String>,
369    node_type: AinlNodeType,
370    #[serde(default)]
371    edges: Vec<AinlEdge>,
372}
373
374impl From<AinlMemoryNodeWire> for AinlMemoryNode {
375    fn from(w: AinlMemoryNodeWire) -> Self {
376        let memory_category = w
377            .memory_category
378            .unwrap_or_else(|| MemoryCategory::from_node_type(&w.node_type));
379        let importance_score = w.importance_score.unwrap_or_else(default_importance_score);
380        Self {
381            id: w.id,
382            memory_category,
383            importance_score,
384            agent_id: w.agent_id.unwrap_or_default(),
385            node_type: w.node_type,
386            edges: w.edges,
387        }
388    }
389}
390
391impl<'de> Deserialize<'de> for AinlMemoryNode {
392    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
393    where
394        D: serde::Deserializer<'de>,
395    {
396        let w = AinlMemoryNodeWire::deserialize(deserializer)?;
397        Ok(Self::from(w))
398    }
399}
400
401/// Typed edge connecting memory nodes
402#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
403pub struct AinlEdge {
404    /// Target node ID
405    pub target_id: Uuid,
406
407    /// Edge label (e.g., "delegated_to", "learned_from", "caused_by")
408    pub label: String,
409}
410
411impl AinlMemoryNode {
412    fn base(
413        memory_category: MemoryCategory,
414        importance_score: f32,
415        agent_id: String,
416        node_type: AinlNodeType,
417    ) -> Self {
418        Self {
419            id: Uuid::new_v4(),
420            memory_category,
421            importance_score,
422            agent_id,
423            node_type,
424            edges: Vec::new(),
425        }
426    }
427
428    /// Create a new episode node
429    pub fn new_episode(
430        turn_id: Uuid,
431        timestamp: i64,
432        tool_calls: Vec<String>,
433        delegation_to: Option<String>,
434        trace_event: Option<serde_json::Value>,
435    ) -> Self {
436        let tools_invoked = tool_calls.clone();
437        let episodic = EpisodicNode {
438            turn_id,
439            timestamp,
440            tool_calls,
441            delegation_to,
442            trace_event,
443            turn_index: 0,
444            user_message_tokens: 0,
445            assistant_response_tokens: 0,
446            tools_invoked,
447            persona_signals_emitted: Vec::new(),
448            sentiment: None,
449            flagged: false,
450            conversation_id: String::new(),
451            follows_episode_id: None,
452            user_message: None,
453            assistant_response: None,
454        };
455        Self::base(
456            MemoryCategory::Episodic,
457            default_importance_score(),
458            String::new(),
459            AinlNodeType::Episode { episodic },
460        )
461    }
462
463    /// Create a new semantic fact node
464    pub fn new_fact(fact: String, confidence: f32, source_turn_id: Uuid) -> Self {
465        let semantic = SemanticNode {
466            fact,
467            confidence,
468            source_turn_id,
469            topic_cluster: None,
470            source_episode_id: String::new(),
471            contradiction_ids: Vec::new(),
472            last_referenced_at: 0,
473            reference_count: 0,
474            decay_eligible: true,
475            tags: Vec::new(),
476            recurrence_count: 0,
477            last_ref_snapshot: 0,
478        };
479        Self::base(
480            MemoryCategory::Semantic,
481            default_importance_score(),
482            String::new(),
483            AinlNodeType::Semantic { semantic },
484        )
485    }
486
487    /// Create a new procedural pattern node
488    pub fn new_pattern(pattern_name: String, compiled_graph: Vec<u8>) -> Self {
489        let mut procedural = ProceduralNode {
490            pattern_name,
491            compiled_graph,
492            tool_sequence: Vec::new(),
493            confidence: None,
494            procedure_type: ProcedureType::default(),
495            trigger_conditions: Vec::new(),
496            success_count: 0,
497            failure_count: 0,
498            success_rate: default_success_rate(),
499            last_invoked_at: 0,
500            reinforcement_episode_ids: Vec::new(),
501            suppression_episode_ids: Vec::new(),
502            patch_version: 1,
503            fitness: None,
504            declared_reads: Vec::new(),
505            retired: false,
506            label: String::new(),
507            trace_id: None,
508        };
509        procedural.recompute_success_rate();
510        Self::base(
511            MemoryCategory::Procedural,
512            default_importance_score(),
513            String::new(),
514            AinlNodeType::Procedural { procedural },
515        )
516    }
517
518    /// Procedural node from a detected tool workflow (no compiled IR).
519    pub fn new_procedural_tools(
520        pattern_name: String,
521        tool_sequence: Vec<String>,
522        confidence: f32,
523    ) -> Self {
524        let mut procedural = ProceduralNode {
525            pattern_name,
526            compiled_graph: Vec::new(),
527            tool_sequence,
528            confidence: Some(confidence),
529            procedure_type: ProcedureType::ToolSequence,
530            trigger_conditions: Vec::new(),
531            success_count: 0,
532            failure_count: 0,
533            success_rate: default_success_rate(),
534            last_invoked_at: 0,
535            reinforcement_episode_ids: Vec::new(),
536            suppression_episode_ids: Vec::new(),
537            patch_version: 1,
538            fitness: None,
539            declared_reads: Vec::new(),
540            retired: false,
541            label: String::new(),
542            trace_id: None,
543        };
544        procedural.recompute_success_rate();
545        Self::base(
546            MemoryCategory::Procedural,
547            default_importance_score(),
548            String::new(),
549            AinlNodeType::Procedural { procedural },
550        )
551    }
552
553    /// Create a new persona trait node
554    pub fn new_persona(trait_name: String, strength: f32, learned_from: Vec<Uuid>) -> Self {
555        let persona = PersonaNode {
556            trait_name,
557            strength,
558            learned_from,
559            layer: PersonaLayer::default(),
560            source: PersonaSource::default(),
561            strength_floor: default_strength_floor(),
562            locked: false,
563            relevance_score: 0.0,
564            provenance_episode_ids: Vec::new(),
565            evolution_log: Vec::new(),
566            axis_scores: HashMap::new(),
567            evolution_cycle: 0,
568            last_evolved: String::new(),
569            agent_id: String::new(),
570            dominant_axes: Vec::new(),
571        };
572        Self::base(
573            MemoryCategory::Persona,
574            default_importance_score(),
575            String::new(),
576            AinlNodeType::Persona { persona },
577        )
578    }
579
580    pub fn episodic(&self) -> Option<&EpisodicNode> {
581        match &self.node_type {
582            AinlNodeType::Episode { episodic } => Some(episodic),
583            _ => None,
584        }
585    }
586
587    pub fn semantic(&self) -> Option<&SemanticNode> {
588        match &self.node_type {
589            AinlNodeType::Semantic { semantic } => Some(semantic),
590            _ => None,
591        }
592    }
593
594    pub fn procedural(&self) -> Option<&ProceduralNode> {
595        match &self.node_type {
596            AinlNodeType::Procedural { procedural } => Some(procedural),
597            _ => None,
598        }
599    }
600
601    pub fn persona(&self) -> Option<&PersonaNode> {
602        match &self.node_type {
603            AinlNodeType::Persona { persona } => Some(persona),
604            _ => None,
605        }
606    }
607
608    /// Add an edge to another node
609    pub fn add_edge(&mut self, target_id: Uuid, label: impl Into<String>) {
610        self.edges.push(AinlEdge {
611            target_id,
612            label: label.into(),
613        });
614    }
615}