use crate::observation::Observation;
use crate::project::ProjectId;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use ulid::Ulid;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct EntityId(pub Ulid);
impl EntityId {
pub fn new() -> Self {
Self(Ulid::new())
}
pub fn from_string(s: &str) -> Result<Self, ulid::DecodeError> {
Ok(Self(Ulid::from_string(s)?))
}
}
impl Default for EntityId {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Display for EntityId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct EntityType(pub String);
impl EntityType {
pub fn new(s: impl Into<String>) -> Self {
Self(s.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<&str> for EntityType {
fn from(s: &str) -> Self {
Self(s.to_string())
}
}
impl From<String> for EntityType {
fn from(s: String) -> Self {
Self(s)
}
}
impl From<&String> for EntityType {
fn from(s: &String) -> Self {
Self(s.clone())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Entity {
pub id: EntityId,
pub project_id: ProjectId,
pub name: String,
pub entity_type: EntityType,
pub observations: Vec<Observation>,
pub tags: Vec<String>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
#[serde(skip_serializing_if = "Option::is_none")]
pub embedding: Option<Vec<f32>>,
}
impl Entity {
pub fn new(
project_id: ProjectId,
name: impl Into<String>,
entity_type: impl Into<EntityType>,
) -> Self {
let now = Utc::now();
Self {
id: EntityId::new(),
project_id,
name: name.into(),
entity_type: entity_type.into(),
observations: Vec::new(),
tags: Vec::new(),
metadata: HashMap::new(),
created_at: now,
updated_at: now,
embedding: None,
}
}
pub fn add_observation(&mut self, content: impl Into<String>) -> &Observation {
let obs = Observation::new(content);
self.observations.push(obs);
self.updated_at = Utc::now();
self.observations.last().expect("observations cannot be empty after push")
}
pub fn add_tag(&mut self, tag: impl Into<String>) {
let tag = tag.into();
if !self.tags.contains(&tag) {
self.tags.push(tag);
self.updated_at = Utc::now();
}
}
pub fn remove_tag(&mut self, tag: &str) -> bool {
if let Some(pos) = self.tags.iter().position(|t| t == tag) {
self.tags.remove(pos);
self.updated_at = Utc::now();
true
} else {
false
}
}
pub fn has_tag(&self, tag: &str) -> bool {
self.tags.iter().any(|t| t == tag)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NewEntity {
pub name: String,
pub entity_type: String,
pub observations: Vec<String>,
#[serde(default)]
pub tags: Vec<String>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
}
impl NewEntity {
pub fn new(name: impl Into<String>, entity_type: impl Into<String>) -> Self {
Self {
name: name.into(),
entity_type: entity_type.into(),
observations: Vec::new(),
tags: Vec::new(),
metadata: HashMap::new(),
}
}
pub fn with_observation(mut self, obs: impl Into<String>) -> Self {
self.observations.push(obs.into());
self
}
pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
self.tags.push(tag.into());
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_entity_creation() {
let project_id = ProjectId::new();
let entity = Entity::new(project_id, "John_Smith", "person");
assert_eq!(entity.name, "John_Smith");
assert_eq!(entity.entity_type.as_str(), "person");
assert!(entity.observations.is_empty());
assert!(entity.tags.is_empty());
}
#[test]
fn test_add_observation() {
let project_id = ProjectId::new();
let mut entity = Entity::new(project_id, "John_Smith", "person");
entity.add_observation("Works at Google");
assert_eq!(entity.observations.len(), 1);
assert_eq!(entity.observations[0].content, "Works at Google");
}
#[test]
fn test_tags() {
let project_id = ProjectId::new();
let mut entity = Entity::new(project_id, "John_Smith", "person");
entity.add_tag("technical");
entity.add_tag("mentor");
entity.add_tag("technical");
assert_eq!(entity.tags.len(), 2);
assert!(entity.has_tag("technical"));
assert!(entity.has_tag("mentor"));
entity.remove_tag("mentor");
assert!(!entity.has_tag("mentor"));
}
}