diff --git a/src/basic_memory/cli/migrate.py b/src/basic_memory/cli/migrate.py deleted file mode 100644 index 0f7279f3..00000000 --- a/src/basic_memory/cli/migrate.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Migration tools for basic-memory""" -import json -import asyncio -from typing import List -from pathlib import Path -import typer -from basic_memory.mcp.server import get_project_services -from basic_memory.config import ProjectConfig -from basic_memory.schemas import EntityIn, ObservationIn, RelationIn - -app = typer.Typer() - -class Migrator: - def __init__(self, memory_service): - self.memory_service = memory_service - self.entity_map = {} # name -> id mapping - - async def migrate_entities(self, entity_data_list: List[dict]): - """Migrate all entities at once""" - entities_in = [ - EntityIn( - name=entity['name'], - entity_type=entity['entityType'], - observations=[ - ObservationIn(content=obs) - for obs in entity['observations'] - ] - ) - for entity in entity_data_list - ] - - entities = await self.memory_service.create_entities(entities_in) - - # Track IDs for relations - for entity in entities: - self.entity_map[entity.name] = entity.id - - return entities - - async def migrate_relations(self, relation_data_list: List[dict]): - """Create all relations at once""" - relations_in = [] - - for rel in relation_data_list: - # Extract source/target IDs handling different formats - from_id = rel.get('from') or rel.get('from_id') - to_id = rel.get('to') or rel.get('to_id') - relation_type = rel.get('relationType') or rel.get('relation_type') - - if from_id and to_id and relation_type: - if from_id in self.entity_map and to_id in self.entity_map: - relations_in.append( - RelationIn( - fromId=self.entity_map[from_id], - toId=self.entity_map[to_id], - relationType=relation_type - ) - ) - else: - typer.echo(f"Skipping relation - missing entities: {from_id} -> {to_id}") - - return await self.memory_service.create_relations(relations_in) - -@app.command() -def migrate_json( - json_path: Path = typer.Argument(..., help="Path to JSON memory store file"), - project_path: Path = typer.Argument(..., help="Path to basic-memory project"), -): - """Migrate data from JSON memory store to basic-memory""" - async def run_migration(): - config = ProjectConfig(path=project_path) - async with get_project_services(config.path) as service: - migrator = Migrator(service) - - # Load JSONL data line by line - typer.echo(f"Loading data from {json_path}") - entities = [] - relations = [] - with open(json_path) as f: - for line in f: - if line.strip(): # Skip empty lines - item = json.loads(line) - if item['type'] == 'entity': - entities.append(item) - elif item['type'] == 'relation': - relations.append(item) - - # Create all entities in parallel - typer.echo(f"Migrating {len(entities)} entities...") - result = await migrator.migrate_entities(entities) - typer.echo(f"Migrated {len(result)} entities") - - # Create all relations in parallel - typer.echo(f"Migrating {len(relations)} relations...") - result = await migrator.migrate_relations(relations) - typer.echo(f"Migrated {len(result)} relations") - - asyncio.run(run_migration()) - -if __name__ == "__main__": - app() \ No newline at end of file diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 7c7e521a..ef49c93f 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -1,8 +1,4 @@ -""" -Core pydantic models for basic-memory entities, observations, and relations. -These models define the schema for our core data types while remaining -independent from storage/persistence concerns. -""" +"""Core pydantic models for basic-memory entities, observations, and relations.""" from typing import List, Optional, Dict, Any, Annotated from annotated_types import Len from pydantic import BaseModel, ConfigDict @@ -13,11 +9,6 @@ class SQLAlchemyOut(BaseModel): model_config = ConfigDict(from_attributes=True) # Base Models -# TODO remove -class ObservationIn(BaseModel): - """Schema for creating a single observation.""" - content: str - class ObservationsIn(BaseModel): """Schema for adding observations to an entity.""" entity_id: str @@ -25,9 +16,10 @@ class ObservationsIn(BaseModel): observations: List[str] model_config = ConfigDict(populate_by_name=True) -class ObservationOut(ObservationIn, SQLAlchemyOut): +class ObservationOut(SQLAlchemyOut): """Schema for observation data returned from the service.""" id: int + content: str class ObservationsOut(SQLAlchemyOut): """Schema for bulk observation operation results.""" @@ -97,11 +89,9 @@ class OpenNodesInput(BaseModel): """Input schema for open_nodes tool.""" names: Annotated[List[str], Len(min_length=1)] -class AddObservationsInput(BaseModel): +class AddObservationsInput(ObservationsIn): """Input schema for add_observations tool.""" - entity_id: str - observations: List[ObservationIn] - model_config = ConfigDict(populate_by_name=True) + pass class CreateRelationsInput(BaseModel): """Input schema for create_relations tool.""" diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index b21180e0..61483208 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -5,7 +5,7 @@ from pathlib import Path from basic_memory.models import Entity, Observation, Relation from basic_memory.schemas import ( - ObservationsIn, EntityIn, RelationIn, ObservationIn + ObservationsIn, EntityIn, RelationIn ) from basic_memory.fileio import write_entity_file, read_entity_file, EntityNotFoundError from basic_memory.services import EntityService, RelationService, ObservationService @@ -166,8 +166,7 @@ class MemoryService: logger.debug(f"Read entity from filesystem: {db_entity.id}") # Create new observations for the entity - for obs in observations_in.observations: - entity.observations.append(ObservationIn(content=obs)) + entity.observations += observations_in.observations logger.debug(f"Added {len(observations_in.observations)} observations to entity") # Write updated entity file diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index e0d5fba0..9e9cfcaa 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -90,12 +90,14 @@ async def test_create_directory_entity(test_directory_entity_data, memory_servic assert result[0].type == "resource" # Verify entity creation - response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) + response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert len(response.entities) == 1 assert response.entities[0].name == "Directory Organization" assert response.entities[0].entity_type == "memory" assert len(response.entities[0].observations) == 3 + +# noinspection DuplicatedCode @pytest.mark.anyio async def test_create_entities_snake_case(test_entity_snake_case, memory_service, test_config): """Test creating an entity with snake_case data (like internal usage).""" @@ -113,7 +115,7 @@ async def test_create_entities_snake_case(test_entity_snake_case, memory_service assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) assert result[0].resource.mimeType == MIME_TYPE - response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) + response = CreateEntitiesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert len(response.entities) == 1 assert response.entities[0].name == "Test Entity" assert response.entities[0].entity_type == "test" @@ -145,7 +147,7 @@ async def test_search_nodes(test_entity_data, memory_service, test_config): assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) assert result[0].resource.mimeType == MIME_TYPE - response = SearchNodesResponse.model_validate_json(result[0].resource.text) + response = SearchNodesResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert len(response.matches) == 1 assert response.matches[0].name == "Test Entity" assert response.query == "Test Entity" @@ -162,7 +164,7 @@ async def test_add_observations(test_entity_data, memory_service, test_config): memory_service=memory_service ) - create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text) + create_response = CreateEntitiesResponse.model_validate_json(create_result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] entity_id = create_response.entities[0].id # Add new observations using camelCase @@ -182,7 +184,7 @@ async def test_add_observations(test_entity_data, memory_service, test_config): assert str(result[0].resource.uri) == str(BASIC_MEMORY_URI) assert result[0].resource.mimeType == MIME_TYPE - response = AddObservationsResponse.model_validate_json(result[0].resource.text) + response = AddObservationsResponse.model_validate_json(result[0].resource.text) # pyright: ignore [reportAttributeAccessIssue] assert response.entity_id == entity_id assert len(response.added_observations) == 1 assert response.added_observations[0].content == "A new observation" diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 9bd3b3f6..f9f70fd9 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -4,7 +4,6 @@ from pydantic import ValidationError from basic_memory.schemas import ( EntityIn, EntityOut, - ObservationIn, RelationIn, CreateEntitiesInput, SearchNodesInput, @@ -61,19 +60,6 @@ def test_entity_in_validation(): with pytest.raises(ValidationError): EntityIn.model_validate({"entityType": "test"}) # Missing name -def test_observation_in_validation(): - """Test ObservationIn validation.""" - # Minimal - obs = ObservationIn.model_validate({"content": "test"}) - assert obs.content == "test" - - # With context - obs = ObservationIn.model_validate({"content": "test", "context": "test context"}) - - # Missing content - with pytest.raises(ValidationError): - ObservationIn.model_validate({}) - def test_relation_in_validation(): """Test RelationIn validation.""" data = {