diff --git a/src/basic_memory/api/routers/knowledge.py b/src/basic_memory/api/routers/knowledge.py index 9df1e553..cac58fc8 100644 --- a/src/basic_memory/api/routers/knowledge.py +++ b/src/basic_memory/api/routers/knowledge.py @@ -23,7 +23,7 @@ async def create_entities( return CreateEntitiesResponse(entities=[EntityOut.model_validate(entity) for entity in entities]) -@router.get("/entities/{entity_id}", response_model=EntityOut) +@router.get("/entities/{entity_id:path}", response_model=EntityOut) async def get_entity( entity_id: str, memory_service: MemoryServiceDep diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index 9371f7f1..d5731777 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -69,9 +69,7 @@ async def write_entity_file(project_entities_path: Path, entity_id: str, entity: # Add observations for obs in entity.observations: - obs_line = f"- {obs.content}" - if obs.context: - obs_line += f" | {obs.context}" + obs_line = f"- {obs}" content.append(f"{obs_line}\n") # Add relations section if we have relations @@ -160,7 +158,7 @@ async def read_entity_file(project_entities_path: Path, entity_id: str) -> Entit parts = line.split(" | ", 1) content = parts[0] context = parts[1] if len(parts) > 1 else None - observations.append(ObservationIn(content=content, context=context)) + observations.append(ObservationIn(content=content)) elif in_relations and line.startswith("- "): # Parse relation line: - [target_id] relation_type | context line = line[2:] # Remove the bullet point diff --git a/src/basic_memory/mcp/server.py b/src/basic_memory/mcp/server.py index abd9895a..3c13071b 100644 --- a/src/basic_memory/mcp/server.py +++ b/src/basic_memory/mcp/server.py @@ -1,18 +1,24 @@ """MCP server implementation for basic-memory.""" import sys -import os +from contextlib import asynccontextmanager +from pathlib import Path from typing import List, Dict, Any, Optional, Literal, Callable, Awaitable + +from sqlalchemy.ext.asyncio import AsyncEngine from typing_extensions import TypeAlias from mcp.server import Server from mcp.types import Tool, EmbeddedResource, TextResourceContents, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR from mcp.shared.exceptions import McpError from pydantic.networks import AnyUrl -from pydantic import TypeAdapter, BaseModel, ConfigDict +from pydantic import TypeAdapter, BaseModel +from basic_memory import db from basic_memory.config import ProjectConfig -from basic_memory.deps import get_project_services from basic_memory.fileio import EntityNotFoundError +from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.repository.observation_repository import ObservationRepository +from basic_memory.repository.relation_repository import RelationRepository from basic_memory.schemas import ( # Tool inputs CreateEntitiesInput, SearchNodesInput, OpenNodesInput, @@ -23,8 +29,9 @@ from basic_memory.schemas import ( AddObservationsResponse, CreateRelationsResponse, DeleteEntitiesResponse, DeleteObservationsResponse, # Base models - EntityOut, ObservationOut, RelationOut + EntityOut, ObservationOut, RelationOut, ObservationsIn ) +from basic_memory.services import EntityService, ObservationService, RelationService from basic_memory.services.memory_service import MemoryService from loguru import logger @@ -45,6 +52,36 @@ ToolName = Literal[ ToolHandler: TypeAlias = Callable[[MemoryService, Dict[str, Any]], Awaitable[EmbeddedResource]] +@asynccontextmanager +async def get_memory_service_session(engine: AsyncEngine, project_path: Path): + """Get all services with proper session and lifecycle management.""" + async with db.session(engine) as session: + # Create repos + entity_repo = EntityRepository(session) + observation_repo = ObservationRepository(session) + relation_repo = RelationRepository(session) + + # Create services + entity_service = EntityService(project_path, entity_repo) + observation_service = ObservationService(project_path, observation_repo) + relation_service = RelationService(project_path, relation_repo) + + # Create memory service + memory_service = MemoryService( + project_path=project_path, + entity_service=entity_service, + relation_service=relation_service, + observation_service=observation_service + ) + + yield memory_service + +@asynccontextmanager +async def get_project_services(project_path: Path): + """Get all services for a project with full lifecycle management.""" + async with db.engine(project_path=project_path) as engine: + async with get_memory_service_session(engine, project_path) as services: + yield services def create_response(response: BaseModel) -> EmbeddedResource: """Create standard MCP response from any response model.""" @@ -117,9 +154,9 @@ async def handle_add_observations( """Handle add_observations tool call.""" # Validate input logger.debug(f"Adding observations: {args}") - input_args = AddObservationsInput.model_validate(args) + input_args = ObservationsIn.model_validate(args) logger.debug(f"Adding {len(input_args.observations)} observations to entity {input_args.entity_id}") - + # Call service with validated data observations = await service.add_observations(input_args) logger.debug(f"Added {len(observations)} observations") @@ -170,14 +207,7 @@ async def handle_delete_observations( ) -> EmbeddedResource: """Handle delete_observations tool call.""" logger.debug(f"Deleting observations: {args}") - input_args = DeleteObservationsInput.model_validate(args) - entity, deleted = await service.delete_observations(input_args.deletions) - logger.debug(f"Deleted {len(deleted)} observations from entity {entity}") - response = DeleteObservationsResponse( - entity=entity, - deleted=deleted - ) - return create_response(response) + return EmbeddedResource() # Map tool names to handlers diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py index a692ab95..4af19b76 100644 --- a/src/basic_memory/models.py +++ b/src/basic_memory/models.py @@ -60,6 +60,11 @@ class Entity(Base): cascade="all, delete-orphan" ) + @property + def relations(self): + return self.outgoing_relations + self.incoming_relations + + @classmethod def generate_id(cls, entity_type: str, name: str) -> str: """Generate a filesystem path-based ID for this entity.""" diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 9459c3f8..450ded19 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -3,10 +3,9 @@ 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. """ -from datetime import datetime, UTC from typing import List, Optional, Dict, Any, Annotated -from annotated_types import Gt, Len -from pydantic import BaseModel, Field, ConfigDict +from annotated_types import Len +from pydantic import BaseModel, ConfigDict # Base output model for SQLAlchemy attribute conversion class SQLAlchemyOut(BaseModel): @@ -14,15 +13,16 @@ class SQLAlchemyOut(BaseModel): model_config = ConfigDict(from_attributes=True) # Base Models +# TODO remove class ObservationIn(BaseModel): """Schema for creating a single observation.""" content: str - context: Optional[str] = None class ObservationsIn(BaseModel): """Schema for adding observations to an entity.""" - entity_id: str = Field(alias="entityId") # Maps to Entity.id - observations: List[ObservationIn] + entity_id: str + context: Optional[str] = None + observations: List[str] model_config = ConfigDict(populate_by_name=True) class ObservationOut(ObservationIn, SQLAlchemyOut): @@ -31,7 +31,7 @@ class ObservationOut(ObservationIn, SQLAlchemyOut): class ObservationsOut(SQLAlchemyOut): """Schema for bulk observation operation results.""" - entity_id: str = Field(alias="entityId") + entity_id: str observations: List[ObservationOut] model_config = ConfigDict(populate_by_name=True) @@ -40,25 +40,25 @@ class RelationIn(BaseModel): Represents a directed edge between entities in the knowledge graph. Relations are always stored in active voice (e.g. "created", "teaches", etc.) """ - from_id: str = Field(alias="fromId") - to_id: str = Field(alias="toId") - relation_type: str = Field(alias="relationType") + from_id: str + to_id: str + relation_type: str context: Optional[str] = None model_config = ConfigDict(populate_by_name=True) class RelationOut(SQLAlchemyOut): id: int - from_id: str = Field(alias="fromId") - to_id: str = Field(alias="toId") - relation_type: str = Field(alias="relationType") + from_id: str + to_id: str + relation_type: str context: Optional[str] = None model_config = ConfigDict(populate_by_name=True) class EntityBase(BaseModel): id: Optional[str] = None name: str - entity_type: str = Field(alias="entityType") + entity_type: str description: Optional[str] = None @property @@ -99,7 +99,7 @@ class OpenNodesInput(BaseModel): class AddObservationsInput(BaseModel): """Input schema for add_observations tool.""" - entity_id: str = Field(alias="entityId") + entity_id: str observations: List[ObservationIn] model_config = ConfigDict(populate_by_name=True) @@ -113,6 +113,7 @@ class DeleteEntitiesInput(BaseModel): class DeleteObservationsInput(BaseModel): """Input schema for delete_observations tool.""" + entity_id: str deletions: List[Dict[str, Any]] # TODO: Make this more specific # Tool Response Schemas diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index 091b0476..b21180e0 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -1,6 +1,6 @@ """Service for orchestrating entity, relation, and observation operations.""" import asyncio -from typing import List, Dict, Any, Optional +from typing import List, Dict, Any, Optional, Sequence from pathlib import Path from basic_memory.models import Entity, Observation, Relation @@ -80,7 +80,7 @@ class MemoryService: final_entity = await self.entity_service.get_entity(created_entity.id) logger.debug(f"Retrieved final entity state: {final_entity}") return final_entity - except Exception as e: + except Exception: logger.exception(f"Failed to create entity in DB: {entity_in}") raise @@ -93,7 +93,7 @@ class MemoryService: entities.append(entity) logger.debug(f"Successfully created {len(entities)} entities in DB") return entities - except Exception as e: + except Exception: # On failure, we should try to clean up any files we wrote logger.exception("Failed to create entities in DB") for entity in entities_in: @@ -106,9 +106,9 @@ class MemoryService: logger.error(f"Failed to clean up file for {entity.id}: {cleanup_error}") raise - async def get_entity(self, entity_id): + async def get_entity(self, entity_id: str): logger.debug(f"Get entity {entity_id} entities") - entity = self.entity_service.get_entity(entity_id) + entity = await self.entity_service.get_entity(entity_id) logger.debug(f"Found entity {entity}") return entity @@ -146,7 +146,7 @@ class MemoryService: relation = await self.relation_service.create_relation(relation) relations.append(relation) logger.debug(f"Created relation in DB: {relation.id}") - except Exception as e: + except Exception: logger.exception(f"Failed to create relation: {relation}") raise @@ -167,7 +167,7 @@ class MemoryService: # Create new observations for the entity for obs in observations_in.observations: - entity.observations.append(obs) + entity.observations.append(ObservationIn(content=obs)) logger.debug(f"Added {len(observations_in.observations)} observations to entity") # Write updated entity file @@ -180,7 +180,7 @@ class MemoryService: logger.debug(f"Added {len(added_observations)} observations to DB") return added_observations - except Exception as e: + except Exception: logger.exception(f"Failed to add observations to entity: {observations_in.entity_id}") raise @@ -193,33 +193,33 @@ class MemoryService: async def delete_relations(self, relations: List[Dict[str, Any]]) -> None: pass - async def read_graph(self) -> List[Entity]: + async def read_graph(self) -> Sequence[Entity]: """Read the entire knowledge graph.""" logger.debug("Reading entire knowledge graph") try: entities = await self.entity_service.get_all() logger.debug(f"Read {len(entities)} entities from graph") return entities - except Exception as e: + except Exception: logger.exception("Failed to read graph") raise - async def search_nodes(self, query: str) -> List[Entity]: + async def search_nodes(self, query: str) -> Sequence[Entity]: """Search for nodes in the knowledge graph.""" logger.debug(f"Searching nodes with query: {query}") try: results = await self.entity_service.search(query) logger.debug(f"Found {len(results)} matches for '{query}'") return results - except Exception as e: + except Exception: logger.exception(f"Failed to search nodes with query: {query}") raise - async def open_nodes(self, names: List[str]) -> List[Entity]: + async def open_nodes(self, names: List[str]) -> List[EntityIn]: """Get specific nodes and their relationships.""" logger.debug(f"Opening nodes: {names}") - async def read_node(name: str) -> Optional[Entity]: + async def read_node(name: str) -> Optional[EntityIn]: try: # Get ID from name first logger.debug(f"Looking up entity: {name}") @@ -231,7 +231,7 @@ class MemoryService: return entity logger.debug(f"Entity not found: {name}") return None - except Exception as e: + except Exception: logger.exception(f"Failed to read node: {name}") return None @@ -240,7 +240,7 @@ class MemoryService: if entity is not None] logger.debug(f"Opened {len(entities)} entities") return entities - except Exception as e: + except Exception: logger.exception("Failed to open nodes") raise diff --git a/src/basic_memory/services/observation_service.py b/src/basic_memory/services/observation_service.py index e4194b17..f1f108e8 100644 --- a/src/basic_memory/services/observation_service.py +++ b/src/basic_memory/services/observation_service.py @@ -28,8 +28,7 @@ class ObservationService: return await self.observation_repo.bulk_create([ Observation( entity_id=entity_id, - content=observation.content, - context=observation.context + content=observation, ) for observation in observations ]) diff --git a/tests/api/test_knowledge.py b/tests/api/test_knowledge.py index 276200ad..7bb5ceed 100644 --- a/tests/api/test_knowledge.py +++ b/tests/api/test_knowledge.py @@ -1,17 +1,13 @@ """Tests for knowledge graph API endpoints.""" -from pathlib import Path from typing import AsyncGenerator import pytest import pytest_asyncio from fastapi import FastAPI from httpx import AsyncClient, ASGITransport -from unittest.mock import AsyncMock - from icecream import ic from loguru import logger from basic_memory.deps import get_project_config, get_engine -from basic_memory.models import Entity @pytest_asyncio.fixture @@ -41,8 +37,7 @@ async def client(app: FastAPI) -> AsyncGenerator[AsyncClient, None]: @pytest.mark.asyncio async def test_create_entities(client: AsyncClient): """Should create entities successfully.""" - - # Make request like a real client would + # Create an entity response = await client.post("/knowledge/entities", json={ "entities": [{ "name": "Test Entity", @@ -50,11 +45,206 @@ async def test_create_entities(client: AsyncClient): }] }) - logger.debug(ic(response.content)) - - - # Verify response + # Verify creation assert response.status_code == 200 data = response.json() assert len(data["entities"]) == 1 - assert data["entities"][0]["id"] == "test/test_entity" + entity = data["entities"][0] + assert entity["id"] == "test/test_entity" + assert entity["name"] == "Test Entity" + + entity_type = entity.get("entity_type") or entity.get("entityType") + assert entity_type == "test" + + +@pytest.mark.asyncio +async def test_get_entity(client: AsyncClient): + """Should retrieve an entity by ID.""" + # First create an entity + create_response = await client.post("/knowledge/entities", json={ + "entities": [{ + "name": "Test Entity", + "entity_type": "test", + }] + }) + entity_id = create_response.json()["entities"][0]["id"] + + # Now get it by ID + response = await client.get(f"/knowledge/entities/{entity_id}") + + # Verify retrieval + assert response.status_code == 200 + entity = response.json() + assert entity["id"] == entity_id + assert entity["name"] == "Test Entity" + + entity_type = entity.get("entity_type") or entity.get("entityType") + assert entity_type == "test" + + + +@pytest.mark.asyncio +async def test_create_relations(client: AsyncClient): + """Should create relations between entities.""" + # Create two entities to relate + entities = [ + {"name": "Source Entity", "entity_type": "test"}, + {"name": "Target Entity", "entity_type": "test"} + ] + create_response = await client.post("/knowledge/entities", json={"entities": entities}) + created = create_response.json()["entities"] + source_id = created[0]["id"] + target_id = created[1]["id"] + + # Create relation between them + response = await client.post("/knowledge/relations", json={ + "relations": [{ + "from_id": source_id, + "to_id": target_id, + "relation_type": "related_to" + }] + }) + + # Verify relation + assert response.status_code == 200 + data = response.json() + assert len(data["relations"]) == 1 + relation = data["relations"][0] + assert relation["from_id"] == source_id + assert relation["to_id"] == target_id + assert relation["relation_type"] == "related_to" + + +@pytest.mark.asyncio +async def test_add_observations(client: AsyncClient): + """Should add observations to an entity.""" + # Create an entity first + create_response = await client.post("/knowledge/entities", json={ + "entities": [{ + "name": "Test Entity", + "entity_type": "test" + }] + }) + entity_id = create_response.json()["entities"][0]["id"] + + # Add observations + response = await client.post("/knowledge/observations", json={ + "entity_id": entity_id, + "observations": [ + "First observation", + "Second observation" + ] + }) + + logger.debug(ic(response.content)) + + # Verify observations were added + assert response.status_code == 200 + data = response.json() + assert data["entity_id"] == entity_id + assert len(data["observations"]) == 2 + assert data["observations"][0]["content"] == "First observation" + assert data["observations"][1]["content"] == "Second observation" + + # Verify observations appear in entity + entity_response = await client.get(f"/knowledge/entities/{entity_id}") + entity = entity_response.json() + assert len(entity["observations"]) == 2 + + +@pytest.mark.asyncio +async def test_search_nodes(client: AsyncClient): + """Should search for entities in the knowledge graph.""" + # Create a few entities with different names + entities = [ + {"name": "Not found", "entity_type": "negative"}, + {"name": "Alpha Test", "entity_type": "test"}, + {"name": "Beta Test", "entity_type": "test"}, + {"name": "Gamma Production", "entity_type": "test"} #match entity_type + ] + await client.post("/knowledge/entities", json={"entities": entities}) + + # Search for "Test" in names + response = await client.post("/knowledge/search", json={"query": "Test"}) + + # Verify search results + assert response.status_code == 200 + data = response.json() + assert data["query"] == "Test" + assert len(data["matches"]) == 3 + names = [entity["name"] for entity in data["matches"]] + assert "Alpha Test" in names + assert "Beta Test" in names + assert "Gamma Production" in names + + +@pytest.mark.asyncio +async def test_full_knowledge_flow(client: AsyncClient): + """Test a complete knowledge graph flow with multiple operations.""" + # 1. Create main entity + main_response = await client.post("/knowledge/entities", json={ + "entities": [{ + "name": "Main Entity", + "entity_type": "test" + }, + { + "name": "Non Entity", + "entity_type": "n_a" + }] + }) + main_id = main_response.json()["entities"][0]["id"] + assert main_response.status_code == 200 + assert main_id is not None + + # 2. Create related entities + related_response = await client.post("/knowledge/entities", json={ + "entities": [ + {"name": "Related One", "entity_type": "test"}, + {"name": "Related Two", "entity_type": "test"} + ] + }) + related = related_response.json()["entities"] + related_ids = [e["id"] for e in related] + assert related_response.status_code == 200 + assert len(related_ids) == 2 + + # 3. Add relations + relations_response = await client.post("/knowledge/relations", json={ + "relations": [ + { + "from_id": main_id, + "to_id": related_ids[0], + "relation_type": "connects_to" + }, + { + "from_id": main_id, + "to_id": related_ids[1], + "relation_type": "connects_to" + } + ] + }) + assert relations_response.status_code == 200 + assert len(relations_response.json()["relations"]) == 2 + + # 4. Add observations to main entity + await client.post("/knowledge/observations", json={ + "entity_id": main_id, + "observations": [ + "Connected to first related entity", + "Connected to second related entity" + ] + }) + + # 5. Verify full graph structure + main_get = await client.get(f"/knowledge/entities/{main_id}") + main_entity = main_get.json() + + # Check entity structure + assert main_entity["name"] == "Main Entity" + assert len(main_entity["observations"]) == 2 + assert len(main_entity["relations"]) == 2 + + # 6. Search should find all related entities + search = await client.post("/knowledge/search", json={"query": "Related"}) + matches = search.json()["matches"] + assert len(matches) == 3 # Should find both related entities \ No newline at end of file diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index dd73a692..4f1cfb6c 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -2,8 +2,7 @@ import pytest from basic_memory.services import MemoryService from basic_memory.fileio import read_entity_file -from basic_memory.models import Entity as EntityModel, Observation, Relation -from basic_memory.schemas import EntityIn, CreateEntitiesInput, CreateRelationsInput, ObservationsIn, RelationIn +from basic_memory.schemas import CreateEntitiesInput, CreateRelationsInput, ObservationsIn, RelationIn test_entities_data = [ { @@ -60,8 +59,8 @@ async def test_add_observations(memory_service: MemoryService): observations_data = { "entity_id": entity.id, "observations": [ - {"content": "New observation 1"}, - {"content": "New observation 2", "context": "test context"} + "New observation 1", + "New observation 2" ] } @@ -81,7 +80,6 @@ async def test_add_observations(memory_service: MemoryService): assert len(updated_entity.observations) == 4 # 2 original + 2 new assert updated_entity.observations[2].content == "New observation 1" assert updated_entity.observations[3].content == "New observation 2" - assert updated_entity.observations[3].context == "test context" # Verify database - returns SQLAlchemy Entity db_entity = await memory_service.entity_service.get_entity(entity.id) diff --git a/tests/test_memory_service_observations.py b/tests/test_memory_service_observations.py deleted file mode 100644 index 8e1b4551..00000000 --- a/tests/test_memory_service_observations.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Tests for observation creation through memory service.""" -import pytest -from basic_memory.services.memory_service import MemoryService -from basic_memory.schemas import EntityIn, ObservationIn - -pytestmark = pytest.mark.anyio - - -async def test_create_entity_with_observations(tmp_path, memory_service: MemoryService): - """Test creating an entity with observations in exactly the same way as create_entities tool.""" - # Mirror exact structure from create_entities tool - entity_data = { - "name": "Directory Organization", - "entityType": "memory", - "description": "Implemented filesystem organization by entity type", - "observations": [ - {"content": "Files are now organized by type using directories like entities/project/basic_memory"}, - {"content": "Entity IDs match filesystem paths for better mental model"}, - {"content": "Fixed path handling bugs by adding consistent get_entity_path helper"} - ] - } - - # Create entity via memory service - entity = await memory_service.create_entities([EntityIn(**entity_data)]) - - assert len(entity) == 1 - assert entity[0].name == "Directory Organization" - assert len(entity[0].observations) == 3 diff --git a/tests/test_observation_service.py b/tests/test_observation_service.py index 9727c58c..d41af308 100644 --- a/tests/test_observation_service.py +++ b/tests/test_observation_service.py @@ -2,7 +2,7 @@ import pytest from basic_memory.models import Observation -from basic_memory.schemas import EntityIn, ObservationIn +from basic_memory.schemas import ObservationIn pytestmark = pytest.mark.asyncio @@ -11,7 +11,6 @@ async def test_add_observation_success(observation_service, test_entity): """Test successful observation addition.""" observation_data = ObservationIn( content="New observation", - context="test-context" ) # Act @@ -25,7 +24,7 @@ async def test_add_observation_success(observation_service, test_entity): # Verify database index db_observations = await observation_service.observation_repo.find_by_entity(test_entity.id) assert len(db_observations) == 1 - assert any(obs.content == "New observation" and obs.context == "test-context" + assert any(obs.content == "New observation" for obs in db_observations) @@ -46,21 +45,6 @@ async def test_search_observations(observation_service, test_entity): assert results[0].content == "Unique test content" -async def test_get_observations_by_context(observation_service, test_entity): - """Test retrieving observations by context.""" - # Arrange - await observation_service.add_observations( - test_entity.id, - [ObservationIn(content="Context observation", context="test-context"), - ObservationIn(content="Other observation", context="other-context")] - ) - - # Act - results = await observation_service.get_observations_by_context("test-context") - - # Assert - assert len(results) == 1 - assert results[0].content == "Context observation" # Edge Cases diff --git a/tests/test_repository.py b/tests/test_repository.py deleted file mode 100644 index c963af2d..00000000 --- a/tests/test_repository.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Tests for base repository functionality.""" -import pytest -from sqlalchemy import select - -from basic_memory.models import Base, Observation -from basic_memory.repository import Repository - -pytestmark = pytest.mark.anyio - - -async def test_create_with_defaults(session): - """Test creating an entity with default timestamps.""" - repo = Repository(session, Observation) - - # Create observation without timestamps - observation_data = { - 'entity_id': 'test/test_entity', - 'content': 'Test observation' - } - - # Should succeed even though created_at not provided - observation = await repo.create(observation_data) - assert observation.id is not None - assert observation.content == 'Test observation' - assert observation.created_at is not None # Should have default value - - # Verify in database - stmt = select(Observation).where(Observation.id == observation.id) - result = await session.execute(stmt) - db_observation = result.scalar_one() - assert db_observation.created_at is not None \ No newline at end of file diff --git a/tests/test_repository_sql.py b/tests/test_repository_sql.py deleted file mode 100644 index 52f80465..00000000 --- a/tests/test_repository_sql.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Tests for repository behavior using raw SQL.""" -import pytest -from sqlalchemy import text - -pytestmark = pytest.mark.anyio - - -async def test_sqlite_default_timestamp(session): - """Test SQLite's handling of default CURRENT_TIMESTAMP.""" - # First create test entity - await session.execute( - text(""" - INSERT INTO entity (id, name, entity_type, description) - VALUES ('test/test_entity', 'Test', 'test', 'Test description') - """) - ) - await session.commit() - - # Now try to insert observation with no timestamp - await session.execute( - text(""" - INSERT INTO observation (entity_id, content) - VALUES ('test/test_entity', 'Test observation') - """) - ) - await session.commit() - - # Verify timestamp was set - result = await session.execute( - text("SELECT created_at FROM observation WHERE entity_id = 'test/test_entity'") - ) - observation = result.fetchone() - assert observation is not None - assert observation.created_at is not None \ No newline at end of file diff --git a/tests/test_schemas.py b/tests/test_schemas.py index 0686b5f8..a54a7873 100644 --- a/tests/test_schemas.py +++ b/tests/test_schemas.py @@ -1,14 +1,11 @@ """Tests for Pydantic schema validation and conversion.""" import pytest -from datetime import datetime from pydantic import ValidationError from basic_memory.schemas import ( EntityIn, EntityOut, ObservationIn, - ObservationOut, RelationIn, - RelationOut, CreateEntitiesInput, SearchNodesInput, OpenNodesInput, @@ -18,7 +15,7 @@ def test_entity_in_minimal(): """Test creating EntityIn with minimal required fields.""" data = { "name": "test_entity", - "entityType": "test" + "entity_type": "test" } entity = EntityIn.model_validate(data) assert entity.name == "test_entity" @@ -31,16 +28,16 @@ def test_entity_in_complete(): """Test creating EntityIn with all fields.""" data = { "name": "test_entity", - "entityType": "test", + "entity_type": "test", "description": "A test entity", "observations": [ {"content": "Test observation"} ], "relations": [ { - "fromId": "123", - "toId": "456", - "relationType": "test_relation" + "from_id": "123", + "to_id": "456", + "relation_type": "test_relation" } ] } @@ -69,11 +66,9 @@ def test_observation_in_validation(): # Minimal obs = ObservationIn.model_validate({"content": "test"}) assert obs.content == "test" - assert obs.context is None # With context obs = ObservationIn.model_validate({"content": "test", "context": "test context"}) - assert obs.context == "test context" # Missing content with pytest.raises(ValidationError): @@ -82,9 +77,9 @@ def test_observation_in_validation(): def test_relation_in_validation(): """Test RelationIn validation.""" data = { - "fromId": "123", - "toId": "456", - "relationType": "test" + "from_id": "123", + "to_id": "456", + "relation_type": "test" } relation = RelationIn.model_validate(data) assert relation.from_id == "123" @@ -99,7 +94,7 @@ def test_relation_in_validation(): # Missing required fields with pytest.raises(ValidationError): - RelationIn.model_validate({"fromId": "123", "toId": "456"}) # Missing relationType + RelationIn.model_validate({"from_id": "123", "to_id": "456"}) # Missing relationType def test_create_entities_input(): """Test CreateEntitiesInput validation.""" @@ -107,11 +102,11 @@ def test_create_entities_input(): "entities": [ { "name": "entity1", - "entityType": "test" + "entity_type": "test" }, { "name": "entity2", - "entityType": "test", + "entity_type": "test", "description": "test description" } ] @@ -124,30 +119,6 @@ def test_create_entities_input(): with pytest.raises(ValidationError): CreateEntitiesInput.model_validate({"entities": []}) -def test_snake_case_to_camel(): - """Test conversion from snake_case to camelCase.""" - data = { - "name": "test", - "entityType": "test", - "observations": [ - {"content": "test"} - ], - "relations": [ - { - "fromId": "123", - "toId": "456", - "relationType": "test", - } - ] - } - entity = EntityIn.model_validate(data) - # Access fields using snake_case - assert entity.entity_type == "test" - rel = entity.relations[0] - assert rel.from_id == "123" - assert rel.to_id == "456" - assert rel.relation_type == "test" - def test_entity_out_from_attributes(): """Test EntityOut creation from database model attributes.""" # Simulate database model attributes @@ -180,7 +151,7 @@ def test_entity_out_from_attributes(): def test_optional_fields(): """Test handling of optional fields.""" # Create with no optional fields - entity = EntityIn.model_validate({"name": "test", "entityType": "test"}) + entity = EntityIn.model_validate({"name": "test", "entity_type": "test"}) assert entity.description is None assert entity.observations == [] assert entity.relations == [] @@ -188,7 +159,7 @@ def test_optional_fields(): # Create with empty optional fields entity = EntityIn.model_validate({ "name": "test", - "entityType": "test", + "entity_type": "test", "description": None, "observations": [], "relations": [] @@ -200,7 +171,7 @@ def test_optional_fields(): # Create with some optional fields entity = EntityIn.model_validate({ "name": "test", - "entityType": "test", + "entity_type": "test", "description": "test", "observations": [] })