klieo-graph-projector 3.3.0

//! Static rules that map each [`Episode`] variant to a hint set of
//! [`EntityRef`]s suitable for indexing into a
//! [`klieo_memory_graph::KnowledgeGraph`].
//!
//! Rules are intentionally minimal and deterministic: they extract
//! structurally-typed information (agent name, tool name, LLM model
//! and provider, agent-prefixed bus subjects). Free-form entity
//! extraction over episode text — names mentioned in tool args,
//! arbitrary subjects — is the job of the
//! [`klieo_memory_graph::EntityExtractor`] supplied to
//! [`crate::EpisodeProjector`].
//!
//! The variant match is non-exhaustive on purpose. [`Episode`] is
//! `#[non_exhaustive]` (see klieo-core memory.rs); a future variant
//! added in a 1.x minor must not break this crate at compile time,
//! and a fall-through arm that returns the empty entity set is the
//! correct default ("project nothing until a rule is added").

use klieo_core::memory::Episode;
use klieo_memory_graph::{EntityRef, EntityType};

const AGENT_SUBJECT_PREFIXES: &[&str] = &["agent.", "klieo.agent."];

/// Map an episode to the structural entities it directly references.
///
/// Returns an empty `Vec` when no structural entities apply: e.g.
/// [`Episode::Completed`], [`Episode::Failed`],
/// [`Episode::SummaryCheckpoint`], [`Episode::Ops`], a bus subject
/// that does not match the agent-naming convention, or an
/// [`Episode::LlmCall`] whose `provider` and `model` are both
/// `None`.
///
/// The returned entities are passed to
/// [`klieo_memory_graph::EntityExtractor::extract`] as hints so the
/// extractor can dedupe against free-form matches in the episode
/// body (when one exists).
pub fn entities_for_episode(episode: &Episode) -> Vec<EntityRef> {
    match episode {
        Episode::Started { agent } => vec![EntityRef::new(EntityType::Agent, agent.as_str())],
        Episode::ToolCall { name, .. } => vec![EntityRef::new(EntityType::Concept, name.as_str())],
        Episode::LlmCall {
            provider, model, ..
        } => llm_call_entities(provider.as_deref(), model.as_deref()),
        Episode::BusPublish { subject } | Episode::BusReceive { subject } => {
            agent_from_subject(subject)
        }
        _ => Vec::new(),
    }
}

fn llm_call_entities(provider: Option<&str>, model: Option<&str>) -> Vec<EntityRef> {
    let mut out = Vec::new();
    if let Some(p) = provider {
        if !p.is_empty() {
            out.push(EntityRef::new(EntityType::Concept, p));
        }
    }
    if let Some(m) = model {
        if !m.is_empty() {
            out.push(EntityRef::new(EntityType::Concept, m));
        }
    }
    out
}

/// Recognised conventions: `agent.<name>.*` and `klieo.agent.<name>.*`.
/// Any other subject shape (e.g. `klieo.a2a.task.<id>`,
/// `klieo.ops.read`) yields the empty set — the projector does not
/// guess agent identity from free-form topics.
fn agent_from_subject(subject: &str) -> Vec<EntityRef> {
    for prefix in AGENT_SUBJECT_PREFIXES {
        if let Some(rest) = subject.strip_prefix(prefix) {
            let name = rest.split('.').next().unwrap_or("");
            if !name.is_empty() {
                return vec![EntityRef::new(EntityType::Agent, name)];
            }
        }
    }
    Vec::new()
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::memory::ToolResult;
    use serde_json::json;

    fn first(entities: &[EntityRef]) -> &EntityRef {
        entities.first().expect("expected at least one entity")
    }

    #[test]
    fn started_produces_agent_entity() {
        let entities = entities_for_episode(&Episode::Started {
            agent: "alpha".into(),
        });
        assert_eq!(entities.len(), 1);
        assert_eq!(first(&entities).entity_type, EntityType::Agent);
        assert_eq!(first(&entities).name, "alpha");
    }

    #[test]
    fn tool_call_produces_concept_entity_lowercased() {
        let entities = entities_for_episode(&Episode::ToolCall {
            name: "WebSearch".into(),
            args: json!({}),
            result: ToolResult::ok(json!("hit")),
        });
        assert_eq!(entities.len(), 1);
        assert_eq!(first(&entities).entity_type, EntityType::Concept);
        assert_eq!(first(&entities).name, "websearch");
    }

    #[test]
    fn llm_call_with_provider_and_model_produces_two_concepts() {
        let entities = entities_for_episode(&Episode::LlmCall {
            tokens: 12,
            latency_ms: 30,
            provider: Some("ollama".into()),
            model: Some("qwen2.5:14b".into()),
            prompt_tokens: None,
            completion_tokens: None,
        });
        assert_eq!(entities.len(), 2);
        assert!(entities
            .iter()
            .any(|e| e.entity_type == EntityType::Concept && e.name == "ollama"));
        assert!(entities
            .iter()
            .any(|e| e.entity_type == EntityType::Concept && e.name == "qwen2.5:14b"));
    }

    #[test]
    fn llm_call_with_only_model_produces_one_concept() {
        let entities = entities_for_episode(&Episode::LlmCall {
            tokens: 0,
            latency_ms: 0,
            provider: None,
            model: Some("gpt-4o-mini".into()),
            prompt_tokens: None,
            completion_tokens: None,
        });
        assert_eq!(entities.len(), 1);
        assert_eq!(first(&entities).name, "gpt-4o-mini");
    }

    #[test]
    fn llm_call_with_neither_provider_nor_model_is_empty() {
        let entities = entities_for_episode(&Episode::llm_call(42, 17));
        assert!(entities.is_empty());
    }

    #[test]
    fn bus_publish_with_agent_prefix_extracts_agent_name() {
        let entities = entities_for_episode(&Episode::BusPublish {
            subject: "agent.alpha.heartbeat".into(),
        });
        assert_eq!(entities.len(), 1);
        assert_eq!(first(&entities).entity_type, EntityType::Agent);
        assert_eq!(first(&entities).name, "alpha");
    }

    #[test]
    fn bus_receive_with_klieo_agent_prefix_extracts_agent_name() {
        let entities = entities_for_episode(&Episode::BusReceive {
            subject: "klieo.agent.bravo.task.assigned".into(),
        });
        assert_eq!(entities.len(), 1);
        assert_eq!(first(&entities).entity_type, EntityType::Agent);
        assert_eq!(first(&entities).name, "bravo");
    }

    #[test]
    fn bus_subject_without_agent_prefix_is_empty() {
        let entities = entities_for_episode(&Episode::BusPublish {
            subject: "klieo.a2a.task.t-1".into(),
        });
        assert!(
            entities.is_empty(),
            "non-agent subjects must produce no entities"
        );

        let entities = entities_for_episode(&Episode::BusReceive {
            subject: "klieo.ops.read".into(),
        });
        assert!(entities.is_empty());
    }

    #[test]
    fn bus_subject_with_prefix_but_empty_name_is_empty() {
        let entities = entities_for_episode(&Episode::BusPublish {
            subject: "agent..stale".into(),
        });
        assert!(
            entities.is_empty(),
            "agent.<empty>.* must produce no entities"
        );
    }

    #[test]
    fn completed_produces_no_entities() {
        assert!(entities_for_episode(&Episode::Completed).is_empty());
    }

    #[test]
    fn failed_produces_no_entities() {
        assert!(entities_for_episode(&Episode::Failed {
            error: "boom".into(),
        })
        .is_empty());
    }

    #[test]
    fn ops_episode_produces_no_entities() {
        assert!(entities_for_episode(&Episode::Ops(json!({"event": "noop"}))).is_empty());
    }

    #[test]
    fn summary_checkpoint_produces_no_entities() {
        assert!(entities_for_episode(&Episode::SummaryCheckpoint {
            input_message_count: 0,
            summary_chars: 0,
            latency_ms: 0,
            tokens: 0,
        })
        .is_empty());
    }
}