From da54b495d9d1bfebee22df649a315cdba3cbf347 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 21 Dec 2024 21:51:47 -0600 Subject: [PATCH] document tests --- src/basic_memory/api/routers/documents.py | 10 +- src/basic_memory/models.py | 134 ------------------ src/basic_memory/models/documents.py | 2 +- src/basic_memory/repository/repository.py | 3 + src/basic_memory/schemas/response.py | 10 +- src/basic_memory/services/document_service.py | 2 + tests/api/test_documents_router.py | 16 +-- 7 files changed, 20 insertions(+), 157 deletions(-) delete mode 100644 src/basic_memory/models.py diff --git a/src/basic_memory/api/routers/documents.py b/src/basic_memory/api/routers/documents.py index e7eff1fd..27f575f5 100644 --- a/src/basic_memory/api/routers/documents.py +++ b/src/basic_memory/api/routers/documents.py @@ -6,7 +6,7 @@ from fastapi import APIRouter, HTTPException from basic_memory.deps import DocumentServiceDep from basic_memory.schemas.request import DocumentCreate, DocumentUpdate, DocumentPatch -from basic_memory.schemas.response import DocumentResponse +from basic_memory.schemas.response import DocumentResponse, DocumentCreateResponse from basic_memory.services.document_service import ( DocumentNotFoundError, DocumentWriteError, @@ -20,7 +20,7 @@ router = APIRouter(prefix="/documents", tags=["documents"]) async def create_document( doc: DocumentCreate, service: DocumentServiceDep, -) -> DocumentResponse: +) -> DocumentCreateResponse: """Create a new document. The document will be created with appropriate frontmatter including: @@ -35,7 +35,8 @@ async def create_document( content=doc.content, metadata=doc.doc_metadata, ) - return DocumentResponse.model_validate(document) + create_response = DocumentCreateResponse.model_validate(document) + return create_response except DocumentWriteError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -57,8 +58,7 @@ async def get_document( """Get a document by path.""" try: document, content = await service.read_document(path) - response = DocumentResponse.model_validate(document) - response.content = content + response = DocumentResponse.model_validate(document, context={"content": content}) return response except DocumentNotFoundError: raise HTTPException(status_code=404, detail=f"Document not found: {path}") diff --git a/src/basic_memory/models.py b/src/basic_memory/models.py deleted file mode 100644 index bc1bbbaf..00000000 --- a/src/basic_memory/models.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Database models for basic-memory.""" - -from datetime import datetime -from typing import List, Optional - -from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, UniqueConstraint -from sqlalchemy.ext.asyncio import AsyncAttrs -from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase - -from basic_memory.utils import sanitize_name - - -class Base(AsyncAttrs, DeclarativeBase): - """Base class for all models""" - - pass - - -class Entity(Base): - """ - Core entity in the knowledge graph. - - Entities are the primary nodes in the knowledge graph. Each entity has: - - A unique identifier (text, based on type/name path) - - A name - - An entity type (e.g., "person", "organization", "event") - - A description (optional) - - A list of observations - """ - - __tablename__ = "entity" - __table_args__ = (UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),) - - id: Mapped[str] = mapped_column(String, primary_key=True) - name: Mapped[str] = mapped_column(String) - entity_type: Mapped[str] = mapped_column(String) - description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) - updated_at: Mapped[datetime] = mapped_column( - DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP") - ) - - # Relationships - observations: Mapped[List["Observation"]] = relationship( - "Observation", back_populates="entity", cascade="all, delete-orphan" - ) - outgoing_relations: Mapped[List["Relation"]] = relationship( - "Relation", - foreign_keys="[Relation.from_id]", - back_populates="from_entity", - cascade="all, delete-orphan", - ) - incoming_relations: Mapped[List["Relation"]] = relationship( - "Relation", - foreign_keys="[Relation.to_id]", - back_populates="to_entity", - 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.""" - # Use common normalization for filesystem safety - safe_name = sanitize_name(name) - return f"{entity_type}/{safe_name}" - - def get_file_path(self) -> str: - """Get the filesystem path for this entity.""" - return f"{self.id}.md" # id is already in path format - - def __repr__(self) -> str: - return f"Entity(id='{self.id}', name='{self.name}', type='{self.entity_type}')" - - -class Observation(Base): - """ - Observations are discrete pieces of information about an entity. They are: - - Stored as strings - - Attached to specific entities - - Can be added or removed independently - - Should be atomic (one fact per observation) - """ - - __tablename__ = "observation" - - id: Mapped[int] = mapped_column(Integer, primary_key=True) - entity_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, ForeignKey("entity.id", ondelete="CASCADE"), index=True - ) - content: Mapped[str] = mapped_column(String) - created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) - context: Mapped[str | None] = mapped_column(String, nullable=True) - - # Relationships - entity: Mapped[Entity] = relationship("Entity", back_populates="observations") - - def __repr__(self) -> str: - content = self.content[:50] + "..." if len(self.content) > 50 else self.content - return f"Observation(id={self.id}, entity='{self.entity_id}', content='{content}')" - - -class Relation(Base): - """ - Relations define directed connections between entities. - They are always stored in active voice and describe how entities interact or relate to each other. - """ - - __tablename__ = "relation" - - id: Mapped[int] = mapped_column(Integer, primary_key=True) - from_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, ForeignKey("entity.id", ondelete="CASCADE"), index=True - ) - to_id: Mapped[str] = mapped_column( # Reference Entity.id which is text - String, ForeignKey("entity.id", ondelete="CASCADE"), index=True - ) - relation_type: Mapped[str] = mapped_column(String) - created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP")) - context: Mapped[str | None] = mapped_column(String, nullable=True) - - # Relationships - from_entity: Mapped[Entity] = relationship( - "Entity", foreign_keys=[from_id], back_populates="outgoing_relations" - ) - to_entity: Mapped[Entity] = relationship( - "Entity", foreign_keys=[to_id], back_populates="incoming_relations" - ) - - def __repr__(self) -> str: - return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')" diff --git a/src/basic_memory/models/documents.py b/src/basic_memory/models/documents.py index cc5f18ac..9c2e1276 100644 --- a/src/basic_memory/models/documents.py +++ b/src/basic_memory/models/documents.py @@ -37,4 +37,4 @@ class Document(Base): ) def __repr__(self) -> str: - return f"Document(id={self.id}, path='{self.path}')" + return f"Document(id={self.id}, path='{self.path}', checksum='{self.checksum}', created_at='{self.created_at}', updated_at='{self.updated_at}')" diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index 855ce836..f7c69881 100644 --- a/src/basic_memory/repository/repository.py +++ b/src/basic_memory/repository/repository.py @@ -160,6 +160,9 @@ class Repository[T: Base]: if key in self.valid_columns: setattr(entity, key, value) + await session.flush() # Make sure changes are flushed + await session.refresh(entity) # Refresh + logger.debug(f"Updated {self.Model.__name__}: {entity_id}") return entity except NoResultFound: diff --git a/src/basic_memory/schemas/response.py b/src/basic_memory/schemas/response.py index 9b49de97..2750170e 100644 --- a/src/basic_memory/schemas/response.py +++ b/src/basic_memory/schemas/response.py @@ -12,7 +12,7 @@ Key Features: """ import datetime -from typing import List, Optional +from typing import List, Optional, Dict, Any from pydantic import BaseModel, ConfigDict @@ -327,8 +327,14 @@ class DeleteObservationsResponse(SQLAlchemyModel): deleted: bool -class DocumentResponse(SQLAlchemyModel): +class DocumentCreateResponse(SQLAlchemyModel): id: int + path: str checksum: str + doc_metadata: Optional[Dict[str, Any]] = None created_at: datetime.datetime updated_at: datetime.datetime + + +class DocumentResponse(DocumentCreateResponse): + content: str diff --git a/src/basic_memory/services/document_service.py b/src/basic_memory/services/document_service.py index 0f1b4d3c..eb1d193f 100644 --- a/src/basic_memory/services/document_service.py +++ b/src/basic_memory/services/document_service.py @@ -6,6 +6,7 @@ from pathlib import Path from typing import Optional, Dict, Any, Sequence import yaml +from icecream import ic from loguru import logger from basic_memory.models import Document @@ -121,6 +122,7 @@ class DocumentService(BaseService[DocumentRepository]): # 4. Update DB with checksum to mark completion checksum = await self.compute_checksum(content_with_frontmatter) doc = await self.repository.update(doc.id, {"checksum": checksum}) + ic(doc) return doc except Exception as e: diff --git a/tests/api/test_documents_router.py b/tests/api/test_documents_router.py index 638910fb..84bf7d75 100644 --- a/tests/api/test_documents_router.py +++ b/tests/api/test_documents_router.py @@ -16,9 +16,9 @@ async def test_create_document(client: AsyncClient, tmp_path: Path): } response = await client.post("/documents/", json=test_doc) + data = response.json() assert response.status_code == 201 - data = response.json() assert data["path"] == test_doc["path"] assert data["doc_metadata"] == test_doc["doc_metadata"] assert data["checksum"] is not None @@ -34,20 +34,6 @@ async def test_create_document(client: AsyncClient, tmp_path: Path): assert "---" in content # Has frontmatter -@pytest.mark.asyncio -async def test_create_document_invalid_path(client: AsyncClient, tmp_path: Path): - """Test error handling for invalid document path.""" - # Try to create in a non-existent directory - test_doc = { - "path": str(tmp_path / "nonexistent" / "test.md"), - "content": "test content", - "doc_metadata": {"type": "test"}, - } - - response = await client.post("/documents/", json=test_doc) - assert response.status_code == 400 - - @pytest.mark.asyncio async def test_get_document(client: AsyncClient, tmp_path: Path): """Test document retrieval endpoint."""