From bcc5f8a165bed9646768d34081b68111c7cebe36 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 30 Dec 2024 13:06:41 -0600 Subject: [PATCH] add get_entity_types endpoint --- src/basic_memory/api/app.py | 2 + src/basic_memory/api/routers/__init__.py | 3 +- .../api/routers/discovery_router.py | 17 +++++ .../repository/entity_repository.py | 2 +- src/basic_memory/repository/repository.py | 4 +- src/basic_memory/schemas/__init__.py | 9 +++ src/basic_memory/schemas/discovery.py | 14 ++++ src/basic_memory/services/entity_service.py | 5 ++ tests/api/test_discovery_router.py | 66 +++++++++++++++++++ 9 files changed, 118 insertions(+), 4 deletions(-) create mode 100644 src/basic_memory/api/routers/discovery_router.py create mode 100644 src/basic_memory/schemas/discovery.py create mode 100644 tests/api/test_discovery_router.py diff --git a/src/basic_memory/api/app.py b/src/basic_memory/api/app.py index 99df9868..b48a1939 100644 --- a/src/basic_memory/api/app.py +++ b/src/basic_memory/api/app.py @@ -5,6 +5,7 @@ from loguru import logger from .routers import documents from .routers import knowledge +from .routers import discovery # Initialize FastAPI app app = FastAPI( @@ -14,6 +15,7 @@ app = FastAPI( # Include routers app.include_router(knowledge.router) app.include_router(documents.router) +app.include_router(discovery.router) # Add startup event diff --git a/src/basic_memory/api/routers/__init__.py b/src/basic_memory/api/routers/__init__.py index 557697f3..b12cc979 100644 --- a/src/basic_memory/api/routers/__init__.py +++ b/src/basic_memory/api/routers/__init__.py @@ -2,5 +2,6 @@ from . import knowledge_router as knowledge from . import documents_router as documents +from . import discovery_router as discovery -__all__ = ["knowledge", "documents"] +__all__ = ["knowledge", "documents", "discovery"] diff --git a/src/basic_memory/api/routers/discovery_router.py b/src/basic_memory/api/routers/discovery_router.py new file mode 100644 index 00000000..2d277d5c --- /dev/null +++ b/src/basic_memory/api/routers/discovery_router.py @@ -0,0 +1,17 @@ +"""Router for knowledge discovery and analytics operations.""" + +from fastapi import APIRouter +from loguru import logger + +from basic_memory.deps import EntityServiceDep +from basic_memory.schemas import EntityTypeList, ObservationCategoryList + +router = APIRouter(prefix="/discovery", tags=["discovery"]) + + +@router.get("/entity-types", response_model=EntityTypeList) +async def get_entity_types(entity_service: EntityServiceDep) -> EntityTypeList: + """Get list of all unique entity types in the system.""" + logger.debug("Getting all entity types") + types = await entity_service.get_entity_types() + return EntityTypeList(types=types) diff --git a/src/basic_memory/repository/entity_repository.py b/src/basic_memory/repository/entity_repository.py index 9be02d1a..e078d0ab 100644 --- a/src/basic_memory/repository/entity_repository.py +++ b/src/basic_memory/repository/entity_repository.py @@ -43,7 +43,7 @@ class EntityRepository(Repository[Entity]): """Get list of distinct entity types.""" query = select(Entity.entity_type).distinct() - result = await self.execute_query(query) + result = await self.execute_query(query, use_query_options=False) return list(result.scalars().all()) async def search(self, query_str: str) -> List[Entity]: diff --git a/src/basic_memory/repository/repository.py b/src/basic_memory/repository/repository.py index 522b7749..6c1904fc 100644 --- a/src/basic_memory/repository/repository.py +++ b/src/basic_memory/repository/repository.py @@ -246,10 +246,10 @@ class Repository[T: Base]: logger.debug(f"Counted {count} {self.Model.__name__} records") return count - async def execute_query(self, query: Executable) -> Result[Any]: + async def execute_query(self, query: Executable, use_query_options:bool = True) -> Result[Any]: """Execute a query asynchronously.""" - query = query.options(*self.get_load_options()) + query = query.options(*self.get_load_options()) if use_query_options else query logger.debug(f"Executing query: {query}") async with db.scoped_session(self.session_maker) as session: diff --git a/src/basic_memory/schemas/__init__.py b/src/basic_memory/schemas/__init__.py index 38fbd867..6daf51be 100644 --- a/src/basic_memory/schemas/__init__.py +++ b/src/basic_memory/schemas/__init__.py @@ -41,6 +41,12 @@ from basic_memory.schemas.response import ( DeleteEntitiesResponse, ) +# Discovery and analytics models +from basic_memory.schemas.discovery import ( + EntityTypeList, + ObservationCategoryList, +) + # For convenient imports, export all models __all__ = [ # Base @@ -67,4 +73,7 @@ __all__ = [ "DeleteEntitiesRequest", "DeleteRelationsRequest", "DeleteObservationsRequest", + # Discovery and Analytics + "EntityTypeList", + "ObservationCategoryList", ] diff --git a/src/basic_memory/schemas/discovery.py b/src/basic_memory/schemas/discovery.py new file mode 100644 index 00000000..2977258d --- /dev/null +++ b/src/basic_memory/schemas/discovery.py @@ -0,0 +1,14 @@ +"""Schemas for knowledge discovery and analytics endpoints.""" + +from typing import List +from pydantic import BaseModel + + +class EntityTypeList(BaseModel): + """List of unique entity types in the system.""" + types: List[str] + + +class ObservationCategoryList(BaseModel): + """List of unique observation categories in the system.""" + categories: List[str] diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 39421474..94cacd2a 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -67,6 +67,11 @@ class EntityService(BaseService[EntityRepository]): async def get_all(self) -> Sequence[EntityModel]: """Get all entities.""" return await self.repository.find_all() + + async def get_entity_types(self) -> List[str]: + """Get list of all distinct entity types in the system.""" + logger.debug("Getting all distinct entity types") + return await self.repository.get_entity_types() async def delete_entity(self, path_id: str) -> bool: """Delete entity from database.""" diff --git a/tests/api/test_discovery_router.py b/tests/api/test_discovery_router.py new file mode 100644 index 00000000..c444bc51 --- /dev/null +++ b/tests/api/test_discovery_router.py @@ -0,0 +1,66 @@ +"""Tests for discovery router endpoints.""" + +import pytest +import pytest_asyncio +from httpx import AsyncClient + +from basic_memory.models.knowledge import Entity +from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.schemas import EntityTypeList + +pytestmark = pytest.mark.asyncio + + +@pytest_asyncio.fixture +async def test_entities(entity_repository: EntityRepository) -> list[Entity]: + """Create test entities with different types.""" + entities = [ + Entity( + name="Memory Service", + entity_type="technical_component", + description="Core memory service", + path_id="component/memory_service", + file_path="component/memory_service.md", + ), + Entity( + name="File Format", + entity_type="specification", + description="File format spec", + path_id="spec/file_format", + file_path="spec/file_format.md", + ), + Entity( + name="Technical Decision", + entity_type="decision", + description="Architecture decision", + path_id="decision/tech_choice", + file_path="decision/tech_choice.md", + ), + ] + + created = await entity_repository.add_all(entities) + return created + + +@pytest.mark.asyncio +async def test_get_entity_types(client: AsyncClient, test_entities): + """Test getting list of entity types.""" + # Get types + response = await client.get("/discovery/entity-types") + assert response.status_code == 200 + + # Parse response + data = EntityTypeList.model_validate(response.json()) + + # Should have types from test data + assert len(data.types) > 0 + assert "technical_component" in data.types + assert "specification" in data.types + assert "decision" in data.types + + # Types should all be strings + assert isinstance(data.types, list) + assert all(isinstance(t, str) for t in data.types) + + # Types should be unique + assert len(data.types) == len(set(data.types))