diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py index ccd2df51..53a4d6b0 100644 --- a/src/basic_memory/deps.py +++ b/src/basic_memory/deps.py @@ -12,6 +12,7 @@ from sqlalchemy.ext.asyncio import ( from basic_memory import db from basic_memory.config import ProjectConfig, config from basic_memory.markdown.knowledge_writer import KnowledgeWriter +from basic_memory.markdown.note_writer import NoteWriter from basic_memory.repository.document_repository import DocumentRepository from basic_memory.repository.entity_repository import EntityRepository from basic_memory.repository.observation_repository import ObservationRepository @@ -196,6 +197,12 @@ async def get_knowledge_writer() -> KnowledgeWriter: KnowledgeWriterDep = Annotated[KnowledgeWriter, Depends(get_knowledge_writer)] +async def get_note_writer() -> NoteWriter: + return NoteWriter() + + +NoteWriterDep = Annotated[NoteWriter, Depends(get_note_writer)] + async def get_knowledge_service( entity_service: EntityServiceDep, diff --git a/src/basic_memory/markdown/knowledge_writer.py b/src/basic_memory/markdown/knowledge_writer.py index 947d5309..6d18107e 100644 --- a/src/basic_memory/markdown/knowledge_writer.py +++ b/src/basic_memory/markdown/knowledge_writer.py @@ -1,10 +1,5 @@ """Writer for knowledge entity markdown files.""" -from typing import Optional, Dict, Any - -import yaml -from loguru import logger - from basic_memory.models import Entity as EntityModel @@ -13,34 +8,17 @@ class KnowledgeWriter: async def format_frontmatter(self, entity: EntityModel) -> dict: """Generate frontmatter metadata for entity.""" - return { + frontmatter = { "id": entity.path_id, "type": entity.entity_type, "created": entity.created_at.isoformat(), "modified": entity.updated_at.isoformat(), } + if entity.entity_metadata: + frontmatter.update(entity.entity_metadata) + return frontmatter - async def format_metadata(self, metadata: Optional[Dict[str, Any]] = None) -> str: - """Format metadata section as YAML block.""" - if not metadata: - return "" - - try: - yaml_block = yaml.dump(metadata, sort_keys=False) - return ( - "# Metadata\n" - "\n\n" - "```yml\n" - f"{yaml_block}" - "```\n" - ) - except Exception as e: - logger.warning(f"Failed to format metadata YAML: {e}") - return "" # Skip metadata on error - - async def format_content( - self, entity: EntityModel, metadata: Optional[Dict[str, Any]] = None - ) -> str: + async def format_content(self, entity: EntityModel, content: str) -> str: """Format entity content as markdown.""" sections = [ f"# {entity.name}\n", @@ -78,9 +56,5 @@ class KnowledgeWriter: # Outgoing relations (entity is "from") for rel in entity.outgoing_relations: sections.append(f"- {rel.relation_type} [[{rel.to_entity.name}]] ") - - if metadata: - sections.append("\n") - sections.append(await self.format_metadata(metadata)) - return "\n".join(sections) + return "\n".join(sections) \ No newline at end of file diff --git a/src/basic_memory/markdown/note_writer.py b/src/basic_memory/markdown/note_writer.py new file mode 100644 index 00000000..47184bba --- /dev/null +++ b/src/basic_memory/markdown/note_writer.py @@ -0,0 +1,44 @@ +"""Writer for note entity markdown files.""" + +from typing import Optional +from basic_memory.models import Entity as EntityModel + + +class NoteWriter: + """Formats notes into markdown files with frontmatter.""" + + async def format_frontmatter(self, entity: EntityModel) -> dict: + """Generate frontmatter for note. + + Args: + entity: The note entity to format frontmatter for + + Returns: + Dictionary of frontmatter fields + """ + frontmatter = { + "id": entity.path_id, + "type": entity.entity_type, + "created": entity.created_at.isoformat(), + "modified": entity.updated_at.isoformat() + } + # Add any entity_metadata that isn't None + if entity.entity_metadata: + frontmatter.update(entity.entity_metadata) + return frontmatter + + async def format_content( + self, + entity: EntityModel, + content: str, + ) -> str: + """Format note content as markdown. + + Args: + entity: The note entity + content: Raw content to format + + Returns: + Formatted markdown content + """ + return content.strip() # Just return the trimmed content \ No newline at end of file diff --git a/src/basic_memory/schemas/base.py b/src/basic_memory/schemas/base.py index 09921545..69204da7 100644 --- a/src/basic_memory/schemas/base.py +++ b/src/basic_memory/schemas/base.py @@ -256,9 +256,9 @@ class Entity(BaseModel): @property def path_id(self) -> PathId: - """Get the path ID in format {type}/{snake_case_name}.""" + """Get the path ID in format {snake_case_name}.""" normalized_name = to_snake_case(self.name) - return f"{self.entity_type}/{normalized_name}" + return normalized_name @property def file_path(self): diff --git a/src/basic_memory/services/knowledge/file_operations.py b/src/basic_memory/services/knowledge/file_operations.py index 0a7fa69e..87785f1b 100644 --- a/src/basic_memory/services/knowledge/file_operations.py +++ b/src/basic_memory/services/knowledge/file_operations.py @@ -1,30 +1,34 @@ """File operations for knowledge service.""" from pathlib import Path -from typing import Tuple +from typing import Tuple, Optional from loguru import logger from basic_memory.markdown.knowledge_writer import KnowledgeWriter +from basic_memory.markdown.note_writer import NoteWriter from basic_memory.models import Entity as EntityModel +from basic_memory.models.knowledge import EntityType from basic_memory.services.entity_service import EntityService from basic_memory.services.exceptions import FileOperationError from basic_memory.services.file_service import FileService class FileOperations: - """File operations for knowledge entities.""" + """File operations for both knowledge and note entities.""" def __init__( self, entity_service: EntityService, file_service: FileService, knowledge_writer: KnowledgeWriter, + note_writer: NoteWriter, base_path: Path, ): self.entity_service = entity_service self.file_service = file_service self.knowledge_writer = knowledge_writer + self.note_writer = note_writer self.base_path = base_path async def file_exists(self, path: Path) -> bool: @@ -35,25 +39,37 @@ class FileOperations: def get_entity_path(self, entity: EntityModel) -> Path: """Generate filesystem path for entity.""" - return self.base_path / entity.entity_type / f"{entity.name}.md" + if entity.file_path: + return self.base_path / entity.file_path + return self.base_path / f"{entity.path_id}.md" - async def write_entity_file(self, entity: EntityModel) -> Tuple[Path, str]: + async def write_entity_file( + self, + entity: EntityModel, + content: Optional[str] = None, + ) -> Tuple[Path, str]: """Write entity to filesystem and return path and checksum.""" try: - # Ensure we have a fresh entity with all relations loaded + # Ensure we have a fresh entity with all relations entity = await self.entity_service.get_by_path_id(entity.path_id) - - frontmatter = await self.knowledge_writer.format_frontmatter(entity) - # Format content - path = self.get_entity_path(entity) - entity_content = await self.knowledge_writer.format_content(entity) - file_content = await self.file_service.add_frontmatter( - frontmatter=frontmatter, - content=entity_content, + + # Select writer based on entity type + writer = self.note_writer if entity.entity_type == EntityType.NOTE else self.knowledge_writer + + # Get frontmatter and content + frontmatter = await writer.format_frontmatter(entity) + file_content = await writer.format_content( + entity=entity, + content=content or entity.description or "", ) - - # Write and get checksum - return path, await self.file_service.write_file(path, file_content) + + # Add frontmatter and write + content_with_frontmatter = await self.file_service.add_frontmatter( + frontmatter=frontmatter, + content=file_content + ) + path = self.get_entity_path(entity) + return path, await self.file_service.write_file(path, content_with_frontmatter) except Exception as e: logger.error(f"Failed to write entity file: {e}") diff --git a/src/basic_memory/services/knowledge/knowledge_service.py b/src/basic_memory/services/knowledge/knowledge_service.py index 7d9b687c..519ee271 100644 --- a/src/basic_memory/services/knowledge/knowledge_service.py +++ b/src/basic_memory/services/knowledge/knowledge_service.py @@ -19,6 +19,7 @@ from .file_operations import FileOperations from .entity_operations import EntityOperations from .relation_operations import RelationOperations from .observation_operations import ObservationOperations +from ...markdown.note_writer import NoteWriter class KnowledgeService: @@ -42,6 +43,7 @@ class KnowledgeService: relation_service: RelationService, file_service: FileService, knowledge_writer: KnowledgeWriter, + note_writer: NoteWriter, base_path: Path, ): @@ -52,6 +54,7 @@ class KnowledgeService: entity_service=entity_service, file_service=file_service, knowledge_writer=knowledge_writer, + note_writer=note_writer, base_path=base_path ) diff --git a/tests/conftest.py b/tests/conftest.py index c66b3fe1..d1d74da8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,6 +12,7 @@ from basic_memory.config import ProjectConfig from basic_memory.db import DatabaseType from basic_memory.markdown.knowledge_parser import KnowledgeParser from basic_memory.markdown.knowledge_writer import KnowledgeWriter +from basic_memory.markdown.note_writer import NoteWriter from basic_memory.models import Base from basic_memory.models.knowledge import Entity, EntityType from basic_memory.repository.document_repository import DocumentRepository @@ -148,6 +149,11 @@ def knowledge_writer(): """Create writer instance.""" return KnowledgeWriter() +@pytest.fixture +def note_writer(): + """Create writer instance.""" + return NoteWriter() + @pytest.fixture def knowledge_parser(): @@ -168,6 +174,7 @@ async def knowledge_service( relation_service: RelationService, file_service: FileService, knowledge_writer: KnowledgeWriter, + note_writer: NoteWriter, test_config: ProjectConfig, ) -> KnowledgeService: """Create KnowledgeService with dependencies.""" @@ -177,6 +184,7 @@ async def knowledge_service( relation_service=relation_service, file_service=file_service, knowledge_writer=knowledge_writer, + note_writer=note_writer, base_path=test_config.knowledge_dir, ) diff --git a/tests/markdown/test_knowledge_writer.py b/tests/markdown/test_knowledge_writer.py index a64ad0db..e8c87d26 100644 --- a/tests/markdown/test_knowledge_writer.py +++ b/tests/markdown/test_knowledge_writer.py @@ -1,199 +1,142 @@ -"""Tests for knowledge entity writer.""" +"""Tests for KnowledgeWriter.""" + from datetime import datetime, UTC import pytest - +from basic_memory.models import Entity, Observation, Relation from basic_memory.markdown.knowledge_writer import KnowledgeWriter -from basic_memory.models import ( - Entity as EntityModel, - Observation, - Relation, ObservationCategory, -) +from basic_memory.models.knowledge import EntityType @pytest.fixture -def writer(): - """Create writer instance.""" +def knowledge_writer() -> KnowledgeWriter: return KnowledgeWriter() @pytest.fixture -def test_entity(): - """Create test entity with observations and relations.""" - now = datetime.now(UTC) - # Create main entity - entity = EntityModel( +def sample_entity() -> Entity: + """Create a sample knowledge entity for testing.""" + return Entity( id=1, - path_id="test/test_entity", - name="TestEntity", - entity_type="test", - description="A test entity", - created_at=now, - updated_at=now + name="test_entity", + entity_type=EntityType.KNOWLEDGE, + path_id="knowledge/test_entity", + file_path="knowledge/test_entity.md", + description="Test description", + created_at=datetime(2025, 1, 1, tzinfo=UTC), + updated_at=datetime(2025, 1, 2, tzinfo=UTC) ) - # Add observations with categories and context - entity.observations = [ - Observation( - content="Technical implementation detail", - category=ObservationCategory.TECH.value, - context="Initial implementation" - ), - Observation( - content="Design pattern choice", - category=ObservationCategory.DESIGN.value - ), - ] - # Create related entity - other_entity = EntityModel( +@pytest.fixture +def entity_with_observations(sample_entity: Entity) -> Entity: + """Create an entity with observations.""" + sample_entity.observations = [ + Observation(entity_id=1, category="tech", content="First observation"), + Observation(entity_id=1, category="design", content="Second observation", context="Some context") + ] + return sample_entity + + +@pytest.fixture +def entity_with_relations(sample_entity: Entity) -> Entity: + """Create an entity with relations.""" + target = Entity( id=2, - path_id="test/other_entity", - name="OtherEntity", - entity_type="test", - created_at=now, - updated_at=now + name="target_entity", + entity_type=EntityType.KNOWLEDGE, + path_id="knowledge/target_entity" ) - - # Create relation from main entity to other - relation = Relation(from_entity=entity, to_entity=other_entity, relation_type="relates_to") - entity.outgoing_relations = [relation] - - return entity - - -@pytest.mark.asyncio -async def test_format_frontmatter(writer: KnowledgeWriter, test_entity: EntityModel): - """Test frontmatter generation.""" - frontmatter = await writer.format_frontmatter(test_entity) - - assert frontmatter["type"] == "test" - assert frontmatter["id"] == "test/test_entity" - assert isinstance(frontmatter["created"], str) - assert isinstance(frontmatter["modified"], str) - - -@pytest.mark.asyncio -async def test_format_content_with_categories(writer: KnowledgeWriter, test_entity: EntityModel): - """Test content formatting with categorized observations.""" - content = await writer.format_content(test_entity) - - # Check observations section header and format comment - assert "## Observations" in content - assert "" in content - - # Check formatted observations - assert "- [tech] Technical implementation detail (Initial implementation)" in content - assert "- [design] Design pattern choice" in content - - -@pytest.mark.asyncio -async def test_format_content_default_category(writer: KnowledgeWriter): - """Test formatting observation with default category.""" - entity = EntityModel(id=1, name="Test", entity_type="test") - entity.observations = [ - Observation(content="Simple note", category=ObservationCategory.NOTE.value) + sample_entity.outgoing_relations = [ + Relation( + from_id=1, + to_id=2, + relation_type="connects_to", + to_entity=target + ) ] - - content = await writer.format_content(entity) - assert "- [note] Simple note" in content + return sample_entity @pytest.mark.asyncio -async def test_format_content_context_handling(writer: KnowledgeWriter): - """Test formatting observations with different context scenarios.""" - entity = EntityModel(id=1, name="Test", entity_type="test") - entity.observations = [ - # With context - Observation( - content="With context", - category=ObservationCategory.TECH.value, - context="Important context" - ), - # Without context - Observation( - content="No context", - category=ObservationCategory.TECH.value - ), - ] - - content = await writer.format_content(entity) - assert "- [tech] With context (Important context)" in content - assert "- [tech] No context" in content - assert "No context ()" not in content # Shouldn't have empty parentheses +async def test_format_frontmatter_basic(knowledge_writer: KnowledgeWriter, sample_entity: Entity): + """Test basic frontmatter formatting.""" + frontmatter = await knowledge_writer.format_frontmatter(sample_entity) + + assert frontmatter["id"] == "knowledge/test_entity" + assert frontmatter["type"] == EntityType.KNOWLEDGE + assert frontmatter["created"] == "2025-01-01T00:00:00+00:00" + assert frontmatter["modified"] == "2025-01-02T00:00:00+00:00" @pytest.mark.asyncio -async def test_format_content_sections_order(writer: KnowledgeWriter, test_entity: EntityModel): - """Test proper order and spacing of sections with new format.""" - content = await writer.format_content(test_entity) - lines = content.split("\n") - - # Find key sections - title_idx = next(i for i, line in enumerate(lines) if line.startswith("# ")) - obs_idx = next(i for i, line in enumerate(lines) if line.strip() == "## Observations") - format_idx = next(i for i, line in enumerate(lines) if "" in content - assert "```yml" in content - assert "ai_generated: true" in content.lower() - assert "confidence: 0.95" in content - assert "tags:" in content - assert "- test" in content - assert "- example" in content +async def test_format_content_basic(knowledge_writer: KnowledgeWriter, sample_entity: Entity): + """Test basic content formatting.""" + content = "" + result = await knowledge_writer.format_content(sample_entity, content) + + assert "# test_entity" in result + assert "Test description" in result @pytest.mark.asyncio -async def test_format_metadata_only(writer: KnowledgeWriter): - """Test metadata formatting alone.""" - metadata = {"test": "value", "nested": {"key": "value"}} - - content = await writer.format_metadata(metadata) - - assert content.startswith("# Metadata\n") - assert "" in content - assert "```yml" in content - assert "test: value" in content - assert "nested:" in content - assert " key: value" in content - assert content.strip().endswith("```") +async def test_format_content_with_observations( + knowledge_writer: KnowledgeWriter, + entity_with_observations: Entity +): + """Test content formatting with observations.""" + content = "" + result = await knowledge_writer.format_content(entity_with_observations, content) + + assert "## Observations" in result + assert "- [tech] First observation" in result + assert "- [design] Second observation (Some context)" in result @pytest.mark.asyncio -async def test_format_metadata_empty(writer: KnowledgeWriter): - """Test metadata formatting with empty/none metadata.""" - assert await writer.format_metadata(None) == "" - assert await writer.format_metadata({}) == "" +async def test_format_content_with_relations( + knowledge_writer: KnowledgeWriter, + entity_with_relations: Entity +): + """Test content formatting with relations.""" + content = "" + result = await knowledge_writer.format_content(entity_with_relations, content) + + assert "## Relations" in result + assert "- connects_to [[target_entity]]" in result @pytest.mark.asyncio -async def test_format_content_minimal_entity(writer: KnowledgeWriter): - """Test formatting with minimal entity.""" - entity = EntityModel(id=1, name="Minimal", entity_type="test") - - content = await writer.format_content(entity) - - # Should have title only - assert content.strip() == "# Minimal" - +async def test_format_content_full_entity( + knowledge_writer: KnowledgeWriter, + entity_with_relations: Entity, + entity_with_observations: Entity +): + """Test content formatting with all entity features.""" + # Combine observations and relations + entity_with_relations.observations = entity_with_observations.observations + content = "" + result = await knowledge_writer.format_content(entity_with_relations, content) + + # Verify all sections present + assert "# test_entity" in result + assert "Test description" in result + assert "## Observations" in result + assert "- [tech] First observation" in result + assert "## Relations" in result + assert "- connects_to [[target_entity]]" in result \ No newline at end of file diff --git a/tests/markdown/test_note_writer.py b/tests/markdown/test_note_writer.py new file mode 100644 index 00000000..fd0b788f --- /dev/null +++ b/tests/markdown/test_note_writer.py @@ -0,0 +1,91 @@ +"""Tests for NoteWriter.""" + +from datetime import datetime, UTC + +import pytest +from basic_memory.models import Entity +from basic_memory.markdown.note_writer import NoteWriter +from basic_memory.models.knowledge import EntityType + + +@pytest.fixture +def note_writer() -> NoteWriter: + return NoteWriter() + + +@pytest.fixture +def sample_note() -> Entity: + """Create a sample note entity for testing.""" + return Entity( + id=1, + name="test_note", + entity_type=EntityType.NOTE, + path_id="notes/test_note", + file_path="notes/test_note.md", + created_at=datetime(2025, 1, 1, tzinfo=UTC), + updated_at=datetime(2025, 1, 2, tzinfo=UTC) + ) + + +@pytest.mark.asyncio +async def test_format_frontmatter_basic(note_writer: NoteWriter, sample_note: Entity): + """Test basic frontmatter formatting.""" + frontmatter = await note_writer.format_frontmatter(sample_note) + + assert frontmatter["id"] == "notes/test_note" + assert frontmatter["type"] == EntityType.NOTE + assert frontmatter["created"] == "2025-01-01T00:00:00+00:00" + assert frontmatter["modified"] == "2025-01-02T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_format_frontmatter_with_metadata(note_writer: NoteWriter, sample_note: Entity): + """Test frontmatter includes entity metadata.""" + sample_note.entity_metadata = { + "category": "research", + "tags": ["python", "testing"] + } + + frontmatter = await note_writer.format_frontmatter(sample_note) + + assert frontmatter["category"] == "research" + assert frontmatter["tags"] == ["python", "testing"] + assert frontmatter["id"] == "notes/test_note" + + +@pytest.mark.asyncio +async def test_format_content_basic(note_writer: NoteWriter, sample_note: Entity): + """Test basic content formatting.""" + content = "# Test Note\n\nThis is a test note." + result = await note_writer.format_content(sample_note, content) + + assert result == content + + +@pytest.mark.asyncio +async def test_format_content_strips_whitespace(note_writer: NoteWriter, sample_note: Entity): + """Test content formatting strips extra whitespace.""" + content = "\n\n# Test Note\n\nThis is a test note.\n\n" + result = await note_writer.format_content(sample_note, content) + + assert result == "# Test Note\n\nThis is a test note." + + +@pytest.mark.asyncio +async def test_format_content_preserves_markdown(note_writer: NoteWriter, sample_note: Entity): + """Test content formatting preserves markdown formatting.""" + content = """# Test Note + +This note has: +- Bullet points +- *Italic text* +- **Bold text** +- `code blocks` + +```python +def test(): + pass +```""" + + result = await note_writer.format_content(sample_note, content) + assert result == content \ No newline at end of file diff --git a/tests/services/test_knowledge_service.py b/tests/services/test_knowledge_service.py index 0aeeaeff..04eee9a0 100644 --- a/tests/services/test_knowledge_service.py +++ b/tests/services/test_knowledge_service.py @@ -8,6 +8,7 @@ import yaml from sqlalchemy.exc import IntegrityError from basic_memory.models import Entity as EntityModel +from basic_memory.models.knowledge import EntityType from basic_memory.schemas import Entity as EntitySchema, Relation as RelationSchema from basic_memory.schemas.base import ObservationCategory from basic_memory.schemas.request import ObservationCreate @@ -19,26 +20,28 @@ from basic_memory.services.knowledge import KnowledgeService @pytest.mark.asyncio async def test_get_entity_path(knowledge_service: KnowledgeService): """Should generate correct filesystem path for entity.""" - entity = EntityModel(id=1, name="test-entity", entity_type="concept", description="Test entity") + entity = EntityModel(id=1, path_id="test-entity", name="test-entity", entity_type=EntityType.KNOWLEDGE, description="Test entity") path = knowledge_service.get_entity_path(entity) - assert path == Path(knowledge_service.base_path / "concept/test-entity.md") + assert path == Path(knowledge_service.base_path / "test-entity.md") @pytest.mark.asyncio async def test_create_entity(knowledge_service: KnowledgeService): """Should create entity in DB and write file correctly.""" # Setup - entity = EntitySchema(name="test-entity", entity_type="concept", description="Test entity") + entity_schema = EntitySchema(name="test-entity", entity_type=EntityType.KNOWLEDGE, description="Test entity") # Execute - created = await knowledge_service.create_entity(entity) + created = await knowledge_service.create_entity(entity_schema) # Verify DB entity - assert created.name == entity.name - assert created.entity_type == entity.entity_type - assert created.description == entity.description + assert created.name == entity_schema.name + assert created.entity_type == entity_schema.entity_type + assert created.description == entity_schema.description assert created.checksum is not None - + assert created.path_id == "test_entity" + assert created.file_path == "test_entity.md" + # Verify file was written file_path = knowledge_service.get_entity_path(created) assert await knowledge_service.file_exists(file_path) @@ -48,8 +51,8 @@ async def test_create_entity(knowledge_service: KnowledgeService): metadata = yaml.safe_load(frontmatter) # Verify frontmatter contents - assert metadata["id"] == entity.path_id - assert metadata["type"] == entity.entity_type + assert metadata["id"] == entity_schema.path_id + assert metadata["type"] == entity_schema.entity_type assert "created" in metadata assert "modified" in metadata @@ -59,7 +62,7 @@ async def test_create_entity(knowledge_service: KnowledgeService): async def test_create_multiple_entities(knowledge_service: KnowledgeService): """Should create multiple entities successfully.""" entities = [ - EntitySchema(name=f"entity-{i}", entity_type="test", description=f"Test entity {i}") + EntitySchema(name=f"entity-{i}", entity_type=EntityType.KNOWLEDGE, description=f"Test entity {i}") for i in range(3) ] @@ -77,10 +80,10 @@ async def test_create_relations(knowledge_service: KnowledgeService, entity_serv """Should create relations and update related entity files.""" # Create test entities entity1 = await knowledge_service.create_entity( - EntitySchema(name="entity1", entity_type="test", description="Test entity 1") + EntitySchema(name="entity1", entity_type=EntityType.KNOWLEDGE, description="Test entity 1") ) entity2 = await knowledge_service.create_entity( - EntitySchema(name="entity2", entity_type="test", description="Test entity 2") + EntitySchema(name="entity2", entity_type=EntityType.KNOWLEDGE, description="Test entity 2") ) # Create relation @@ -114,7 +117,7 @@ async def test_add_observations_observation(knowledge_service: KnowledgeService) """Should add observations and update entity file.""" # Create test entity entity = await knowledge_service.create_entity( - EntitySchema(name="test", entity_type="test", description="Test entity") + EntitySchema(name="test", entity_type=EntityType.KNOWLEDGE, description="Test entity") ) # Add observations @@ -153,7 +156,7 @@ async def test_delete_entity(knowledge_service: KnowledgeService): """Should delete entity and its file.""" # Create test entity entity = await knowledge_service.create_entity( - EntitySchema(name="test", entity_type="test", description="Test entity") + EntitySchema(name="test", entity_type=EntityType.KNOWLEDGE, description="Test entity") ) file_path = knowledge_service.get_entity_path(entity) @@ -179,7 +182,7 @@ async def test_delete_multiple_entities(knowledge_service: KnowledgeService): entities = [] for i in range(3): entity = await knowledge_service.create_entity( - EntitySchema(name=f"test-{i}", entity_type="test", description=f"Test entity {i}") + EntitySchema(name=f"test-{i}", entity_type=EntityType.KNOWLEDGE, description=f"Test entity {i}") ) entities.append(entity) @@ -206,7 +209,7 @@ async def test_handle_file_operation_errors(knowledge_service: KnowledgeService, with pytest.raises(FileOperationError): await knowledge_service.create_entity( - EntitySchema(name="test", entity_type="test", description="Test entity") + EntitySchema(name="test", entity_type=EntityType.KNOWLEDGE, description="Test entity") ) @@ -240,7 +243,7 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, # Attempt creation (should fail) with pytest.raises(FileOperationError): await knowledge_service.create_entity( - EntitySchema(name="test", entity_type="test", description="Test entity") + EntitySchema(name="test", entity_type=EntityType.KNOWLEDGE, description="Test entity") ) # Verify entity was cleaned up @@ -253,9 +256,9 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, async def test_skip_failed_batch_operations(knowledge_service: KnowledgeService): """Should continue processing batch operations if some fail.""" entities = [ - EntitySchema(name="test-1", entity_type="test", description="Test entity 1"), - EntitySchema(name="test-1", entity_type="test", description="Duplicate name - should fail"), - EntitySchema(name="test-2", entity_type="test", description="Test entity 2"), + EntitySchema(name="test-1", entity_type=EntityType.KNOWLEDGE, description="Test entity 1"), + EntitySchema(name="test-1", entity_type=EntityType.KNOWLEDGE, description="Duplicate name - should fail"), + EntitySchema(name="test-2", entity_type=EntityType.KNOWLEDGE, description="Test entity 2"), ] with pytest.raises(IntegrityError): @@ -267,10 +270,10 @@ async def test_update_relations_in_files(knowledge_service: KnowledgeService): """Should update both entity files when creating relations.""" # Create test entities entity1 = await knowledge_service.create_entity( - EntitySchema(name="source", entity_type="test", description="Source entity") + EntitySchema(name="source", entity_type=EntityType.KNOWLEDGE, description="Source entity") ) entity2 = await knowledge_service.create_entity( - EntitySchema(name="target", entity_type="test", description="Target entity") + EntitySchema(name="target", entity_type=EntityType.KNOWLEDGE, description="Target entity") ) # Create relation