poimen-memory/crates/mem-core/src/entity.rs

/// Entity domain model for temporal graph-RAG.
/// Single Responsibility: Entity identity and metadata.
/// Open/Closed: EntityType enum extensible.
/// Dependencies: Uses time::OffsetDateTime (consistent with mem-core).

use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use std::fmt;

/// Entity type classification (extensible enum).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum EntityType {
    Person,
    Tool,
    Concept,
    Location,
    Event,
    Organization,
    /// Agent prompt template tracked as a first-class entity.
    /// Enables the agent to learn which prompts produce good results.
    AgentPrompt,
    /// Agent skill — a reusable capability the agent has learned.
    AgentSkill,
    /// Agent decision — a recorded choice with reasoning and outcome.
    AgentDecision,
    Unknown,
}

impl EntityType {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Person => "person",
            Self::Tool => "tool",
            Self::Concept => "concept",
            Self::Location => "location",
            Self::Event => "event",
            Self::Organization => "organization",
            Self::AgentPrompt => "agent_prompt",
            Self::AgentSkill => "agent_skill",
            Self::AgentDecision => "agent_decision",
            Self::Unknown => "unknown",
        }
    }

    pub fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "person" => Self::Person,
            "tool" => Self::Tool,
            "concept" => Self::Concept,
            "location" => Self::Location,
            "event" => Self::Event,
            "organization" => Self::Organization,
            "agent_prompt" => Self::AgentPrompt,
            "agent_skill" => Self::AgentSkill,
            "agent_decision" => Self::AgentDecision,
            _ => Self::Unknown,
        }
    }
}

impl fmt::Display for EntityType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

/// Entity: Named concept in the knowledge graph.
/// Dependency Inversion: Depends on abstractions (String for id, OffsetDateTime for time).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Entity {
    pub id: String,  // UUID as string for serialization
    pub project_id: String,
    
    // Identity
    pub name: String,
    pub name_embedding: Option<Vec<f32>>,
    pub summary: Option<String>,
    pub summary_embedding: Option<Vec<f32>>,
    pub entity_type: EntityType,
    
    // Temporal (transaction time)
    #[serde(with = "time::serde::rfc3339")]
    pub t_created: OffsetDateTime,
    #[serde(with = "time::serde::rfc3339::option")]
    pub t_expired: Option<OffsetDateTime>,
    
    // Provenance: Which episodes mention this entity
    pub source_episodes: Vec<i64>,  // memory_node.id references
    
    // Access tracking (for LRU)
    pub access_count: i64,
    #[serde(with = "time::serde::rfc3339::option")]
    pub last_accessed: Option<OffsetDateTime>,
    
    // Community reference (nullable until Phase 4)
    pub community_id: Option<String>,  // UUID as string
}

impl Entity {
    /// Create new entity with minimal fields.
    pub fn new(project_id: &str, name: &str, entity_type: EntityType) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            project_id: project_id.to_string(),
            name: name.to_string(),
            name_embedding: None,
            summary: None,
            summary_embedding: None,
            entity_type,
            t_created: OffsetDateTime::now_utc(),
            t_expired: None,
            source_episodes: vec![],
            access_count: 0,
            last_accessed: None,
            community_id: None,
        }
    }

    /// Is this entity currently active (not soft-deleted)?
    pub fn is_active(&self) -> bool {
        self.t_expired.is_none()
    }

    /// Builder pattern: Set summary.
    pub fn with_summary(mut self, summary: &str) -> Self {
        self.summary = Some(summary.to_string());
        self
    }

    /// Builder pattern: Set embedding.
    pub fn with_name_embedding(mut self, embedding: Vec<f32>) -> Self {
        self.name_embedding = Some(embedding);
        self
    }

    /// Builder pattern: Set summary embedding.
    pub fn with_summary_embedding(mut self, embedding: Vec<f32>) -> Self {
        self.summary_embedding = Some(embedding);
        self
    }

    /// Builder pattern: Link source episode.
    pub fn with_source_episode(mut self, episode_id: i64) -> Self {
        if !self.source_episodes.contains(&episode_id) {
            self.source_episodes.push(episode_id);
        }
        self
    }

    /// Builder pattern: Set confidence (unused for now, placeholder for extraction)
    pub fn with_confidence(self, _confidence: f32) -> Self {
        self // Placeholder for extraction confidence
    }

    /// DRY: Normalized name for deduplication
    pub fn name_normalized(&self) -> String {
        self.name.to_lowercase().trim().to_string()
    }
}

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

    #[test]
    fn test_entity_creation() {
        let entity = Entity::new("proj1", "Kubernetes", EntityType::Tool);
        assert_eq!(entity.name, "Kubernetes");
        assert_eq!(entity.entity_type, EntityType::Tool);
        assert!(entity.is_active());
        assert_eq!(entity.access_count, 0);
    }

    #[test]
    fn test_entity_builder_pattern() {
        let entity = Entity::new("proj1", "Rock", EntityType::Person)
            .with_summary("SRE and Rust developer")
            .with_name_embedding(vec![0.1, 0.2, 0.3]);

        assert_eq!(entity.summary, Some("SRE and Rust developer".to_string()));
        assert_eq!(entity.name_embedding.as_ref().unwrap().len(), 3);
    }

    #[test]
    fn test_entity_type_round_trip() {
        for ty in &[
            EntityType::Person,
            EntityType::Tool,
            EntityType::Concept,
            EntityType::AgentPrompt,
            EntityType::AgentSkill,
            EntityType::AgentDecision,
        ] {
            let s = ty.as_str();
            assert_eq!(EntityType::from_str(s), *ty);
        }
    }

    #[test]
    fn test_entity_normalized_name() {
        let entity = Entity::new("proj1", "  Kubernetes  ", EntityType::Tool);
        assert_eq!(entity.name_normalized(), "kubernetes");
    }

    #[test]
    fn test_entity_serialization() {
        let entity = Entity::new("proj1", "Test", EntityType::Concept);
        let json = serde_json::to_string(&entity).unwrap();
        let deserialized: Entity = serde_json::from_str(&json).unwrap();
        assert_eq!(entity.name, deserialized.name);
        assert_eq!(entity.entity_type, deserialized.entity_type);
    }
}