mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
phernandez
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
"""
|
|
Core pydantic models for basic-memory entities, observations, and relations.
|
|
These models define the schema for our core data types while remaining
|
|
independent from storage/persistence concerns.
|
|
"""
|
|
|
|
from datetime import datetime, UTC
|
|
from uuid import uuid4
|
|
from typing import List, Optional, ForwardRef, Dict, Any
|
|
from pydantic import BaseModel, model_validator
|
|
|
|
|
|
class Observation(BaseModel):
|
|
"""An atomic piece of information about an entity."""
|
|
content: str
|
|
|
|
|
|
class Relation(BaseModel):
|
|
"""
|
|
Represents a directed edge between entities in the knowledge graph.
|
|
Relations are always stored in active voice (e.g. "created", "teaches", etc.)
|
|
"""
|
|
id: str
|
|
from_entity: 'Entity'
|
|
to_entity: 'Entity'
|
|
relation_type: str
|
|
context: Optional[str] = None
|
|
|
|
@model_validator(mode='before')
|
|
@classmethod
|
|
def generate_id_if_needed(cls, data: dict) -> dict:
|
|
"""Generate an ID if one wasn't provided"""
|
|
if not data.get('id'):
|
|
data['id'] = f"rel-{uuid4().hex[:8]}"
|
|
return data
|
|
|
|
def model_dump(self, **kwargs) -> Dict[str, Any]:
|
|
"""Serialize to storage format with entity IDs"""
|
|
return {
|
|
'id': self.id,
|
|
'from_id': self.from_entity.id,
|
|
'to_id': self.to_entity.id,
|
|
'relation_type': self.relation_type,
|
|
'context': self.context
|
|
}
|
|
|
|
|
|
class Entity(BaseModel):
|
|
"""
|
|
Represents a node in our knowledge graph - could be a person, project,
|
|
concept, etc. Each entity has a unique name, a type, and a list of
|
|
associated observations.
|
|
"""
|
|
id: str
|
|
name: str
|
|
entity_type: str
|
|
observations: List[Observation] = []
|
|
relations: List[Relation] = []
|
|
|
|
@model_validator(mode='before')
|
|
@classmethod
|
|
def generate_id_if_needed(cls, data: dict) -> dict:
|
|
"""Generate an ID if one wasn't provided during instantiation"""
|
|
if not data.get('id') and data.get('name'):
|
|
timestamp = datetime.now(UTC).strftime("%Y%m%d")
|
|
normalized_name = data['name'].lower().replace(" ", "-")
|
|
data['id'] = f"{timestamp}-{normalized_name}-{uuid4().hex[:8]}"
|
|
return data
|
|
|
|
def model_dump(self, **kwargs) -> Dict[str, Any]:
|
|
"""Serialize entity, handling relations to prevent circular references"""
|
|
# Get basic data without relations
|
|
exclude = kwargs.pop('exclude', set())
|
|
exclude.add('relations')
|
|
basic_data = super().model_dump(exclude=exclude, **kwargs)
|
|
|
|
# Add serialized relations if we have any
|
|
if 'relations' not in exclude and self.relations:
|
|
basic_data['relations'] = [
|
|
relation.model_dump(**kwargs)
|
|
for relation in self.relations
|
|
]
|
|
|
|
return basic_data
|
|
|
|
|
|
# Update forward refs
|
|
Entity.model_rebuild()
|
|
Relation.model_rebuild() |