From 55f7997c23d9bfea08fe7b05b7a3a0e4ee7c19be Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 22 Dec 2024 21:42:45 -0600 Subject: [PATCH] add knowledge_writer --- src/basic_memory/markdown/knowledge_writer.py | 73 ++++++++++ .../services/knowledge_service.py | 3 +- tests/markdown/test_knowledge_writer.py | 136 ++++++++++++++++++ 3 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 src/basic_memory/markdown/knowledge_writer.py create mode 100644 tests/markdown/test_knowledge_writer.py diff --git a/src/basic_memory/markdown/knowledge_writer.py b/src/basic_memory/markdown/knowledge_writer.py new file mode 100644 index 00000000..91ec913c --- /dev/null +++ b/src/basic_memory/markdown/knowledge_writer.py @@ -0,0 +1,73 @@ +"""Writer for knowledge entity markdown files.""" + +from datetime import datetime, UTC +from typing import Optional, Dict, Any + +import yaml +from loguru import logger + +from basic_memory.models import Entity as EntityModel + + +class KnowledgeWriter: + """Formats entities into markdown files.""" + + async def format_frontmatter(self, entity: EntityModel) -> dict: + """Generate frontmatter metadata for entity.""" + now = datetime.now(UTC).isoformat() + return { + "type": entity.entity_type, + "id": entity.id, + "created": now, + "modified": now + } + + 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: + """Format entity content as markdown.""" + sections = [ + f"# {entity.name}\n" + ] + + if entity.description: + sections.extend([ + entity.description, + "" + ]) + + if entity.observations: + sections.extend([ + "## Observations", + *[f"- {obs.content}" for obs in entity.observations], + "" + ]) + + if entity.relations: + sections.extend([ + "## Relations", + *[f"- [[{rel.to_entity.name}]] {rel.relation_type}" for rel in entity.relations], + "" + ]) + + if metadata: + sections.append("\n") + sections.append(await self.format_metadata(metadata)) + + return "\n".join(sections) diff --git a/src/basic_memory/services/knowledge_service.py b/src/basic_memory/services/knowledge_service.py index b9922b19..7245d6d5 100644 --- a/src/basic_memory/services/knowledge_service.py +++ b/src/basic_memory/services/knowledge_service.py @@ -10,6 +10,7 @@ from basic_memory.models import Entity as EntityModel from basic_memory.schemas import Entity as EntitySchema, Relation as RelationSchema from basic_memory.services.entity_service import EntityService from basic_memory.services.exceptions import EntityNotFoundError, FileOperationError +from basic_memory.services.file_service import FileService from basic_memory.services.observation_service import ObservationService from basic_memory.services.relation_service import RelationService @@ -31,7 +32,7 @@ class KnowledgeService: entity_service: EntityService, observation_service: ObservationService, relation_service: RelationService, - file_service, # FileService + file_service: FileService, # FileService knowledge_parser: KnowledgeParser, ): self.entity_service = entity_service diff --git a/tests/markdown/test_knowledge_writer.py b/tests/markdown/test_knowledge_writer.py new file mode 100644 index 00000000..ae10afc0 --- /dev/null +++ b/tests/markdown/test_knowledge_writer.py @@ -0,0 +1,136 @@ +"""Tests for knowledge entity writer.""" + +import pytest + +from basic_memory.markdown.knowledge_writer import KnowledgeWriter +from basic_memory.models import Entity as EntityModel, Observation, Relation + + +@pytest.fixture +def writer(): + """Create writer instance.""" + return KnowledgeWriter() + + +@pytest.fixture +def test_entity(): + """Create test entity with observations and relations.""" + # Create main entity + entity = EntityModel(id=123, name="TestEntity", entity_type="test", description="A test entity") + + # Add observations + entity.observations = [ + Observation(content="First observation"), + Observation(content="Second observation"), + ] + + # Create related entity + other_entity = EntityModel(id=456, name="OtherEntity", entity_type="test") + + # Create relation from main entity to other + relation = Relation(from_entity=entity, to_entity=other_entity, relation_type="relates_to") + entity.from_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"] == 123 + assert isinstance(frontmatter["created"], str) + assert isinstance(frontmatter["modified"], str) + + +@pytest.mark.asyncio +async def test_format_content_basic(writer: KnowledgeWriter, test_entity: EntityModel): + """Test basic content formatting without metadata.""" + content = await writer.format_content(test_entity) + + # Check sections + assert content.startswith("# TestEntity\n") + assert "A test entity" in content + assert "## Observations" in content + assert "- First observation" in content + assert "- Second observation" in content + assert "## Relations" in content + assert "- [[OtherEntity]] relates_to" in content + + # Should not have metadata section + assert "# Metadata" not in content + assert "```yml" not in content + + +@pytest.mark.asyncio +async def test_format_content_with_metadata(writer: KnowledgeWriter, test_entity: EntityModel): + """Test content formatting with metadata section.""" + metadata = {"ai_generated": True, "confidence": 0.95, "tags": ["test", "example"]} + + content = await writer.format_content(test_entity, metadata) + + # Regular content should be there + assert "# TestEntity" in content + assert "## Observations" in content + + # Metadata section should be properly formatted + assert "# Metadata" in content + assert "" 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 + + +@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("```") + + +@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({}) == "" + + +@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" + + +@pytest.mark.asyncio +async def test_content_section_spacing(writer: KnowledgeWriter, test_entity: EntityModel): + """Test proper spacing between sections.""" + content = await writer.format_content(test_entity) + lines = content.split("\n") + + # Find section headers + section_indexes = [i for i, line in enumerate(lines) if line.startswith("##")] + + for idx in section_indexes: + # Should be blank line before header + assert lines[idx - 1] == "", f"No blank line before section at line {idx}" + # Content should start right after header + assert lines[idx + 1].strip(), f"No content after section at line {idx}"