admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

update entity/repostitory/service

Browse Source
This commit is contained in:
phernandez
2025-01-06 18:57:04 -06:00
parent fb330fe745
commit ea359d2484
10 changed files with 115 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/api/routers/knowledge_router.py
-13
View File
@@ -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:
src/basic_memory/models/documents.py
-5
View File
@@ -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}')"
src/basic_memory/models/knowledge.py
+41 -23
View File
@@ -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}')"
return f"Relation(id={self.id}, from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')"
src/basic_memory/repository/entity_repository.py
-40
View File
@@ -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."""
src/basic_memory/schemas/base.py
+10 -11
View File
@@ -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] = []
src/basic_memory/schemas/response.py
+1
View File
@@ -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] = []
src/basic_memory/services/entity_service.py
+1 -5
View File
@@ -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}")
tests/conftest.py
+2 -2
View File
@@ -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",
tests/repository/test_entity_repository.py
+17 -110
View File
@@ -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
tests/services/test_entity_service.py
+43 -44
View File
@@ -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=[],
)