pe_core/

cognitive.rs

//! Cognitive architecture — optional inner subgraph for agent reasoning.
//!
//! When present on an [`Agent`](crate::agent::Agent), the cognitive architecture
//! decomposes the agent's reasoning into parallel cognitive streams that process
//! the same input through different lenses (emotional, logical, devil's advocate, etc.)
//! and synthesize the results before the main LLM call.
//!
//! ## How it works
//!
//! - `None` → normal execution. Single LLM call, zero overhead.
//! - `Some(CognitiveArchitecture)` → the agent carries an inner subgraph.
//!   When a node calls the LLM, the cognitive architecture runs first:
//!   parallel streams process the agent's self-context, a hippocampus node
//!   synthesizes them, and the main LLM receives an enriched prompt.
//!
//! ## Multi-model
//!
//! Each stream can use a different (typically smaller/cheaper) model.
//! The main model is always the agent's `model_preference.primary`.
//! Streams do pre-processing — the main model gets richer input.
//!
//! ## Runtime wiring
//!
//! The types here define the cognitive architecture configuration and state.
//! The actual inner graph execution is wired by pe-runtime (future plan).
//! The inner graph uses the same `CompiledGraph<CognitiveState>` primitives
//! from pe-graph — no new engine needed.

use crate::cognitive_budget::CognitiveBudget;
use crate::cognitive_memory::{MemoryConfig, WorkingNote};
use crate::cognitive_signal::CognitiveSignal;
use crate::lobe::LobeOutput;
use crate::self_model::{FailureRecord, NegativeKnowledge, SelfModel};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Defines how an agent's reasoning is decomposed into parallel streams.
///
/// This is a property of the agent, not the graph. The graph is oblivious —
/// a node calls `llm.complete()`, and the cognitive architecture intercepts
/// to run the inner subgraph transparently.
///
/// # Example
///
/// ```ignore
/// let cognitive = CognitiveArchitecture::new()
///     .add_stream(CognitiveStream::new("logical", "Analyze purely logically."))
///     .add_stream(CognitiveStream::new("antithink", "Challenge every assumption."))
///     .with_synthesis(SynthesisStrategy::Integrate);
///
/// let agent = Agent::new("analyst", "You analyze data.")
///     .with_cognitive_architecture(cognitive);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CognitiveArchitecture {
    /// Parallel cognitive streams — each processes input through a different lens.
    pub streams: Vec<CognitiveStream>,

    /// How stream outputs are combined into the final enriched prompt.
    pub synthesis: SynthesisStrategy,

    /// Resource limits for cognitive processing.
    #[serde(default)]
    pub budget: CognitiveBudget,

    /// Memory tier configuration.
    #[serde(default)]
    pub memory_config: MemoryConfig,

    /// The agent's self-model (self/user/collective context).
    #[serde(default)]
    pub self_model: SelfModel,
}

impl CognitiveArchitecture {
    /// Create a new cognitive architecture with no streams.
    pub fn new() -> Self {
        Self {
            streams: Vec::new(),
            synthesis: SynthesisStrategy::Integrate,
            budget: CognitiveBudget::default(),
            memory_config: MemoryConfig::default(),
            self_model: SelfModel::default(),
        }
    }

    /// Add a cognitive stream.
    #[must_use]
    pub fn add_stream(mut self, stream: CognitiveStream) -> Self {
        self.streams.push(stream);
        self
    }

    /// Set the synthesis strategy.
    #[must_use]
    pub fn with_synthesis(mut self, strategy: SynthesisStrategy) -> Self {
        self.synthesis = strategy;
        self
    }

    /// Set the cognitive budget.
    #[must_use]
    pub fn with_budget(mut self, budget: CognitiveBudget) -> Self {
        self.budget = budget;
        self
    }

    /// Set the memory configuration.
    #[must_use]
    pub fn with_memory_config(mut self, config: MemoryConfig) -> Self {
        self.memory_config = config;
        self
    }

    /// Set the self-model.
    #[must_use]
    pub fn with_self_model(mut self, model: SelfModel) -> Self {
        self.self_model = model;
        self
    }
}

impl Default for CognitiveArchitecture {
    fn default() -> Self {
        Self::new()
    }
}

/// A single cognitive processing stream — one lens on the agent's identity.
///
/// All streams share the same agent identity, memory, and boundaries.
/// Each stream modifies the system prompt to focus on a specific cognitive
/// function, and can optionally use a different (typically smaller) model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CognitiveStream {
    /// What cognitive function this stream handles.
    /// Examples: "emotional", "logical", "antithink", "past-errors", "creative"
    pub lens: String,

    /// How the agent's system prompt is modified for this stream.
    /// Appended to the agent's base system prompt.
    pub prompt_modifier: String,

    /// Optional model override — use a smaller/cheaper model for this stream.
    /// When `None`, uses the agent's primary model.
    /// Fast streams (haiku, local models) cost ~1% of the main model.
    #[serde(default)]
    pub model_override: Option<String>,

    /// Stream-specific metadata.
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl CognitiveStream {
    /// Create a new cognitive stream with a lens name and prompt modifier.
    pub fn new(lens: impl Into<String>, prompt_modifier: impl Into<String>) -> Self {
        Self {
            lens: lens.into(),
            prompt_modifier: prompt_modifier.into(),
            model_override: None,
            metadata: HashMap::new(),
        }
    }

    /// Use a different model for this stream (e.g., a fast/cheap model).
    #[must_use]
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model_override = Some(model.into());
        self
    }
}

/// How cognitive stream outputs are combined into the final result.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[non_exhaustive]
pub enum SynthesisStrategy {
    /// LLM synthesizes all stream outputs into a unified perspective.
    /// The hippocampus node receives all outputs and produces one result.
    #[default]
    Integrate,

    /// Majority vote — streams produce discrete choices, most common wins.
    Vote,

    /// Weighted combination — streams have confidence scores, higher weight wins.
    Weighted,
}

/// State for the inner cognitive subgraph.
///
/// Defines the fields the cognitive architecture operates on.
/// Separate from the outer task state. Implements [`State`](crate::State)
/// so it can be used with `CompiledGraph<CognitiveState>` directly.
///
/// ## Merge semantics
///
/// **Last-value-wins (replace):**
/// `input`, `synthesis_result`, `budget_tokens`, `current_plan`, `confidence`
///
/// **Merge (HashMap extend):**
/// `stream_outputs` — new keys overwrite, old keys preserved.
/// Key is `CognitiveStream::lens` or `Lobe::name()`.
///
/// **Append (Vec extend):**
/// `loaded_memories`, `working_notes`, `quality_trend`, `error_history`,
/// `signals`, `negative_knowledge`, `failure_records`
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CognitiveState {
    /// The input prompt/task the agent is processing.
    pub input: String,

    /// Full outputs from each cognitive lobe, keyed by lobe name.
    ///
    /// Stores the complete [`LobeOutput`] (content + confidence + signals + metadata)
    /// so the synthesizer receives lossless data. The key is `Lobe::name()`.
    #[serde(default)]
    pub stream_outputs: HashMap<String, LobeOutput>,

    /// The synthesized result from synthesis node.
    #[serde(default)]
    pub synthesis_result: Option<String>,

    /// Token budget for cognitive processing (fraction of agent's budget).
    #[serde(default)]
    pub budget_tokens: Option<u32>,

    /// Relevant memories loaded for this cognitive cycle.
    #[serde(default)]
    pub loaded_memories: Vec<String>,

    // --- Working Memory ---
    /// Agent's scratchpad — structured notes with categories.
    #[serde(default)]
    pub working_notes: Vec<WorkingNote>,

    /// What the agent thinks it should do next.
    #[serde(default)]
    pub current_plan: Option<String>,

    // --- Self-Awareness ---
    /// Confidence level (0.0-1.0), fed from matrix C value when available.
    #[serde(default)]
    pub confidence: f64,

    /// Rolling quality scores — is the agent getting better or worse?
    #[serde(default)]
    pub quality_trend: Vec<f64>,

    /// Record of errors/failures during this session.
    #[serde(default)]
    pub error_history: Vec<String>,

    // --- Signals ---
    /// Signals emitted by lobes for the outer graph to read.
    #[serde(default)]
    pub signals: Vec<CognitiveSignal>,

    // --- Structured Stores ---
    /// Negative knowledge — things the agent learned NOT to do.
    #[serde(default)]
    pub negative_knowledge: Vec<NegativeKnowledge>,

    /// Structured failure records for pattern recognition.
    #[serde(default)]
    pub failure_records: Vec<FailureRecord>,
}

/// Partial update for [`CognitiveState`].
///
/// Each field is `Option<T>` — `None` means "no change", `Some(v)` means "apply v".
/// - `input`, `synthesis_result`, `budget_tokens`: last-value-wins (replace).
/// - `stream_outputs`: merge — new entries overwrite existing keys, old keys preserved.
/// - `loaded_memories`: append — new memories are added to the existing list.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CognitiveStateUpdate {
    /// Replace the input prompt.
    #[serde(default)]
    pub input: Option<String>,

    /// Merge lobe outputs — new keys overwrite, existing keys preserved.
    #[serde(default)]
    pub stream_outputs: Option<HashMap<String, LobeOutput>>,

    /// Replace the synthesis result.
    #[serde(default)]
    pub synthesis_result: Option<Option<String>>,

    /// Replace the token budget.
    #[serde(default)]
    pub budget_tokens: Option<Option<u32>>,

    /// Append to loaded memories.
    #[serde(default)]
    pub loaded_memories: Option<Vec<String>>,

    /// Append working notes.
    #[serde(default)]
    pub working_notes: Option<Vec<WorkingNote>>,

    /// Replace current plan.
    #[serde(default)]
    pub current_plan: Option<Option<String>>,

    /// Replace confidence.
    #[serde(default)]
    pub confidence: Option<f64>,

    /// Append quality scores.
    #[serde(default)]
    pub quality_trend: Option<Vec<f64>>,

    /// Append error history entries.
    #[serde(default)]
    pub error_history: Option<Vec<String>>,

    /// Append cognitive signals.
    #[serde(default)]
    pub signals: Option<Vec<CognitiveSignal>>,

    /// Append negative knowledge entries.
    #[serde(default)]
    pub negative_knowledge: Option<Vec<NegativeKnowledge>>,

    /// Append failure records.
    #[serde(default)]
    pub failure_records: Option<Vec<FailureRecord>>,

    /// Replace working notes entirely (used by meditate/consolidation).
    ///
    /// When `Some`, the entire `working_notes` vector is replaced (not appended).
    /// Takes precedence over `working_notes` if both are set.
    #[serde(default)]
    pub replace_working_notes: Option<Vec<WorkingNote>>,

    /// Replace failure records entirely (used by meditate/consolidation).
    ///
    /// When `Some`, the entire `failure_records` vector is replaced.
    /// Takes precedence over `failure_records` if both are set.
    #[serde(default)]
    pub replace_failure_records: Option<Vec<FailureRecord>>,

    /// Replace negative knowledge entirely (used by meditate/consolidation).
    ///
    /// When `Some`, the entire `negative_knowledge` vector is replaced.
    /// Takes precedence over `negative_knowledge` if both are set.
    #[serde(default)]
    pub replace_negative_knowledge: Option<Vec<NegativeKnowledge>>,
}

impl crate::state::StateUpdate for CognitiveStateUpdate {}

impl crate::state::State for CognitiveState {
    type Update = CognitiveStateUpdate;

    fn apply(&mut self, update: CognitiveStateUpdate) {
        // Last-value-wins fields
        if let Some(input) = update.input {
            self.input = input;
        }
        if let Some(synthesis) = update.synthesis_result {
            self.synthesis_result = synthesis;
        }
        if let Some(budget) = update.budget_tokens {
            self.budget_tokens = budget;
        }
        if let Some(plan) = update.current_plan {
            self.current_plan = plan;
        }
        if let Some(confidence) = update.confidence {
            self.confidence = confidence;
        }

        // Merge fields (HashMap extend)
        if let Some(outputs) = update.stream_outputs {
            self.stream_outputs.extend(outputs);
        }

        // Replace fields (meditate/consolidation) — takes precedence over append
        if let Some(notes) = update.replace_working_notes {
            self.working_notes = notes;
        } else if let Some(notes) = update.working_notes {
            self.working_notes.extend(notes);
        }
        if let Some(failures) = update.replace_failure_records {
            self.failure_records = failures;
        } else if let Some(failures) = update.failure_records {
            self.failure_records.extend(failures);
        }
        if let Some(nk) = update.replace_negative_knowledge {
            self.negative_knowledge = nk;
        } else if let Some(nk) = update.negative_knowledge {
            self.negative_knowledge.extend(nk);
        }

        // Append fields (Vec extend)
        if let Some(memories) = update.loaded_memories {
            self.loaded_memories.extend(memories);
        }
        if let Some(scores) = update.quality_trend {
            self.quality_trend.extend(scores);
        }
        if let Some(errors) = update.error_history {
            self.error_history.extend(errors);
        }
        if let Some(signals) = update.signals {
            self.signals.extend(signals);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cognitive_memory::NoteCategory;
    use crate::self_model::{FailureRecord, NegativeKnowledge};
    use crate::state::State;

    fn base_state() -> CognitiveState {
        CognitiveState {
            input: "analyze this".to_string(),
            stream_outputs: HashMap::from([(
                "logical".to_string(),
                LobeOutput::new("logic output", 0.9).with_lobe_name("logical"),
            )]),
            synthesis_result: None,
            budget_tokens: Some(1000),
            loaded_memories: vec!["mem-1".to_string()],
            working_notes: Vec::new(),
            current_plan: None,
            confidence: 0.5,
            quality_trend: Vec::new(),
            error_history: Vec::new(),
            signals: Vec::new(),
            negative_knowledge: Vec::new(),
            failure_records: Vec::new(),
        }
    }

    #[test]
    fn test_apply_input_replaces() {
        let mut state = base_state();
        state.apply(CognitiveStateUpdate {
            input: Some("new input".to_string()),
            ..Default::default()
        });
        assert_eq!(state.input, "new input");
        // Other fields unchanged
        assert_eq!(state.stream_outputs.len(), 1);
        assert_eq!(state.budget_tokens, Some(1000));
    }

    #[test]
    fn test_apply_stream_outputs_merges() {
        let mut state = base_state();
        state.apply(CognitiveStateUpdate {
            stream_outputs: Some(HashMap::from([
                (
                    "emotional".to_string(),
                    LobeOutput::new("emotion output", 0.7).with_lobe_name("emotional"),
                ),
                (
                    "logical".to_string(),
                    LobeOutput::new("updated logic", 0.95).with_lobe_name("logical"),
                ),
            ])),
            ..Default::default()
        });
        // New key added
        assert_eq!(
            state.stream_outputs.get("emotional").unwrap().content,
            "emotion output"
        );
        // Existing key overwritten
        assert_eq!(
            state.stream_outputs.get("logical").unwrap().content,
            "updated logic"
        );
        // Confidence preserved
        assert!(
            (state.stream_outputs.get("logical").unwrap().confidence - 0.95).abs() < f64::EPSILON
        );
        assert_eq!(state.stream_outputs.len(), 2);
    }

    #[test]
    fn test_apply_synthesis_result_replaces() {
        let mut state = base_state();
        assert!(state.synthesis_result.is_none());

        state.apply(CognitiveStateUpdate {
            synthesis_result: Some(Some("synthesized answer".to_string())),
            ..Default::default()
        });
        assert_eq!(
            state.synthesis_result.as_deref(),
            Some("synthesized answer")
        );

        // Can also clear it
        state.apply(CognitiveStateUpdate {
            synthesis_result: Some(None),
            ..Default::default()
        });
        assert!(state.synthesis_result.is_none());
    }

    #[test]
    fn test_apply_loaded_memories_appends() {
        let mut state = base_state();
        assert_eq!(state.loaded_memories, vec!["mem-1"]);

        state.apply(CognitiveStateUpdate {
            loaded_memories: Some(vec!["mem-2".to_string(), "mem-3".to_string()]),
            ..Default::default()
        });
        assert_eq!(state.loaded_memories, vec!["mem-1", "mem-2", "mem-3"]);
    }

    #[test]
    fn test_apply_none_fields_no_change() {
        let mut state = base_state();
        let original = state.clone();
        state.apply(CognitiveStateUpdate::default());
        assert_eq!(state.input, original.input);
        assert_eq!(state.stream_outputs, original.stream_outputs);
        assert_eq!(state.synthesis_result, original.synthesis_result);
        assert_eq!(state.budget_tokens, original.budget_tokens);
        assert_eq!(state.loaded_memories, original.loaded_memories);
    }

    #[test]
    fn test_apply_multiple_fields_at_once() {
        let mut state = base_state();
        state.apply(CognitiveStateUpdate {
            input: Some("new prompt".to_string()),
            stream_outputs: Some(HashMap::from([(
                "creative".to_string(),
                LobeOutput::new("creative out", 0.8).with_lobe_name("creative"),
            )])),
            synthesis_result: Some(Some("final".to_string())),
            budget_tokens: Some(Some(500)),
            loaded_memories: Some(vec!["mem-4".to_string()]),
            ..Default::default()
        });
        assert_eq!(state.input, "new prompt");
        assert_eq!(state.stream_outputs.len(), 2); // logical + creative
        assert_eq!(state.synthesis_result.as_deref(), Some("final"));
        assert_eq!(state.budget_tokens, Some(500));
        assert_eq!(state.loaded_memories, vec!["mem-1", "mem-4"]);
    }

    #[test]
    fn test_apply_working_notes_appends() {
        let mut state = base_state();
        let note = WorkingNote::new("important finding", NoteCategory::Discovery);
        state.apply(CognitiveStateUpdate {
            working_notes: Some(vec![note]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 1);
        assert_eq!(state.working_notes[0].content, "important finding");

        // Second append
        state.apply(CognitiveStateUpdate {
            working_notes: Some(vec![WorkingNote::new("concern", NoteCategory::Concern)]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 2);
    }

    #[test]
    fn test_apply_confidence_replaces() {
        let mut state = base_state();
        assert!((state.confidence - 0.5).abs() < f64::EPSILON);
        state.apply(CognitiveStateUpdate {
            confidence: Some(0.9),
            ..Default::default()
        });
        assert!((state.confidence - 0.9).abs() < f64::EPSILON);
    }

    #[test]
    fn test_apply_signals_appends() {
        let mut state = base_state();
        state.apply(CognitiveStateUpdate {
            signals: Some(vec![CognitiveSignal::Proceed]),
            ..Default::default()
        });
        state.apply(CognitiveStateUpdate {
            signals: Some(vec![CognitiveSignal::SimplifyMode]),
            ..Default::default()
        });
        assert_eq!(state.signals.len(), 2);
    }

    #[test]
    fn test_apply_negative_knowledge_appends() {
        use crate::self_model::Severity;
        let mut state = base_state();
        let nk = NegativeKnowledge::new("api", "max 100 items", Severity::High);
        state.apply(CognitiveStateUpdate {
            negative_knowledge: Some(vec![nk]),
            ..Default::default()
        });
        assert_eq!(state.negative_knowledge.len(), 1);
        assert_eq!(state.negative_knowledge[0].category, "api");
    }

    #[test]
    fn test_apply_failure_records_appends() {
        let mut state = base_state();
        let record = FailureRecord::new("db_migration", "ALTER TABLE");
        state.apply(CognitiveStateUpdate {
            failure_records: Some(vec![record]),
            ..Default::default()
        });
        assert_eq!(state.failure_records.len(), 1);
    }

    #[test]
    fn test_cognitive_state_update_serialization() {
        let update = CognitiveStateUpdate {
            input: Some("test".to_string()),
            ..Default::default()
        };
        let json = serde_json::to_string(&update).unwrap();
        let deserialized: CognitiveStateUpdate = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.input.as_deref(), Some("test"));
        assert!(deserialized.stream_outputs.is_none());
    }

    #[test]
    fn test_replace_working_notes_overrides_existing() {
        let mut state = base_state();
        // First add some notes via append.
        state.apply(CognitiveStateUpdate {
            working_notes: Some(vec![
                WorkingNote::new("old note 1", NoteCategory::Observation),
                WorkingNote::new("old note 2", NoteCategory::Concern),
                WorkingNote::new("old note 3", NoteCategory::Discovery),
            ]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 3);

        // Now replace with a smaller set (meditate consolidation).
        state.apply(CognitiveStateUpdate {
            replace_working_notes: Some(vec![WorkingNote::new(
                "consolidated",
                NoteCategory::Reflection,
            )]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 1);
        assert_eq!(state.working_notes[0].content, "consolidated");
    }

    #[test]
    fn test_replace_takes_precedence_over_append() {
        let mut state = base_state();
        state.apply(CognitiveStateUpdate {
            working_notes: Some(vec![WorkingNote::new("old", NoteCategory::Observation)]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 1);

        // Both replace and append set — replace wins.
        state.apply(CognitiveStateUpdate {
            working_notes: Some(vec![WorkingNote::new("appended", NoteCategory::Concern)]),
            replace_working_notes: Some(vec![WorkingNote::new("replaced", NoteCategory::Plan)]),
            ..Default::default()
        });
        assert_eq!(state.working_notes.len(), 1);
        assert_eq!(state.working_notes[0].content, "replaced");
    }
}