klieo_core/

memory.rs

//! Memory traits — short-term, long-term, episodic.

use crate::error::MemoryError;
use crate::ids::{FactId, RunId, ThreadId};
use crate::llm::Message;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

/// Outcome of a tool invocation as recorded in the episodic event stream.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum ToolResult {
    /// Tool returned a successful JSON result.
    Ok {
        /// Result payload.
        value: serde_json::Value,
    },
    /// Tool returned an error message.
    Err {
        /// Error string.
        message: String,
    },
}

/// Conversation buffer scoped to a single thread.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryShortTerm;
/// use klieo_core::{ShortTermMemory, Message, Role, ThreadId};
///
/// let m = InMemoryShortTerm::default();
/// let thread = ThreadId::new("t1");
/// m.append(thread.clone(), Message {
///     role: Role::User, content: "hi".into(),
///     tool_calls: vec![], tool_call_id: None,
/// }).await.unwrap();
/// let loaded = m.load(thread, 1024).await.unwrap();
/// assert_eq!(loaded.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait ShortTermMemory: Send + Sync {
    /// Append a message to the thread's history.
    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError>;

    /// Load up to `max_tokens` of the most-recent messages, oldest first.
    /// Implementations approximate token counts (provider-specific).
    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError>;

    /// Drop all messages for `thread`.
    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError>;
}

/// Namespacing for long-term memory facts.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Scope {
    /// Workspace-scoped (multi-agent shared).
    Workspace(String),
    /// Per-agent scoped.
    Agent(String),
    /// Process-global (use sparingly).
    Global,
}

/// One stored fact in long-term memory.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Fact {
    /// Plain-text body, embedded for retrieval.
    pub text: String,
    /// Caller-supplied metadata, opaque to the store.
    #[serde(default)]
    pub metadata: serde_json::Value,
}

/// Long-term semantic memory.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryLongTerm;
/// use klieo_core::{Fact, LongTermMemory, Scope};
///
/// let m = InMemoryLongTerm::default();
/// let scope = Scope::Workspace("ws".into());
/// m.remember(scope.clone(), Fact {
///     text: "the sky is blue".into(),
///     metadata: serde_json::Value::Null,
/// }).await.unwrap();
/// let hits = m.recall(scope, "sky", 1).await.unwrap();
/// assert_eq!(hits.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait LongTermMemory: Send + Sync {
    /// Store a fact under `scope`. Returns a stable id.
    async fn remember(&self, scope: Scope, fact: Fact) -> Result<FactId, MemoryError>;

    /// Top-`k` semantic recall under `scope` for the supplied query.
    async fn recall(&self, scope: Scope, query: &str, k: usize) -> Result<Vec<Fact>, MemoryError>;

    /// Remove a stored fact.
    async fn forget(&self, id: FactId) -> Result<(), MemoryError>;
}

/// One event in the episodic event stream of a single agent run.
///
/// Marked `#[non_exhaustive]` so additive variants (e.g.
/// [`Self::SummaryCheckpoint`]) can be introduced without forcing a
/// SemVer-major bump. Match arms in downstream crates must include a
/// fallback `_ => …`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Episode {
    /// Run started.
    Started {
        /// Agent name.
        agent: String,
    },
    /// LLM call completed.
    LlmCall {
        /// Total tokens used.
        tokens: u32,
        /// Wall-clock latency.
        latency_ms: u32,
    },
    /// Tool call completed.
    ToolCall {
        /// Tool name.
        name: String,
        /// JSON arguments.
        args: serde_json::Value,
        /// Tool outcome.
        result: ToolResult,
    },
    /// Agent published a bus message.
    BusPublish {
        /// Subject.
        subject: String,
    },
    /// Agent received a bus message.
    BusReceive {
        /// Subject.
        subject: String,
    },
    /// Run completed successfully.
    Completed,
    /// Run failed.
    Failed {
        /// Error message.
        error: String,
    },
    /// Summarizer checkpoint completed.
    ///
    /// Emitted by [`crate::summarize::summarize_history`] in lieu of
    /// [`Self::LlmCall`] so the audit trail can distinguish summarizer
    /// overhead from substantive agent reasoning. Downstream
    /// observability (e.g. `klieo-runlog`) typically projects this as
    /// a separate step kind so cost / latency attribution stays
    /// faithful.
    SummaryCheckpoint {
        /// Number of older messages folded into the summary call.
        input_message_count: u32,
        /// Length of the resulting summary, in Unicode scalar values.
        summary_chars: u32,
        /// Wall-clock latency of the summarizer call.
        latency_ms: u32,
        /// Total tokens reported by the summarizer LLM (prompt +
        /// completion).
        tokens: u32,
    },
    /// Operational-layer event (klieo-ops). Body is an opaque
    /// `serde_json::Value` to keep klieo-core free of an ops dependency.
    /// klieo-ops provides typed serde conversion helpers via `OpsEvent`.
    Ops(serde_json::Value),
}

/// Filter passed to `EpisodicMemory::list_runs`.
#[derive(Debug, Clone, Default)]
pub struct RunFilter {
    /// Filter by agent name (substring match).
    pub agent: Option<String>,
    /// Inclusive lower bound on `started_at`.
    pub since: Option<DateTime<Utc>>,
    /// Inclusive upper bound on `started_at`.
    pub until: Option<DateTime<Utc>>,
    /// Maximum rows returned.
    pub limit: Option<usize>,
}

/// Index-row summary returned by `list_runs`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunSummary {
    /// Run id.
    pub run_id: RunId,
    /// Agent name.
    pub agent: String,
    /// First-event timestamp.
    pub started_at: DateTime<Utc>,
    /// Last-event timestamp, if completed/failed.
    pub finished_at: Option<DateTime<Utc>>,
    /// Number of episodes.
    pub episode_count: u32,
}

/// Append-only event log of agent runs.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::InMemoryEpisodic;
/// use klieo_core::{Episode, EpisodicMemory, RunId};
///
/// let m = InMemoryEpisodic::default();
/// let run = RunId::new();
/// m.record(run, Episode::Started { agent: "a".into() }).await.unwrap();
/// let events = m.replay(run).await.unwrap();
/// assert_eq!(events.len(), 1);
/// # });
/// ```
#[async_trait]
pub trait EpisodicMemory: Send + Sync {
    /// Record an episode for `run`.
    async fn record(&self, run: RunId, event: Episode) -> Result<(), MemoryError>;

    /// Replay all episodes for `run` in order.
    async fn replay(&self, run: RunId) -> Result<Vec<Episode>, MemoryError>;

    /// List run summaries matching `filter`.
    async fn list_runs(&self, filter: RunFilter) -> Result<Vec<RunSummary>, MemoryError>;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[allow(dead_code)]
    fn _assert_dyn_short(_: &dyn ShortTermMemory) {}
    #[allow(dead_code)]
    fn _assert_dyn_long(_: &dyn LongTermMemory) {}
    #[allow(dead_code)]
    fn _assert_dyn_episodic(_: &dyn EpisodicMemory) {}
}