diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index dc6d62da..e19f92ab 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -7,10 +7,10 @@ from basic_memory.schemas import ( CreateEntityRequest, CreateEntityResponse, SearchNodesRequest, SearchNodesResponse, CreateRelationsRequest, CreateRelationsResponse, - EntityResponse, RelationResponse, AddObservationsRequest, ObservationResponse, + EntityResponse, AddObservationsRequest, ObservationResponse, OpenNodesRequest, OpenNodesResponse, DeleteEntityResponse, - DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse + DeleteObservationsRequest, DeleteObservationsResponse, AddObservationsResponse, RelationResponse ) router = APIRouter(prefix="/knowledge", tags=["knowledge"]) diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index 8c26c0cb..dd3371c5 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -6,7 +6,7 @@ from pathlib import Path from loguru import logger -from basic_memory.schemas import EntityRequest, RelationRequest +from basic_memory.schemas import Entity, Relation class FileOperationError(Exception): @@ -34,7 +34,7 @@ def get_entity_path(project_entities_path: Path, entity_id: str) -> Path: return Path(f"{project_entities_path}/{entity_id}.md") -async def write_entity_file(project_entities_path: Path, entity_id: str, entity: EntityRequest) -> bool: +async def write_entity_file(project_entities_path: Path, entity_id: str, entity: Entity) -> bool: """ Write entity to filesystem in markdown format. @@ -102,7 +102,7 @@ async def write_entity_file(project_entities_path: Path, entity_id: str, entity: return True -async def read_entity_file(project_entities_path: Path, entity_id: str) -> EntityRequest: +async def read_entity_file(project_entities_path: Path, entity_id: str) -> Entity: """ Read entity data from filesystem. @@ -174,14 +174,14 @@ async def read_entity_file(project_entities_path: Path, entity_id: str) -> Entit relation_type = parts[0] context = parts[1] if len(parts) > 1 else None - relations.append(RelationRequest( # pyright: ignore [reportCallIssue] + relations.append(Relation( # pyright: ignore [reportCallIssue] from_id=entity_id, # pyright: ignore [reportCallIssue] to_id=target_id, # pyright: ignore [reportCallIssue] relation_type=relation_type, # pyright: ignore [reportCallIssue] context=context )) - return EntityRequest( # pyright: ignore [reportCallIssue] + return Entity( # pyright: ignore [reportCallIssue] id=entity_id, name=name, entity_type=entity_type, # pyright: ignore [reportCallIssue] diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index 98a5495b..b3fab06e 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -27,7 +27,7 @@ from basic_memory.schemas import ( # Tool responses CreateEntityResponse, SearchNodesResponse, OpenNodesResponse, AddObservationsResponse, CreateRelationsResponse, DeleteEntityResponse, - EntityResponse, ObservationResponse, RelationResponse, AddObservationsRequest + EntityResponse, ObservationResponse, Relation, AddObservationsRequest ) from basic_memory.services import EntityService, ObservationService, RelationService from basic_memory.services.memory_service import MemoryService @@ -182,7 +182,7 @@ async def handle_create_relations( logger.debug(f"Created {len(created)} relations") # Format response - response = CreateRelationsResponse(relations=[RelationResponse.model_validate(relation) for relation in created]) + response = CreateRelationsResponse(relations=[Relation.model_validate(relation) for relation in created]) return create_response(response) diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index e32c4d94..ae3a90c9 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -1,31 +1,12 @@ """Core pydantic models for basic-memory entities, observations, and relations.""" -from typing import List, Optional, Annotated +from typing import List, Optional, Annotated, TypeAlias from annotated_types import Len from pydantic import BaseModel, ConfigDict -# Base output model for SQLAlchemy attribute conversion -class SQLAlchemyModel(BaseModel): - """Base class for models that read from SQLAlchemy attributes.""" - model_config = ConfigDict(from_attributes=True) - # Base Models -class AddObservationsRequest(BaseModel): - """Schema for adding observations to an entity.""" - entity_id: str - context: Optional[str] = None - observations: List[str] +Observation: TypeAlias = str -class ObservationResponse(SQLAlchemyModel): - """Schema for observation data returned from the service.""" - id: int - content: str - -class ObservationsResponse(SQLAlchemyModel): - """Schema for bulk observation operation results.""" - entity_id: str - observations: List[ObservationResponse] - -class RelationRequest(BaseModel): +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.) @@ -35,44 +16,35 @@ class RelationRequest(BaseModel): relation_type: str context: Optional[str] = None - -class RelationResponse(SQLAlchemyModel): - id: int - from_id: str - to_id: str - relation_type: str - context: Optional[str] = None - -class EntityBase(BaseModel): +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: Optional[str] = None name: str entity_type: str description: Optional[str] = None + observations: List[Observation] = [] + relations: List[Relation] = [] @property def file_path(self) -> str: """The relative file path for this entity.""" return f"{id}.md" +# request input models -class EntityRequest(EntityBase): - """ - 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. - """ - observations: List[str] = [] - relations: List[RelationRequest] = [] +class AddObservationsRequest(BaseModel): + """Schema for adding observations to an entity.""" + entity_id: str + context: Optional[str] = None + observations: List[Observation] -class EntityResponse(EntityBase, SQLAlchemyModel): - """Schema for entity data returned from the service.""" - observations: List[ObservationResponse] = [] - relations: List[RelationResponse] = [] - -# Tool Request schemas class CreateEntityRequest(BaseModel): """Request schema for create_entities tool.""" - entities: Annotated[List[EntityRequest], Len(min_length=1)] + entities: Annotated[List[Entity], Len(min_length=1)] class SearchNodesRequest(BaseModel): """Request schema for search_nodes tool.""" @@ -84,7 +56,7 @@ class OpenNodesRequest(BaseModel): class CreateRelationsRequest(BaseModel): """Request schema for create_relations tool.""" - relations: List[RelationRequest] + relations: List[Relation] class DeleteEntityRequest(BaseModel): """Request schema for delete_entities tool.""" @@ -95,6 +67,35 @@ class DeleteObservationsRequest(BaseModel): entity_id: str deletions: List[str] # TODO: Make this more specific +# response output models + +# Base output model for SQLAlchemy attribute conversion +class SQLAlchemyModel(BaseModel): + """Base class for models that read from SQLAlchemy attributes.""" + model_config = ConfigDict(from_attributes=True) + +class ObservationResponse(SQLAlchemyModel): + """Schema for observation data returned from the service.""" + id: int + content: Observation + +class ObservationsResponse(SQLAlchemyModel): + """Schema for bulk observation operation results.""" + entity_id: str + observations: List[ObservationResponse] + +class RelationResponse(Relation,SQLAlchemyModel): + id: int + +class EntityResponse(SQLAlchemyModel): + """Schema for entity data returned from the service.""" + id: str + name: str + entity_type: str + description: Optional[str] = None + observations: List[ObservationResponse] = [] + relations: List[RelationResponse] = [] + class CreateEntityResponse(SQLAlchemyModel): """Response for create_entities tool.""" entities: List[EntityResponse] @@ -115,7 +116,7 @@ class AddObservationsResponse(SQLAlchemyModel): class CreateRelationsResponse(SQLAlchemyModel): """Response for create_relations tool.""" - relations: List[RelationResponse] + relations: List[Relation] class DeleteEntityResponse(SQLAlchemyModel): """Response for delete_entities tool.""" @@ -124,4 +125,4 @@ class DeleteEntityResponse(SQLAlchemyModel): class DeleteObservationsResponse(SQLAlchemyModel): """Response for delete_observations tool.""" entity_id: str - deleted: List[str] \ No newline at end of file + deleted: List[Observation] \ No newline at end of file diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index e57e5a49..e0c04dbd 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -3,8 +3,8 @@ from pathlib import Path from typing import Dict, Any, Sequence from basic_memory.repository.entity_repository import EntityRepository -from basic_memory.schemas import EntityRequest -from basic_memory.models import Entity +from basic_memory.schemas import Entity +from basic_memory.models import Entity as EntityModel from basic_memory.fileio import EntityNotFoundError from loguru import logger @@ -20,7 +20,7 @@ class EntityService: self.entity_repo = entity_repo logger.debug(f"Initialized EntityService with path: {project_path}") - async def search(self, query: str) -> Sequence[Entity]: + async def search(self, query: str) -> Sequence[EntityModel]: """Search entities using LIKE pattern matching.""" logger.debug(f"Searching entities with query: {query}") try: @@ -31,7 +31,7 @@ class EntityService: logger.exception(f"Failed to search entities with query: {query}") raise - async def create_entity(self, entity: EntityRequest) -> Entity: + async def create_entity(self, entity: Entity) -> EntityModel: """Create a new entity in the database.""" logger.debug(f"Creating entity in DB: {entity}") try: @@ -46,7 +46,7 @@ class EntityService: logger.exception(f"Failed to create entity: {entity}") raise - async def update_entity(self, entity_id: str, update_data: Dict[str, Any]) -> Entity: + async def update_entity(self, entity_id: str, update_data: Dict[str, Any]) -> EntityModel: """Update an entity's fields.""" logger.debug(f"Updating entity {entity_id} with data: {update_data}") try: @@ -62,7 +62,7 @@ class EntityService: logger.exception(f"Failed to update entity: {entity_id}") raise - async def get_entity(self, entity_id: str) -> Entity: + async def get_entity(self, entity_id: str) -> EntityModel: """Get entity by ID.""" logger.debug(f"Getting entity by ID: {entity_id}") try: @@ -79,7 +79,7 @@ class EntityService: logger.exception(f"Failed to get entity: {entity_id}") raise - async def get_by_type_and_name(self, entity_type: str, name: str) -> Entity: + async def get_by_type_and_name(self, entity_type: str, name: str) -> EntityModel: """Get entity by type and name combination.""" logger.debug(f"Getting entity by type/name: {entity_type}/{name}") try: @@ -96,7 +96,7 @@ class EntityService: logger.exception(f"Failed to get entity by type/name: {entity_type}/{name}") raise - async def get_all(self) -> Sequence[Entity]: + async def get_all(self) -> Sequence[EntityModel]: return await self.entity_repo.find_all() async def delete_entity(self, entity_id: str) -> bool: diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 762c3ab5..43410464 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -3,9 +3,9 @@ import asyncio from typing import List, Dict, Any, Optional, Sequence from pathlib import Path -from basic_memory.models import Entity, Observation +from basic_memory.models import Entity as EntityModel, Observation as ObservationModel, Relation as RelationModel from basic_memory.schemas import ( - AddObservationsRequest, EntityRequest, RelationRequest + AddObservationsRequest, Entity, Relation ) from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError from basic_memory.services import EntityService, RelationService, ObservationService @@ -32,12 +32,12 @@ class MemoryService: self.observation_service = observation_service logger.debug(f"Initialized MemoryService with path: {project_path}") - async def create_entities(self, entities_in: List[EntityRequest]) -> List[Entity]: + async def create_entities(self, entities_in: List[Entity]) -> List[EntityModel]: """Create multiple entities with their observations.""" logger.debug(f"Creating {len(entities_in)} entities") # Write files in parallel (filesystem is source of truth) - async def write_file(entity: EntityRequest): + async def write_file(entity: Entity): try: existing = await self.entity_service.get_by_type_and_name( entity.entity_type, @@ -52,7 +52,7 @@ class MemoryService: pass # Generate ID and write file - entity_id = Entity.generate_id(entity.entity_type, entity.name) + entity_id = EntityModel.generate_id(entity.entity_type, entity.name) await write_entity_file(self.entities_path, entity_id, entity) file_writes = [write_file(entity) for entity in entities_in] @@ -60,7 +60,7 @@ class MemoryService: await asyncio.gather(*file_writes) logger.debug("Completed all file writes") - async def create_entity_in_db(entity_in: EntityRequest): + async def create_entity_in_db(entity_in: Entity): logger.debug(f"Creating entity in DB: {entity_in}") try: # Create base entity @@ -98,7 +98,7 @@ class MemoryService: logger.exception("Failed to create entities in DB") for entity in entities_in: try: - entity_id = Entity.generate_id(entity.entity_type, entity.name) + entity_id = EntityModel.generate_id(entity.entity_type, entity.name) path = self.entities_path / entity_id if path.exists(): path.unlink() @@ -112,7 +112,7 @@ class MemoryService: logger.debug(f"Found entity {entity}") return entity - async def create_relations(self, relations_data: List[RelationRequest]) -> List[RelationRequest]: + async def create_relations(self, relations_data: List[Relation]) -> List[RelationModel]: """Create multiple relations between entities.""" logger.debug(f"Creating {len(relations_data)} relations") @@ -153,7 +153,7 @@ class MemoryService: logger.debug(f"Successfully created {len(relations)} relations") return relations - async def add_observations(self, observations_in: AddObservationsRequest) -> List[Observation]: + async def add_observations(self, observations_in: AddObservationsRequest) -> List[ObservationModel]: """Add observations to an existing entity.""" logger.debug(f"Adding observations to entity: {observations_in.entity_id}") try: @@ -192,7 +192,7 @@ class MemoryService: async def delete_relations(self, relations: List[Dict[str, Any]]) -> None: pass - async def read_graph(self) -> Sequence[Entity]: + async def read_graph(self) -> Sequence[EntityModel]: """Read the entire knowledge graph.""" logger.debug("Reading entire knowledge graph") try: @@ -203,7 +203,7 @@ class MemoryService: logger.exception("Failed to read graph") raise - async def search_nodes(self, query: str) -> Sequence[Entity]: + async def search_nodes(self, query: str) -> Sequence[EntityModel]: """Search for nodes in the knowledge graph.""" logger.debug(f"Searching nodes with query: {query}") try: @@ -214,11 +214,11 @@ class MemoryService: logger.exception(f"Failed to search nodes with query: {query}") raise - async def open_nodes(self, names: List[str]) -> List[EntityRequest]: + async def open_nodes(self, names: List[str]) -> List[Entity]: """Get specific nodes and their relationships.""" logger.debug(f"Opening nodes: {names}") - async def read_node(name: str) -> Optional[EntityRequest]: + async def read_node(name: str) -> Optional[Entity]: try: # Get ID from name first logger.debug(f"Looking up entity: {name}") diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index c79c1c46..f868b778 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -3,11 +3,14 @@ from pathlib import Path from typing import List, Sequence from sqlalchemy import select -from basic_memory.models import Observation +from basic_memory.models import Observation as ObservationModel from basic_memory.repository.observation_repository import ObservationRepository from . import DatabaseSyncError +from basic_memory.schemas import Observation +#from basic_memory.schemas import Observation + class ObservationService: """ Service for managing observations in the database. @@ -18,14 +21,14 @@ class ObservationService: self.project_path = project_path self.observation_repo = observation_repo - async def add_observations(self, entity_id: str, observations: List[str]) -> List[Observation]: + async def add_observations(self, entity_id: str, observations: List[Observation]) -> List[ObservationModel]: """ Add multiple observations to an entity. Returns the created observations with IDs set. """ try: return await self.observation_repo.bulk_create([ - Observation( + ObservationModel( entity_id=entity_id, content=observation, ) @@ -34,7 +37,7 @@ class ObservationService: except Exception as e: raise DatabaseSyncError(f"Failed to add observations to database: {str(e)}") from e - async def search_observations(self, query: str) -> List[Observation]: + async def search_observations(self, query: str) -> List[ObservationModel]: """ Search for observations across all entities. @@ -45,15 +48,15 @@ class ObservationService: List of matching observations with their entity contexts """ result = await self.observation_repo.execute_query( - select(Observation).filter( - Observation.content.contains(query) + select(ObservationModel).filter( + ObservationModel.content.contains(query) ) ) return [ - Observation(content=obs.content) + ObservationModel(content=obs.content) for obs in result.scalars().all() ] - async def get_observations_by_context(self, context: str) -> Sequence[Observation]: + async def get_observations_by_context(self, context: str) -> Sequence[ObservationModel]: """Get all observations with a specific context.""" return await self.observation_repo.find_by_context(context) diff --git a/src/basic_memory/services/relation_service.py b/src/basic_memory/services/relation_service.py index 9d50b700..f9f61216 100644 --- a/src/basic_memory/services/relation_service.py +++ b/src/basic_memory/services/relation_service.py @@ -2,9 +2,9 @@ from pathlib import Path from basic_memory.repository.relation_repository import RelationRepository -from basic_memory.schemas import EntityRequest, RelationRequest +from basic_memory.schemas import Entity, Relation from . import DatabaseSyncError - +from basic_memory.models import Relation as RelationModel class RelationService: """ @@ -16,14 +16,14 @@ class RelationService: self.project_path = project_path self.relation_repo = relation_repo - async def create_relation(self, relation: RelationRequest) -> RelationRequest: + async def create_relation(self, relation: Relation) -> RelationModel: """Create a new relation in the database.""" try: return await self.relation_repo.create(relation.model_dump()) except Exception as e: raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e - async def delete_relation(self, from_entity: EntityRequest, to_entity: EntityRequest, relation_type: str) -> bool: + async def delete_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str) -> bool: """Delete a specific relation between entities.""" try: # Use repository to find and delete the relation diff --git a/tests/api/test_knowledge.py b/tests/api/test_knowledge.py index 7bb5ceed..183a85b3 100644 --- a/tests/api/test_knowledge.py +++ b/tests/api/test_knowledge.py @@ -105,6 +105,8 @@ async def test_create_relations(client: AsyncClient): }] }) + logger.debug("response: %s", response.json()) + # Verify relation assert response.status_code == 200 data = response.json() diff --git a/tests/conftest.py b/tests/conftest.py index 4616fa85..e78e41d9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,7 +14,7 @@ from basic_memory.deps import ( get_relation_service, get_relation_repo, get_observation_repo, get_entity_repo ) -from basic_memory.schemas import EntityRequest +from basic_memory.schemas import Entity from basic_memory.config import ProjectConfig from basic_memory.services import MemoryService @@ -115,7 +115,7 @@ async def sample_entity(entity_repository: EntityRepository): @pytest_asyncio.fixture async def test_entity(entity_service): """Create a test entity for reuse in tests.""" - entity_data = EntityRequest( # pyright: ignore [reportCallIssue] + entity_data = Entity( # pyright: ignore [reportCallIssue] name="Test Entity", entity_type="test", # pyright: ignore [reportCallIssue] ) diff --git a/tests/test_entity_service.py b/tests/test_entity_service.py index bececc89..9660fe75 100644 --- a/tests/test_entity_service.py +++ b/tests/test_entity_service.py @@ -2,15 +2,15 @@ import pytest from basic_memory.fileio import EntityNotFoundError -from basic_memory.models import Entity -from basic_memory.schemas import EntityRequest +from basic_memory.models import Entity as EntityModel +from basic_memory.schemas import Entity pytestmark = pytest.mark.asyncio async def test_create_entity_success(entity_service): """Test successful entity creation.""" - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="A test entity description" @@ -20,7 +20,7 @@ async def test_create_entity_success(entity_service): entity = await entity_service.create_entity(entity_data) # Assert Entity - assert isinstance(entity, Entity) + assert isinstance(entity, EntityModel) assert entity.name == "Test Entity" assert entity.entity_type == "test" assert entity.description == "A test entity description" @@ -33,14 +33,14 @@ async def test_create_entity_success(entity_service): async def test_get_by_type_and_name(entity_service): """Test finding entity by type and name combination.""" # Create two entities with same name but different types - entity1_data = EntityRequest( + entity1_data = Entity( name="Test Entity", entity_type="type1", description="First test entity" ) entity1 = await entity_service.create_entity(entity1_data) - entity2_data = EntityRequest( + entity2_data = Entity( name="Test Entity", # Same name entity_type="type2", # Different type description="Second test entity" @@ -67,7 +67,7 @@ async def test_get_by_type_and_name(entity_service): async def test_create_entity_no_description(entity_service): """Test creating entity without description (should be None).""" - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", ) @@ -82,7 +82,7 @@ async def test_create_entity_no_description(entity_service): async def test_get_entity_success(entity_service): """Test successful entity retrieval.""" # Arrange - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="Test description" @@ -93,7 +93,7 @@ async def test_get_entity_success(entity_service): retrieved = await entity_service.get_entity(created.id) # Assert - assert isinstance(retrieved, Entity) + assert isinstance(retrieved, EntityModel) assert retrieved.id == created.id assert retrieved.name == created.name assert retrieved.entity_type == created.entity_type @@ -103,7 +103,7 @@ async def test_get_entity_success(entity_service): async def test_update_entity_description(entity_service): """Test updating an entity's description.""" # Create entity with description - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="Initial description" @@ -121,7 +121,7 @@ async def test_update_entity_description(entity_service): async def test_update_entity_description_to_none(entity_service): """Test updating an entity's description to None.""" # Create entity with description - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="Initial description" @@ -139,7 +139,7 @@ async def test_update_entity_description_to_none(entity_service): async def test_delete_entity_success(entity_service): """Test successful entity deletion.""" # Arrange - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", ) @@ -167,7 +167,7 @@ async def test_create_entity_db_error(entity_service, monkeypatch): raise Exception("Mock DB error") monkeypatch.setattr(entity_service.entity_repo, "create", mock_create) - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="Test description" @@ -188,7 +188,7 @@ async def test_create_entity_with_special_chars(entity_service): """Test entity creation with special characters in name and description.""" name = "Test & Entity! With @ Special #Chars" description = "Description with $pecial chars & symbols!" - entity_data = EntityRequest( + entity_data = Entity( name=name, entity_type="test", description=description @@ -204,7 +204,7 @@ async def test_create_entity_with_special_chars(entity_service): async def test_entity_id_generation(entity_service): """Test that entities get unique IDs generated correctly.""" - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description="Test description", @@ -219,7 +219,7 @@ async def test_entity_id_generation(entity_service): async def test_create_entity_long_description(entity_service): """Test creating entity with a long description.""" long_description = "A" * 1000 # 1000 character description - entity_data = EntityRequest( + entity_data = Entity( name="Test Entity", entity_type="test", description=long_description diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index ed3baeb2..06442abd 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -2,7 +2,7 @@ import pytest from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file -from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest +from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, Relation test_entities_data = [ { @@ -195,4 +195,4 @@ async def test_create_relations_with_invalid_entity_id(memory_service: MemorySer } with pytest.raises(Exception) as exc: # We might want to define a specific error type - await memory_service.create_relations([RelationRequest.model_validate(bad_relation)]) \ No newline at end of file + await memory_service.create_relations([Relation.model_validate(bad_relation)]) \ No newline at end of file diff --git a/tests/test_relation_service.py b/tests/test_relation_service.py index f76f05d8..04b9ed65 100644 --- a/tests/test_relation_service.py +++ b/tests/test_relation_service.py @@ -2,7 +2,7 @@ import pytest import pytest_asyncio -from basic_memory.schemas import EntityRequest, RelationRequest +from basic_memory.schemas import Entity, Relation pytestmark = pytest.mark.asyncio @@ -10,13 +10,13 @@ pytestmark = pytest.mark.asyncio @pytest_asyncio.fixture async def sample_entities(entity_service): """Create two sample entities for testing relations""" - entity1_data = EntityRequest( + entity1_data = Entity( name="test_entity_1", entity_type="test_type", observations=[], relations=[] ) - entity2_data = EntityRequest( + entity2_data = Entity( name="test_entity_2", entity_type="test_type", observations=[], @@ -36,7 +36,7 @@ async def test_create_relation(relation_service, sample_entities): """Test creating a basic relation between two entities""" entity1, entity2 = sample_entities - relation_data = RelationRequest( + relation_data = Relation( from_id=entity1.id, to_id=entity2.id, relation_type="test_relation" @@ -62,7 +62,7 @@ async def test_create_relation_with_context(relation_service, sample_entities): """Test creating a relation with context information""" entity1, entity2 = sample_entities - relation_data = RelationRequest( + relation_data = Relation( from_id=entity1.id, to_id=entity2.id, relation_type="test_relation", diff --git a/tests/test_schemas.py b/tests/test_schemas.py index fa5f72ea..f52b51e8 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -2,9 +2,9 @@ import pytest from pydantic import ValidationError from basic_memory.schemas import ( - EntityRequest, + Entity, EntityResponse, - RelationRequest, + Relation, CreateEntityRequest, SearchNodesRequest, OpenNodesRequest, @@ -16,7 +16,7 @@ def test_entity_in_minimal(): "name": "test_entity", "entity_type": "test" } - entity = EntityRequest.model_validate(data) + entity = Entity.model_validate(data) assert entity.name == "test_entity" assert entity.entity_type == "test" assert entity.description is None @@ -40,7 +40,7 @@ def test_entity_in_complete(): } ] } - entity = EntityRequest.model_validate(data) + entity = Entity.model_validate(data) assert entity.name == "test_entity" assert entity.entity_type == "test" assert entity.description == "A test entity" @@ -52,13 +52,13 @@ def test_entity_in_complete(): def test_entity_in_validation(): """Test validation errors for EntityIn.""" with pytest.raises(ValidationError): - EntityRequest.model_validate({}) # Missing required fields + Entity.model_validate({}) # Missing required fields with pytest.raises(ValidationError): - EntityRequest.model_validate({"name": "test"}) # Missing entityType + Entity.model_validate({"name": "test"}) # Missing entityType with pytest.raises(ValidationError): - EntityRequest.model_validate({"entityType": "test"}) # Missing name + Entity.model_validate({"entityType": "test"}) # Missing name def test_relation_in_validation(): """Test RelationIn validation.""" @@ -67,7 +67,7 @@ def test_relation_in_validation(): "to_id": "456", "relation_type": "test" } - relation = RelationRequest.model_validate(data) + relation = Relation.model_validate(data) assert relation.from_id == "123" assert relation.to_id == "456" assert relation.relation_type == "test" @@ -75,12 +75,12 @@ def test_relation_in_validation(): # With context data["context"] = "test context" - relation = RelationRequest.model_validate(data) + relation = Relation.model_validate(data) assert relation.context == "test context" # Missing required fields with pytest.raises(ValidationError): - RelationRequest.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType + Relation.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType def test_create_entities_input(): """Test CreateEntitiesInput validation.""" @@ -132,18 +132,17 @@ def test_entity_out_from_attributes(): assert len(entity.observations) == 1 assert entity.observations[0].id == 1 assert len(entity.relations) == 1 - assert entity.relations[0].id == 1 def test_optional_fields(): """Test handling of optional fields.""" # Create with no optional fields - entity = EntityRequest.model_validate({"name": "test", "entity_type": "test"}) + entity = Entity.model_validate({"name": "test", "entity_type": "test"}) assert entity.description is None assert entity.observations == [] assert entity.relations == [] # Create with empty optional fields - entity = EntityRequest.model_validate({ + entity = Entity.model_validate({ "name": "test", "entity_type": "test", "description": None, @@ -155,7 +154,7 @@ def test_optional_fields(): assert entity.relations == [] # Create with some optional fields - entity = EntityRequest.model_validate({ + entity = Entity.model_validate({ "name": "test", "entity_type": "test", "description": "test",