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