klieo-graph-projector 3.3.0

//! [`EpisodeProjector`] — broadcast subscriber that translates each
//! recorded [`Episode`] into typed entities and indexes them into a
//! [`KnowledgeGraph`].
//!
//! Two entry points:
//!
//! - [`EpisodeProjector::spawn`] — owns a
//!   `tokio::sync::broadcast::Receiver<(RunId, Episode)>` and projects
//!   each delivered episode asynchronously. Returns a `JoinHandle<()>`;
//!   the caller drives shutdown by cancelling the supplied
//!   [`CancellationToken`] or by dropping every [`broadcast::Sender`]
//!   (`RecvError::Closed`).
//! - [`EpisodeProjector::project_run`] — on-demand replay of a single
//!   run via [`klieo_core::memory::EpisodicMemory::replay`]. Used for
//!   cold-start backfill or one-off re-indexing.
//!
//! Persistent failures inside an individual per-episode projection are
//! logged at `ERROR` and the spawn loop continues — projection is
//! best-effort and must never starve the agent's event stream. Lag
//! (when the receiver falls behind producer rate) is recorded on
//! [`RecallMetrics::record_projector_lag`] so ops dashboards stay
//! honest about dropped projection coverage.

use crate::error::ProjectionError;
use crate::rules;
use klieo_core::ids::{FactId, RunId};
use klieo_core::memory::{Episode, EpisodicMemory, Scope};
use klieo_memory_graph::{EntityExtractor, IndexEntry, KnowledgeGraph, RecallMetrics};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::sync::broadcast::{self, error::RecvError};
use tokio_util::sync::CancellationToken;

/// Subscriber that projects [`Episode`]s into a [`KnowledgeGraph`].
///
/// Multiple instances may share the same graph + extractor as long as
/// each carries its own [`Scope`] (or a distinct
/// [`RecallMetrics`] when scope is shared). The internal monotonic
/// `seq` counter is **per-projector** — two projectors over the same
/// run produce overlapping `FactId`s, which is the correct behaviour
/// when re-indexing the same run after the projector restarts.
///
/// # FactId namespace — separate scope from RAG vector stores
///
/// The projector mints synthetic `FactId`s of the form
/// `{run_id}-{kind}-{seq}` (see [`Self`]'s internal `next_fact_id`) and
/// writes them ONLY to the graph — never to a `LongTermMemory`. These ids
/// therefore do not exist in any vector store. A graph used by the projector
/// MUST NOT be shared as the `KnowledgeGraph` of a vector-backed
/// `GraphAwareLongTerm`: that composer feeds `neighbors()` ids straight into
/// `LongTermMemory::recall_filtered`, where the projector's synthetic ids
/// (graph-only) match nothing and the projected facts silently vanish from
/// recall. Use a dedicated graph scope for projection, distinct from any
/// RAG graph.
#[non_exhaustive]
pub struct EpisodeProjector {
    graph: Arc<dyn KnowledgeGraph>,
    extractor: Arc<dyn EntityExtractor>,
    scope: Scope,
    metrics: Arc<RecallMetrics>,
    /// Monotonic counter so two events of the same kind in the same
    /// run get distinct [`FactId`]s. Plain `Relaxed` ordering is
    /// sufficient: the projector spawn loop is single-task, and
    /// `project_run` is the only other writer (also serial). Cross-
    /// projector overlap is allowed by design.
    seq: AtomicU64,
}

impl EpisodeProjector {
    /// Construct a new projector over `graph` + `extractor` with the
    /// given write [`Scope`] and shared [`RecallMetrics`] sink.
    pub fn new(
        graph: Arc<dyn KnowledgeGraph>,
        extractor: Arc<dyn EntityExtractor>,
        scope: Scope,
        metrics: Arc<RecallMetrics>,
    ) -> Self {
        Self {
            graph,
            extractor,
            scope,
            metrics,
            seq: AtomicU64::new(0),
        }
    }

    /// Capability-shaped builder facade — defaults the shared
    /// [`RecallMetrics`] sink so callers only supply the three
    /// domain-required values (`graph`, `extractor`, `scope`).
    pub fn builder() -> EpisodeProjectorBuilder {
        EpisodeProjectorBuilder::default()
    }

    /// Spawn a background task that consumes `rx` until the
    /// `cancel` token fires, every [`broadcast::Sender`] is dropped,
    /// or the receiver is permanently closed.
    ///
    /// Per-episode projection failures are logged at `ERROR` and the
    /// loop continues — projection coverage is best-effort. Receiver
    /// lag (capacity overflow) is recorded on
    /// [`RecallMetrics::record_projector_lag`].
    pub fn spawn(
        self: Arc<Self>,
        mut rx: broadcast::Receiver<(RunId, Episode)>,
        cancel: CancellationToken,
    ) -> tokio::task::JoinHandle<()> {
        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = cancel.cancelled() => break,
                    msg = rx.recv() => match msg {
                        Ok((run_id, episode)) => {
                            let kind = episode_kind(&episode);
                            if let Err(e) = self.project_episode(run_id, &episode).await {
                                tracing::error!(
                                    error = %e,
                                    run_id = %run_id,
                                    episode_kind = kind,
                                    "episode projection failed",
                                );
                            }
                        }
                        Err(RecvError::Lagged(skipped)) => {
                            tracing::warn!(
                                skipped,
                                scope = ?self.scope,
                                "projector lagged; episodes dropped",
                            );
                            self.metrics.record_projector_lag(skipped);
                        }
                        Err(RecvError::Closed) => break,
                    }
                }
            }
        })
    }

    /// Replay every episode of `run_id` from `episodic` and project
    /// them as one batched graph write. Returns the number of
    /// episodes processed (including episodes that produce no
    /// entities — those are counted but contribute nothing to the
    /// batch).
    ///
    /// Extraction is sequential per episode (per-episode entity
    /// discovery is intentionally serial), but the graph write
    /// flushes the whole run in a single
    /// [`KnowledgeGraph::index_many`] call so backends with native
    /// bulk-insert (e.g. Neo4j UNWIND) collapse N facts into O(1)
    /// round-trips.
    pub async fn project_run(
        &self,
        episodic: &dyn EpisodicMemory,
        run_id: RunId,
    ) -> Result<usize, ProjectionError> {
        let episodes = episodic
            .replay(run_id)
            .await
            .map_err(ProjectionError::replay)?;
        let mut batch: Vec<IndexEntry> = Vec::with_capacity(episodes.len());
        for episode in &episodes {
            if let Some(entry) = self.build_index_entry(run_id, episode).await? {
                batch.push(entry);
            }
        }
        if !batch.is_empty() {
            self.graph
                .index_many(self.scope.clone(), &batch)
                .await
                .map_err(ProjectionError::graph)?;
        }
        Ok(episodes.len())
    }

    async fn project_episode(
        &self,
        run_id: RunId,
        episode: &Episode,
    ) -> Result<(), ProjectionError> {
        // Live broadcast path: can't batch across unknown future
        // episodes; route through the single-fact `index` so each
        // event lands as soon as it arrives.
        let Some(entry) = self.build_index_entry(run_id, episode).await? else {
            return Ok(());
        };
        self.graph
            .index(
                self.scope.clone(),
                &entry.fact_id,
                &entry.entities,
                &entry.text,
                entry.valid_from,
            )
            .await
            .map_err(ProjectionError::graph)
    }

    async fn build_index_entry(
        &self,
        run_id: RunId,
        episode: &Episode,
    ) -> Result<Option<IndexEntry>, ProjectionError> {
        let hints = rules::entities_for_episode(episode);
        let entities = self
            .extractor
            .extract("", &hints)
            .await
            .map_err(ProjectionError::extraction)?;
        if entities.is_empty() {
            return Ok(None);
        }
        let fact_id = self.next_fact_id(run_id, episode);
        Ok(Some(IndexEntry::new(fact_id, entities, "", None)))
    }

    fn next_fact_id(&self, run_id: RunId, episode: &Episode) -> FactId {
        let seq = self.seq.fetch_add(1, Ordering::Relaxed);
        FactId::new(format!("{run_id}-{}-{seq}", episode_kind(episode)))
    }
}

/// Builder facade for [`EpisodeProjector`] — defaults the
/// [`RecallMetrics`] sink so callers only supply the three
/// domain-required values at [`Self::build`] time.
#[derive(Default)]
pub struct EpisodeProjectorBuilder {
    metrics: Option<Arc<RecallMetrics>>,
}

impl EpisodeProjectorBuilder {
    /// Supply a shared [`RecallMetrics`] sink — set this when the
    /// caller wants to read `metrics.snapshot()` from an ops endpoint.
    /// Without this, the builder allocates a fresh counter the caller
    /// cannot observe.
    pub fn metrics(mut self, metrics: Arc<RecallMetrics>) -> Self {
        self.metrics = Some(metrics);
        self
    }

    /// Assemble the [`EpisodeProjector`]. `metrics` defaults to a fresh
    /// [`RecallMetrics`] the caller cannot observe unless supplied via
    /// [`Self::metrics`].
    pub fn build(
        self,
        graph: Arc<dyn KnowledgeGraph>,
        extractor: Arc<dyn EntityExtractor>,
        scope: Scope,
    ) -> EpisodeProjector {
        EpisodeProjector::new(
            graph,
            extractor,
            scope,
            self.metrics
                .unwrap_or_else(|| Arc::new(RecallMetrics::default())),
        )
    }
}

/// Short, stable kind label embedded in projected [`FactId`]s. New
/// [`Episode`] variants default to `"unknown"`; matching by `kind`
/// downstream stays compile-stable across non-exhaustive additions.
fn episode_kind(episode: &Episode) -> &'static str {
    match episode {
        Episode::Started { .. } => "started",
        Episode::LlmCall { .. } => "llmcall",
        Episode::ToolCall { .. } => "toolcall",
        Episode::BusPublish { .. } => "buspublish",
        Episode::BusReceive { .. } => "busreceive",
        Episode::Completed => "completed",
        Episode::Failed { .. } => "failed",
        Episode::SummaryCheckpoint { .. } => "summarycheckpoint",
        Episode::Ops(_) => "ops",
        _ => "unknown",
    }
}