mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
search service implementation
This commit is contained in:
@@ -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}')"
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user