search service implementation

This commit is contained in:
phernandez
2025-01-04 16:37:02 -06:00
parent 35c21c1891
commit 905ff78aea
7 changed files with 343 additions and 3 deletions
+3 -3
View File
@@ -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}')"
+15
View File
@@ -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
);
""")
@@ -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()
+16
View File
@@ -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
@@ -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)