2025-04-13 14:52:26 -07:00

69 lines
2.2 KiB
Python

from typing import Dict, Any, List, Optional
from datetime import UTC, datetime
from pydantic import BaseModel, Field
import uuid
class Entity(BaseModel):
"""Represents an entity in a knowledge graph"""
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
label: str
type: str
properties: Dict[str, Any] = Field(default_factory=dict)
document_ids: List[str] = Field(default_factory=list)
chunk_sources: Dict[str, List[int]] = Field(default_factory=dict)
def __hash__(self):
return hash(self.id)
def __eq__(self, other):
if not isinstance(other, Entity):
return False
return self.id == other.id
class Relationship(BaseModel):
"""Represents a relationship between entities in a knowledge graph"""
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
source_id: str
target_id: str
type: str
document_ids: List[str] = Field(default_factory=list)
chunk_sources: Dict[str, List[int]] = Field(default_factory=dict)
def __hash__(self):
return hash(self.id)
def __eq__(self, other):
if not isinstance(other, Relationship):
return False
return self.id == other.id
class Graph(BaseModel):
"""Represents a knowledge graph"""
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
name: str
entities: List[Entity] = Field(default_factory=list)
relationships: List[Relationship] = Field(default_factory=list)
metadata: Dict[str, Any] = Field(default_factory=dict)
system_metadata: Dict[str, Any] = Field(
default_factory=lambda: {
"created_at": datetime.now(UTC),
"updated_at": datetime.now(UTC),
"folder_name": None,
"end_user_id": None,
}
)
document_ids: List[str] = Field(default_factory=list)
filters: Optional[Dict[str, Any]] = None
created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
owner: Dict[str, str] = Field(default_factory=dict)
access_control: Dict[str, List[str]] = Field(
default_factory=lambda: {"readers": [], "writers": [], "admins": []}
)