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