refactor schemas

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"])

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]

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)
 
 

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]
+    deleted: List[Observation]

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:

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}")

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)

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