From 905ff78aea7e354a7edd263c5216c995bcf4b8b7 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 4 Jan 2025 16:37:02 -0600 Subject: [PATCH] search service implementation --- src/basic_memory/models/knowledge.py | 6 +- src/basic_memory/models/search.py | 15 ++ .../repository/search_repository.py | 130 ++++++++++++++++++ src/basic_memory/schemas/search.py | 16 +++ src/basic_memory/services/search_service.py | 69 ++++++++++ tests/{ => schemas}/test_activity_schemas.py | 0 tests/services/test_search_service.py | 110 +++++++++++++++ 7 files changed, 343 insertions(+), 3 deletions(-) create mode 100644 src/basic_memory/models/search.py create mode 100644 src/basic_memory/repository/search_repository.py create mode 100644 src/basic_memory/schemas/search.py create mode 100644 src/basic_memory/services/search_service.py rename tests/{ => schemas}/test_activity_schemas.py (100%) create mode 100644 tests/services/test_search_service.py diff --git a/src/basic_memory/models/knowledge.py b/src/basic_memory/models/knowledge.py index 0ea11dc9..2275e505 100644 --- a/src/basic_memory/models/knowledge.py +++ b/src/basic_memory/models/knowledge.py @@ -25,6 +25,7 @@ class Entity(Base): __tablename__ = "entity" __table_args__ = ( UniqueConstraint("entity_type", "name", name="uix_entity_type_name"), + UniqueConstraint("path_id", name="uix_entity_path_id"), # Make path_id unique Index("ix_entity_type", "entity_type"), Index("ix_entity_doc_id", "doc_id"), Index("ix_entity_created_at", "created_at"), # For timeline queries @@ -35,8 +36,7 @@ class Entity(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True) name: Mapped[str] = mapped_column(String) entity_type: Mapped[str] = mapped_column(String) - # Normalized path for URIs - # (entity_type, path_id) are unique + # Normalized path for URIs - must be unique path_id: Mapped[str] = mapped_column(String, index=True) # Actual filesystem relative path file_path: Mapped[str] = mapped_column(String, unique=True, index=True) @@ -157,4 +157,4 @@ class Relation(Base): to_entity = relationship("Entity", foreign_keys=[to_id], back_populates="incoming_relations") def __repr__(self) -> str: - return f"Relation(id={self.id}, from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')" + return f"Relation(id={self.id}, from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')" \ No newline at end of file diff --git a/src/basic_memory/models/search.py b/src/basic_memory/models/search.py new file mode 100644 index 00000000..49b9132d --- /dev/null +++ b/src/basic_memory/models/search.py @@ -0,0 +1,15 @@ +"""Search models and tables.""" + +from sqlalchemy import DDL + +# Define FTS5 virtual table creation +CREATE_SEARCH_INDEX = DDL(""" +CREATE VIRTUAL TABLE IF NOT EXISTS search_index USING fts5( + content, -- Searchable text content + path_id UNINDEXED, -- Link to entity/document (must be unique) + file_path UNINDEXED, -- Filesystem path + type UNINDEXED, -- 'entity' or 'document' + metadata UNINDEXED, -- JSON with timestamps, types, etc. + tokenize='porter unicode61' -- Enable stemming + unicode +); +""") \ No newline at end of file diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py new file mode 100644 index 00000000..0e5116db --- /dev/null +++ b/src/basic_memory/repository/search_repository.py @@ -0,0 +1,130 @@ +"""Repository for search operations.""" + +import json +from typing import List, Optional +from datetime import datetime +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.repository.repository import Repository +from basic_memory.schemas.search import SearchQuery, SearchResult +from basic_memory.models.search import CREATE_SEARCH_INDEX + +class SearchRepository(): + """Repository for search index operations.""" + + def __init__(self, session_maker: async_sessionmaker[AsyncSession]): + self.session_maker = session_maker + + async def init_search_index(self): + """Create or recreate the search index.""" + async with db.scoped_session(self.session_maker) as session: + await session.execute(CREATE_SEARCH_INDEX) + await session.commit() + + async def search( + self, + query: SearchQuery, + context: Optional[List[str]] = None + ) -> List[SearchResult]: + """Search across all indexed content.""" + conditions = [] + params = {} + + # Handle text search + if query.text: + conditions.append(f"content MATCH '{query.text}'") + + # Handle type filter + if query.types: + type_list = ", ".join(f"'{t}'" for t in query.types) + conditions.append(f"type IN ({type_list})") + + # Handle entity type filter + if query.entity_types: + entity_type_list = ", ".join(f"'{t}'" for t in query.entity_types) + conditions.append( + f"json_extract(metadata, '$.entity_type') IN ({entity_type_list})" + ) + + # Handle date filter + if query.after_date: + params["after_date"] = query.after_date.isoformat() + conditions.append( + "json_extract(metadata, '$.created_at') > :after_date" + ) + + # Build WHERE clause + where_clause = " AND ".join(conditions) if conditions else "1=1" + + sql = f""" + SELECT + path_id, + file_path, + type, + metadata, + bm25(search_index) as score + FROM search_index + WHERE {where_clause} + ORDER BY score DESC + """ + + async with db.scoped_session(self.session_maker) as session: + result = await session.execute(text(sql), params) + rows = result.fetchall() + + return [ + SearchResult( + path_id=row.path_id, + file_path=row.file_path, + type=row.type, + score=row.score, + metadata=json.loads(row.metadata) + ) + for row in rows + ] + + async def index_item( + self, + content: str, + path_id: str, + file_path: str, + type: str, + metadata: dict + ): + """Index or update a single item.""" + async with db.scoped_session(self.session_maker) as session: + # Delete existing record if any + await session.execute( + text("DELETE FROM search_index WHERE path_id = :path_id"), + {"path_id": path_id} + ) + + # Insert new record + await session.execute( + text(""" + INSERT INTO search_index ( + content, path_id, file_path, type, metadata + ) VALUES ( + :content, :path_id, :file_path, :type, :metadata + ) + """), + { + "content": content, + "path_id": path_id, + "file_path": file_path, + "type": type, + "metadata": json.dumps(metadata) + } + ) + await session.commit() + + async def delete_by_path(self, path_id: str): + """Delete an item from the search index.""" + async with db.scoped_session(self.session_maker) as session: + await session.execute( + text("DELETE FROM search_index WHERE path_id = :path_id"), + {"path_id": path_id} + ) + await session.commit() \ No newline at end of file diff --git a/src/basic_memory/schemas/search.py b/src/basic_memory/schemas/search.py new file mode 100644 index 00000000..07ed749d --- /dev/null +++ b/src/basic_memory/schemas/search.py @@ -0,0 +1,16 @@ +from typing import Optional, List +from datetime import datetime +from pydantic import BaseModel + +class SearchQuery(BaseModel): + text: str + types: Optional[List[str]] = None + entity_types: Optional[List[str]] = None + after_date: Optional[datetime] = None + +class SearchResult(BaseModel): + path_id: str + file_path: str + type: str + score: float + metadata: dict \ No newline at end of file diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py new file mode 100644 index 00000000..85518486 --- /dev/null +++ b/src/basic_memory/services/search_service.py @@ -0,0 +1,69 @@ +"""Service for search operations.""" + +from typing import List, Optional + +from basic_memory.repository.search_repository import SearchRepository +from basic_memory.schemas.search import SearchQuery, SearchResult + + +class SearchService: + """Service for search operations.""" + + def __init__(self, search_repository: SearchRepository): + self.repository = search_repository + + async def init_search_index(self): + """Create FTS5 virtual table if it doesn't exist.""" + await self.repository.init_search_index() + + async def search( + self, query: SearchQuery, context: Optional[List[str]] = None + ) -> List[SearchResult]: + """Search across all indexed content.""" + return await self.repository.search(query, context) + + async def index_entity(self, entity, background_tasks=None): + """Index an entity and its components.""" + # Build searchable content + content = "\n".join( + [ + entity.name, + entity.description or "", + # Add observations + *[f"{obs.category}: {obs.content}" for obs in entity.observations], + # Add relations + *[ + f"{rel.relation_type} {rel.to_id}: {rel.context or ''}" + for rel in entity.relations + ], + ] + ) + + metadata = { + "entity_type": entity.entity_type, + "created_at": entity.created_at.isoformat(), + "updated_at": entity.updated_at.isoformat(), + } + + # Queue indexing if background_tasks provided + if background_tasks: + background_tasks.add_task( + self._do_index, + content=content, + path_id=entity.path_id, + file_path=entity.file_path, + type="entity", + metadata=metadata, + ) + else: + await self._do_index( + content=content, + path_id=entity.path_id, + file_path=entity.file_path, + type="entity", + metadata=metadata, + ) + + async def _do_index(self, **kwargs): + """Actually perform the indexing.""" + await self.repository.index_item(**kwargs) diff --git a/tests/test_activity_schemas.py b/tests/schemas/test_activity_schemas.py similarity index 100% rename from tests/test_activity_schemas.py rename to tests/schemas/test_activity_schemas.py diff --git a/tests/services/test_search_service.py b/tests/services/test_search_service.py new file mode 100644 index 00000000..dbe7d080 --- /dev/null +++ b/tests/services/test_search_service.py @@ -0,0 +1,110 @@ +import pytest +from datetime import datetime, timezone + +import pytest_asyncio +from sqlalchemy import text + +from basic_memory import db +from basic_memory.repository.search_repository import SearchRepository +from basic_memory.schemas.search import SearchQuery +from basic_memory.services.search_service import SearchService + +@pytest_asyncio.fixture +async def search_repository(session_maker): + """Create SearchRepository instance""" + return SearchRepository(session_maker) + +@pytest_asyncio.fixture +async def search_service(search_repository: SearchRepository): + """Create and initialize search service""" + service = SearchService(search_repository) + await service.init_search_index() + return service + +@pytest.fixture +def test_entity(): + """Create a test entity""" + class Entity: + id = 1 + name = "TestComponent" + entity_type = "component" + path_id = "component/test_component" + file_path = "entities/component/test_component.md" + description = "A test component for search" + created_at = datetime.now(timezone.utc) + updated_at = datetime.now(timezone.utc) + observations = [] + relations = [] + return Entity() + +@pytest.mark.asyncio +async def test_init_search_index(search_service, session_maker): + """Test search index initialization""" + # Check that table exists + async with db.scoped_session(session_maker) as session: + result = await session.execute(text( + "SELECT name FROM sqlite_master WHERE type='table' AND name='search_index';" + )) + assert result.scalar() == "search_index" + +@pytest.mark.asyncio +async def test_index_entity(search_service, test_entity): + """Test indexing an entity""" + await search_service.index_entity(test_entity) + + # Search for the entity + results = await search_service.search(SearchQuery(text="test component")) + assert len(results) == 1 + assert results[0].path_id == test_entity.path_id + +@pytest.mark.asyncio +async def test_search_filtering(search_service, test_entity): + """Test search with filters""" + await search_service.index_entity(test_entity) + + # Search with type filter + results = await search_service.search( + SearchQuery( + text="test", + types=["entity"], + entity_types=["component"] + ) + ) + assert len(results) == 1 + + # Search with wrong type (should return no results) + results = await search_service.search( + SearchQuery( + text="test", + types=["document"] + ) + ) + assert len(results) == 0 + +@pytest.mark.asyncio +async def test_update_index(search_service, test_entity): + """Test updating indexed content""" + await search_service.index_entity(test_entity) + + # Update entity + test_entity.description = "Updated description with new terms" + await search_service.index_entity(test_entity) + + # Search for new terms + results = await search_service.search(SearchQuery(text="new terms")) + assert len(results) == 1 + +@pytest.mark.asyncio +async def test_search_date_filter(search_service, test_entity): + """Test searching with date filter""" + await search_service.index_entity(test_entity) + + # Search with future date (should return no results) + future = datetime.now(timezone.utc).replace(year=2026) + results = await search_service.search( + SearchQuery( + text="test", + after_date=future + ) + ) + assert len(results) == 0 \ No newline at end of file