From ea359d2484cfc18db55a5e498cc609e07de5ee76 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 6 Jan 2025 18:57:04 -0600 Subject: [PATCH] update entity/repostitory/service --- .../api/routers/knowledge_router.py | 13 -- src/basic_memory/models/documents.py | 5 - src/basic_memory/models/knowledge.py | 64 +++++---- .../repository/entity_repository.py | 40 ------ src/basic_memory/schemas/base.py | 21 ++- src/basic_memory/schemas/response.py | 1 + src/basic_memory/services/entity_service.py | 6 +- tests/conftest.py | 4 +- tests/repository/test_entity_repository.py | 127 +++--------------- tests/services/test_entity_service.py | 87 ++++++------ 10 files changed, 115 insertions(+), 253 deletions(-) diff --git a/src/basic_memory/api/routers/knowledge_router.py b/src/basic_memory/api/routers/knowledge_router.py index e19d5818..f59f767b 100644 --- a/src/basic_memory/api/routers/knowledge_router.py +++ b/src/basic_memory/api/routers/knowledge_router.py @@ -100,19 +100,6 @@ async def get_entity(path_id: PathId, entity_service: EntityServiceDep) -> Entit raise HTTPException(status_code=404, detail=f"Entity with {path_id} not found") -@router.post("/search", response_model=SearchNodesResponse) -async def search_nodes( - data: SearchNodesRequest, entity_service: EntityServiceDep -) -> SearchNodesResponse: - """Search for entities in the knowledge graph.""" - logger.debug(f"Searching nodes with query: {data.query}") - matches = await entity_service.search(data.query) - logger.debug(f"Found {len(matches)} matches for '{data.query}'") - - return SearchNodesResponse( - matches=[EntityResponse.model_validate(entity) for entity in matches], query=data.query - ) - @router.post("/nodes", response_model=EntityListResponse) async def open_nodes(data: OpenNodesRequest, entity_service: EntityServiceDep) -> EntityListResponse: diff --git a/src/basic_memory/models/documents.py b/src/basic_memory/models/documents.py index 6b1c5086..da7cec8b 100644 --- a/src/basic_memory/models/documents.py +++ b/src/basic_memory/models/documents.py @@ -40,10 +40,5 @@ class Document(Base): DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") ) - # Relationships - entities: Mapped[List["Entity"]] = relationship( # pyright: ignore [reportUndefinedVariable] # noqa: F821 - "Entity", back_populates="document", cascade="all, delete-orphan" - ) - def __repr__(self) -> str: return f"Document(id={self.id}, path_id='{self.path_id}', checksum='{self.checksum}', created_at='{self.created_at}', updated_at='{self.updated_at}')" diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 2275e505..75e23a63 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -3,47 +3,69 @@ from datetime import datetime from typing import Optional -from sqlalchemy import Integer, String, Text, ForeignKey, UniqueConstraint, text, DateTime, Index +from sqlalchemy import ( + Integer, + String, + Text, + ForeignKey, + UniqueConstraint, + text, + DateTime, + Index, + JSON, + CheckConstraint, +) from sqlalchemy.orm import Mapped, mapped_column, relationship from basic_memory.models.base import Base -from basic_memory.models.documents import Document from enum import Enum +class EntityType(str, Enum): + """Types of knowledge nodes.""" + + KNOWLEDGE = "knowledge" + NOTE = "note" + + class Entity(Base): """ Core entity in the knowledge graph. Entities represent semantic nodes maintained by the AI layer. Each entity: - Has a unique numeric ID (database-generated) - - Maps to a document file on disk (optional) + - Maps to a file on disk - Maintains a checksum for change detection - Tracks both source document and semantic properties """ __tablename__ = "entity" __table_args__ = ( - UniqueConstraint("entity_type", "name", name="uix_entity_type_name"), UniqueConstraint("path_id", name="uix_entity_path_id"), # Make path_id unique Index("ix_entity_type", "entity_type"), - Index("ix_entity_doc_id", "doc_id"), Index("ix_entity_created_at", "created_at"), # For timeline queries - Index("ix_entity_updated_at", "updated_at") # For timeline queries + Index("ix_entity_updated_at", "updated_at"), # For timeline queries + CheckConstraint( + f"entity_type IN {tuple(t.value for t in EntityType)}", name="check_entity_type" + ), ) # Core identity id: Mapped[int] = mapped_column(Integer, primary_key=True) name: Mapped[str] = mapped_column(String) - entity_type: Mapped[str] = mapped_column(String) - # Normalized path for URIs - must be unique - path_id: Mapped[str] = mapped_column(String, index=True) + entity_type: Mapped[EntityType] = mapped_column(String, default=EntityType.KNOWLEDGE) + entity_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) + + + # Normalized path for URIs + path_id: Mapped[str] = mapped_column(String, unique=True, index=True) # Actual filesystem relative path file_path: Mapped[str] = mapped_column(String, unique=True, index=True) - - # Content and validation - description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + # checksum of file checksum: Mapped[Optional[str]] = mapped_column(String, nullable=True) + + # Content for knowledge entity_type + description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # Metadata and tracking created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) @@ -51,11 +73,6 @@ class Entity(Base): DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") ) - # Relations - doc_id: Mapped[Optional[int]] = mapped_column( - Integer, ForeignKey("document.id", ondelete="SET NULL"), nullable=True - ) - # Relationships observations = relationship( "Observation", back_populates="entity", cascade="all, delete-orphan" @@ -72,7 +89,6 @@ class Entity(Base): foreign_keys="[Relation.to_id]", cascade="all, delete-orphan", ) - document: Mapped[Optional[Document]] = relationship(Document, back_populates="entities") @property def relations(self): @@ -99,8 +115,8 @@ class Observation(Base): __tablename__ = "observation" __table_args__ = ( - Index("ix_observation_entity_id", "entity_id"), # Add FK index - Index("ix_observation_category", "category"), # Add category index + Index("ix_observation_entity_id", "entity_id"), # Add FK index + Index("ix_observation_category", "category"), # Add category index Index("ix_observation_created_at", "created_at"), # For timeline queries Index("ix_observation_updated_at", "updated_at"), # For timeline queries ) @@ -112,7 +128,7 @@ class Observation(Base): String, nullable=False, default=ObservationCategory.NOTE.value, - server_default=ObservationCategory.NOTE.value + server_default=ObservationCategory.NOTE.value, ) context: Mapped[str] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) @@ -153,8 +169,10 @@ class Relation(Base): ) # Relationships - from_entity = relationship("Entity", foreign_keys=[from_id], back_populates="outgoing_relations") + from_entity = relationship( + "Entity", foreign_keys=[from_id], back_populates="outgoing_relations" + ) to_entity = relationship("Entity", foreign_keys=[to_id], back_populates="incoming_relations") def __repr__(self) -> str: - return f"Relation(id={self.id}, from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')" \ No newline at end of file + return f"Relation(id={self.id}, from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')" diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 18e6d2cf..423528c0 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -26,7 +26,6 @@ class EntityRepository(Repository[Entity]): async def list_entities( self, entity_type: Optional[str] = None, - doc_id: Optional[int] = None, sort_by: Optional[str] = "updated_at", include_related: bool = False, ) -> Sequence[Entity]: @@ -52,9 +51,6 @@ class EntityRepository(Repository[Entity]): else: query = query.where(Entity.entity_type == entity_type) - if doc_id: - query = query.where(Entity.doc_id == doc_id) - # Apply sorting if sort_by: sort_field = getattr(Entity, sort_by, Entity.updated_at) @@ -63,42 +59,6 @@ class EntityRepository(Repository[Entity]): result = await self.execute_query(query) return list(result.scalars().all()) - async def get_entity_types(self) -> List[str]: - """Get list of distinct entity types.""" - query = select(Entity.entity_type).distinct() - - result = await self.execute_query(query, use_query_options=False) - return list(result.scalars().all()) - - async def search(self, query_str: str) -> List[Entity]: - """ - Search for entities. - - Searches across: - - Entity names - - Entity types - - Entity descriptions - - Associated Observations content - """ - search_term = f"%{query_str}%" - query = ( - self.select() - .where( - or_( - Entity.name.ilike(search_term), - Entity.entity_type.ilike(search_term), - Entity.description.ilike(search_term), - Entity.observations.any(Observation.content.ilike(search_term)), - ) - ) - .options(*self.get_load_options()) - ) - result = await self.execute_query(query) - return list(result.scalars().all()) - - async def delete_entities_by_doc_id(self, doc_id: int) -> bool: - """Delete all entities associated with a document.""" - return await self.delete_by_fields(doc_id=doc_id) async def delete_by_file_path(self, file_path: str) -> bool: """Delete entity with the provided file_path.""" diff --git a/src/basic_memory/schemas/base.py b/src/basic_memory/schemas/base.py index 431224e9..09921545 100644 --- a/src/basic_memory/schemas/base.py +++ b/src/basic_memory/schemas/base.py @@ -34,10 +34,10 @@ Common Relation Types: import re from enum import Enum -from typing import List, Optional, Annotated +from typing import List, Optional, Annotated, Dict from annotated_types import MinLen, MaxLen -from pydantic import BaseModel, BeforeValidator +from pydantic import BaseModel, BeforeValidator, Field def to_snake_case(name: str) -> str: @@ -122,17 +122,15 @@ Examples: - "Depends on SQLAlchemy for database operations" """ -EntityType = Annotated[str, BeforeValidator(to_snake_case), MinLen(1), MaxLen(200)] -"""Classification of entity (e.g., 'person', 'project', 'concept'). -The type serves multiple purposes: -1. Organizes entities in the filesystem -2. Enables filtering and querying -3. Provides context for relations -4. Helps generate meaningful IDs +class EntityType(str, Enum): + """Type of entity. -Common types are listed in the module docstring. -""" + - knowledge: Contain information used in the semantic graph + - note: Free form information + """ + KNOWLEDGE = "knowledge" + NOTE = "note" RelationType = Annotated[str, BeforeValidator(to_snake_case), MinLen(1), MaxLen(200)] """Type of relationship between entities. Always use active voice present tense. @@ -252,6 +250,7 @@ class Entity(BaseModel): name: str entity_type: EntityType + entity_metadata: Optional[Dict] = Field(default=None, description="Optional metadata") description: Optional[str] = None observations: List[Observation] = [] diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index c2e55603..d11ba3d9 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -121,6 +121,7 @@ class EntityResponse(SQLAlchemyModel): path_id: PathId name: str entity_type: EntityType + entity_metadata: Optional[Dict] = None description: Optional[str] = None observations: List[ObservationResponse] = [] relations: List[RelationResponse] = [] diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 20015e4f..e180e209 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -15,6 +15,7 @@ def entity_model(entity: EntitySchema): model = EntityModel( name=entity.name, entity_type=entity.entity_type, + entity_metadata=entity.entity_metadata, path_id=entity.path_id, file_path=entity.file_path, description=entity.description, @@ -29,11 +30,6 @@ class EntityService(BaseService[EntityModel]): def __init__(self, entity_repository: EntityRepository): super().__init__(entity_repository) - async def search(self, query: str) -> Sequence[EntityModel]: - """Search entities using LIKE pattern matching.""" - logger.debug(f"Searching entities with query: {query}") - return await self.repository.search(query) - async def create_entity(self, entity: EntitySchema) -> EntityModel: """Create a new entity in the database.""" logger.debug(f"Creating entity in DB: {entity}") diff --git a/tests/conftest.py b/tests/conftest.py index 6497d087..c66b3fe1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,7 +13,7 @@ from basic_memory.db import DatabaseType from basic_memory.markdown.knowledge_parser import KnowledgeParser from basic_memory.markdown.knowledge_writer import KnowledgeWriter from basic_memory.models import Base -from basic_memory.models.knowledge import Entity +from basic_memory.models.knowledge import Entity, EntityType from basic_memory.repository.document_repository import DocumentRepository from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.observation_repository import ObservationRepository @@ -243,7 +243,7 @@ async def sample_entity(entity_repository: EntityRepository) -> Entity: """Create a sample entity for testing.""" entity_data = { "name": "Test Entity", - "entity_type": "test", + "entity_type": EntityType.KNOWLEDGE, "description": "A test entity", "path_id": "test/test_entity", "file_path": "test/test_entity.md", diff --git a/tests/repository/test_entity_repository.py b/tests/repository/test_entity_repository.py index 6347cff3..e8297c3a 100644 --- a/tests/repository/test_entity_repository.py +++ b/tests/repository/test_entity_repository.py @@ -9,6 +9,7 @@ from sqlalchemy.exc import IntegrityError from basic_memory import db from basic_memory.models import Entity, Observation, Relation +from basic_memory.models.knowledge import EntityType from basic_memory.repository.entity_repository import EntityRepository @@ -30,14 +31,14 @@ async def related_entities(session_maker): async with db.scoped_session(session_maker) as session: source = Entity( name="source", - entity_type="source", + entity_type=EntityType.KNOWLEDGE, path_id="source/source", file_path="source/source.md", description="Source entity", ) target = Entity( name="target", - entity_type="target", + entity_type=EntityType.KNOWLEDGE, path_id="target/target", file_path="target/target.md", description="Target entity", @@ -57,7 +58,7 @@ async def test_create_entity(entity_repository: EntityRepository): """Test creating a new entity""" entity_data = { "name": "Test", - "entity_type": "test", + "entity_type": EntityType.KNOWLEDGE, "path_id": "test/test", "file_path": "test/test.md", "description": "Test description", @@ -90,14 +91,14 @@ async def test_create_all(entity_repository: EntityRepository): entity_data = [ { "name": "Test_1", - "entity_type": "test", + "entity_type": EntityType.KNOWLEDGE, "path_id": "test/test_1", "file_path": "test/test_1.md", "description": "Test description", }, { "name": "Test-2", - "entity_type": "test", + "entity_type": EntityType.KNOWLEDGE, "path_id": "test/test_2", "file_path": "test/test_2.md", "description": "Test description", @@ -121,40 +122,12 @@ async def test_create_all(entity_repository: EntityRepository): assert len(entity.relations) == 0 -@pytest.mark.asyncio -async def test_entity_type_name_unique_constraint(entity_repository: EntityRepository): - """Test the unique constraint on entity_type + name combination.""" - # Create first entity - entity1_data = { - "name": "Test Entity", - "entity_type": "type1", - "path_id": "type1/test_entity", - "file_path": "type1/test_entity1.md", - "description": "First entity", - } - await entity_repository.create(entity1_data) - - # Try to create another entity with same type and name - entity2_data = { - "name": "Test Entity", # Same name - "entity_type": "type1", # Same type - "path_id": "type1/test_entity", - "file_path": "type1/test_entity2.md", - "description": "Second entity", - } - - # Should raise IntegrityError - with pytest.raises(IntegrityError) as exc_info: - await entity_repository.create(entity2_data) - assert "UNIQUE constraint failed" in str(exc_info.value) - - @pytest.mark.asyncio async def test_create_entity_null_description(session_maker, entity_repository: EntityRepository): """Test creating an entity with null description""" entity_data = { "name": "Test", - "entity_type": "test", + "entity_type": EntityType.KNOWLEDGE, "path_id": "test/test", "file_path": "test/test.md", "description": None, @@ -296,61 +269,6 @@ async def test_delete_nonexistent_entity(entity_repository: EntityRepository): assert result is False -@pytest.mark.asyncio -async def test_search(session_maker, entity_repository: EntityRepository): - """Test searching entities""" - # First create and commit the entities - async with db.scoped_session(session_maker) as session: - entity1 = Entity( - name="Search Test 1", - entity_type="test", - path_id="test/search_test_1", - file_path="test/search_test_1.md", - description="First test entity", - ) - entity2 = Entity( - name="Search Test 2", - entity_type="other", - path_id="other/search_test_2", - file_path="other/search_test_2.md", - description="Second test entity", - ) - session.add_all([entity1, entity2]) - - # Then add observations in a new transaction - async with db.scoped_session(session_maker) as session: - ts = datetime.now(UTC) - stmt = text(""" - INSERT INTO observation (entity_id, content, created_at) - VALUES (:e1_id, :e1_obs, :ts), (:e2_id, :e2_obs, :ts) - """) - await session.execute( - stmt, - { - "e1_id": entity1.id, - "e1_obs": "First observation with searchable content", - "e2_id": entity2.id, - "e2_obs": "Another observation to find", - "ts": ts, - }, - ) - - # Test search by name - results = await entity_repository.search("Search Test") - assert len(results) == 2 - names = {e.name for e in results} - assert "Search Test 1" in names - assert "Search Test 2" in names - - # Test search by type - results = await entity_repository.search("other") - assert len(results) == 1 - assert results[0].entity_type == "other" - - # Test search by observation content - results = await entity_repository.search("searchable") - assert len(results) == 1 - assert results[0].id == entity1.id @pytest_asyncio.fixture @@ -360,21 +278,21 @@ async def test_entities(session_maker): entities = [ Entity( name="entity1", - entity_type="type1", + entity_type=EntityType.KNOWLEDGE, description="First test entity", path_id="type1/entity1", file_path="type1/entity1.md", ), Entity( name="entity2", - entity_type="type1", + entity_type=EntityType.KNOWLEDGE, description="Second test entity", path_id="type1/entity2", file_path="type1/entity2.md", ), Entity( name="entity3", - entity_type="type2", + entity_type=EntityType.KNOWLEDGE, description="Third test entity", path_id="type2/entity3", file_path="type2/entity3.md", @@ -472,14 +390,14 @@ async def test_list_entities_with_related(entity_repository: EntityRepository, s # Core entities core = Entity( name="core_service", - entity_type="service", + entity_type=EntityType.NOTE, path_id="service/core", file_path="service/core.md", description="Core service" ) dbe = Entity( name="db_service", - entity_type="service", + entity_type=EntityType.KNOWLEDGE, path_id="service/db", file_path="service/db.md", description="Database service" @@ -487,7 +405,7 @@ async def test_list_entities_with_related(entity_repository: EntityRepository, s # Related entity of different type config = Entity( name="service_config", - entity_type="configuration", + entity_type=EntityType.KNOWLEDGE, path_id="config/service", file_path="config/service.md", description="Service configuration" @@ -504,18 +422,18 @@ async def test_list_entities_with_related(entity_repository: EntityRepository, s ] session.add_all(relations) - # Test 1: List services without related entities + # Test 1: List without related entities services = await entity_repository.list_entities( - entity_type="service", + entity_type=EntityType.KNOWLEDGE, include_related=False ) assert len(services) == 2 service_names = {s.name for s in services} - assert service_names == {"core_service", "db_service"} + assert service_names == {"service_config", "db_service"} # Test 2: List services with related entities services_and_related = await entity_repository.list_entities( - entity_type="service", + entity_type=EntityType.KNOWLEDGE, include_related=True ) assert len(services_and_related) == 3 @@ -527,14 +445,3 @@ async def test_list_entities_with_related(entity_repository: EntityRepository, s core_service = next(e for e in services_and_related if e.name == "core_service") assert len(core_service.outgoing_relations) > 0 # Has incoming relation from config assert len(core_service.incoming_relations) > 0 # Has outgoing relation to db - - # Test 4: List configurations with related - configs = await entity_repository.list_entities( - entity_type="configuration", - include_related=True, - sort_by="name" - ) - config_names = {c.name for c in configs} - # Should include both config and the services it relates to - assert "service_config" in config_names - assert "core_service" in config_names # Related via configures relation \ No newline at end of file diff --git a/tests/services/test_entity_service.py b/tests/services/test_entity_service.py index ea561168..8d65a150 100644 --- a/tests/services/test_entity_service.py +++ b/tests/services/test_entity_service.py @@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from basic_memory.models import Entity as EntityModel from basic_memory.repository.entity_repository import EntityRepository -from basic_memory.schemas import Entity +from basic_memory.schemas import Entity as EntitySchema, EntityType from basic_memory.services.entity_service import EntityService from basic_memory.services.exceptions import EntityNotFoundError @@ -27,9 +27,9 @@ async def entity_service(entity_repository: EntityRepository) -> EntityService: async def test_create_entity(entity_service: EntityService): """Test successful entity creation.""" - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="A test entity description", observations=["this is a test observation"], ) @@ -42,7 +42,7 @@ async def test_create_entity(entity_service: EntityService): assert entity.name == "TestEntity" assert entity.path_id == entity_data.path_id assert entity.file_path == entity_data.file_path - assert entity.entity_type == "test" + assert entity.entity_type == EntityType.KNOWLEDGE assert entity.description == "A test entity description" assert entity.created_at is not None assert entity.observations[0].content == "this is a test observation" @@ -52,7 +52,7 @@ async def test_create_entity(entity_service: EntityService): retrieved = await entity_service.get_by_path_id(entity_data.path_id) assert retrieved.description == "A test entity description" assert retrieved.name == "TestEntity" - assert retrieved.entity_type == "test" + assert retrieved.entity_type == EntityType.KNOWLEDGE assert retrieved.description == "A test entity description" assert retrieved.created_at is not None assert retrieved.observations[0].content == "this is a test observation" @@ -61,15 +61,15 @@ async def test_create_entity(entity_service: EntityService): async def test_create_entities(entity_service: EntityService): """Test successful entity creation.""" entity_data = [ - Entity( + EntitySchema( name="TestEntity1", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="A test entity description", observations=["this is a test observation"], ), - Entity( + EntitySchema( name="TestEntity2", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="A test entity description", observations=["this is a test observation"], ), @@ -83,7 +83,7 @@ async def test_create_entities(entity_service: EntityService): entity1 = entities[0] assert isinstance(entity1, EntityModel) assert entity1.name == "TestEntity1" - assert entity1.entity_type == "test" + assert entity1.entity_type == EntityType.KNOWLEDGE assert entity1.description == "A test entity description" assert entity1.created_at is not None assert entity1.observations[0].content == "this is a test observation" @@ -92,7 +92,7 @@ async def test_create_entities(entity_service: EntityService): entity2 = entities[1] assert isinstance(entity1, EntityModel) assert entity2.name == "TestEntity2" - assert entity2.entity_type == "test" + assert entity2.entity_type == EntityType.KNOWLEDGE assert entity2.description == "A test entity description" assert entity2.created_at is not None assert entity2.observations[0].content == "this is a test observation" @@ -107,18 +107,17 @@ async def test_create_entities(entity_service: EntityService): async def test_get_by_path_id(entity_service: EntityService): """Test finding entity by type and name combination.""" - # Create two entities with same name but different types - entity1_data = Entity( - name="TestEntity", - entity_type="type1", + entity1_data = EntitySchema( + name="TestEntity1", + entity_type=EntityType.KNOWLEDGE, description="First test entity", observations=[], ) entity1 = await entity_service.create_entity(entity1_data) - entity2_data = Entity( - name="TestEntity", # Same name - entity_type="type2", # Different type + entity2_data = EntitySchema( + name="TestEntity2", + entity_type=EntityType.KNOWLEDGE, description="Second test entity", observations=[], ) @@ -128,14 +127,14 @@ async def test_get_by_path_id(entity_service: EntityService): found = await entity_service.get_by_path_id(entity1_data.path_id) assert found is not None assert found.id == entity1.id - assert found.entity_type == "type1" + assert found.entity_type == entity1.entity_type assert found.description == "First test entity" # Find by type2 and name found = await entity_service.get_by_path_id(entity2_data.path_id) assert found is not None assert found.id == entity2.id - assert found.entity_type == "type2" + assert found.entity_type == entity2.entity_type assert found.description == "Second test entity" # Test not found case @@ -145,7 +144,7 @@ async def test_get_by_path_id(entity_service: EntityService): async def test_create_entity_no_description(entity_service: EntityService): """Test creating entity without description (should be None).""" - entity_data = Entity(name="TestEntity", entity_type="test", observations=[]) + entity_data = EntitySchema(name="TestEntity", entity_type=EntityType.KNOWLEDGE, observations=[]) entity = await entity_service.create_entity(entity_data) assert entity.description is None @@ -157,9 +156,9 @@ async def test_create_entity_no_description(entity_service: EntityService): async def test_get_entity_success(entity_service: EntityService): """Test successful entity retrieval.""" - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="Test description", observations=[], ) @@ -170,15 +169,15 @@ async def test_get_entity_success(entity_service: EntityService): assert isinstance(retrieved, EntityModel) assert retrieved.name == "TestEntity" - assert retrieved.entity_type == "test" + assert retrieved.entity_type == EntityType.KNOWLEDGE assert retrieved.description == "Test description" async def test_update_entity_description(entity_service: EntityService): """Test updating an entity's description.""" - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="Initial description", observations=[], ) @@ -197,9 +196,9 @@ async def test_update_entity_description(entity_service: EntityService): async def test_update_entity_description_to_none(entity_service: EntityService): """Test updating an entity's description to None.""" - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description="Initial description", observations=[], ) @@ -216,9 +215,9 @@ async def test_update_entity_description_to_none(entity_service: EntityService): async def test_delete_entity_success(entity_service: EntityService): """Test successful entity deletion.""" - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, observations=[], ) await entity_service.create_entity(entity_data) @@ -248,9 +247,9 @@ async def test_create_entity_with_special_chars(entity_service: EntityService): """Test entity creation with special characters in name and description.""" name = "TestEntity_Special" # Note: Using valid path characters description = "Description with $pecial chars & symbols!" - entity_data = Entity( + entity_data = EntitySchema( name=name, - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description=description, ) entity = await entity_service.create_entity(entity_data) @@ -266,9 +265,9 @@ async def test_create_entity_with_special_chars(entity_service: EntityService): async def test_create_entity_long_description(entity_service: EntityService): """Test creating entity with a long description.""" long_description = "A" * 1000 # 1000 character description - entity_data = Entity( + entity_data = EntitySchema( name="TestEntity", - entity_type="test", + entity_type=EntityType.KNOWLEDGE, description=long_description, observations=[], ) @@ -284,15 +283,15 @@ async def test_create_entity_long_description(entity_service: EntityService): async def test_open_nodes_by_path_ids(entity_service: EntityService): """Test opening multiple nodes by path IDs.""" # Create test entities - entity1_data = Entity( + entity1_data = EntitySchema( name="Entity1", - entity_type="type1", + entity_type=EntityType.KNOWLEDGE, description="First entity", observations=[], ) - entity2_data = Entity( + entity2_data = EntitySchema( name="Entity2", - entity_type="type2", + entity_type=EntityType.KNOWLEDGE, description="Second entity", observations=[], ) @@ -317,9 +316,9 @@ async def test_open_nodes_empty_input(entity_service: EntityService): async def test_open_nodes_some_not_found(entity_service: EntityService): """Test opening nodes with mix of existing and non-existent path IDs.""" # Create one test entity - entity_data = Entity( + entity_data = EntitySchema( name="Entity1", - entity_type="type1", + entity_type=EntityType.KNOWLEDGE, description="Test entity", observations=[], ) @@ -336,15 +335,15 @@ async def test_open_nodes_some_not_found(entity_service: EntityService): async def test_delete_entities_by_path_ids(entity_service: EntityService): """Test deleting multiple entities by path IDs.""" # Create test entities - entity1_data = Entity( + entity1_data = EntitySchema( name="Entity1", - entity_type="type1", + entity_type=EntityType.KNOWLEDGE, description="First entity", observations=[], ) - entity2_data = Entity( + entity2_data = EntitySchema( name="Entity2", - entity_type="type2", + entity_type=EntityType.KNOWLEDGE, description="Second entity", observations=[], )