add index_document to search_service.py

This commit is contained in:
phernandez
2025-01-04 19:53:50 -06:00
parent 0b918ee9e8
commit 6b94ec57a4
4 changed files with 114 additions and 15 deletions
@@ -8,9 +8,10 @@ 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.schemas.search import SearchQuery, SearchResult, SearchItemType
from basic_memory.models.search import CREATE_SEARCH_INDEX
class SearchRepository():
"""Repository for search index operations."""
@@ -38,7 +39,8 @@ class SearchRepository():
# Handle type filter
if query.types:
type_list = ", ".join(f"'{t}'" for t in query.types)
# Get string values from enums
type_list = ", ".join(f"'{t.value}'" for t in query.types)
conditions.append(f"type IN ({type_list})")
# Handle entity type filter
@@ -78,7 +80,7 @@ class SearchRepository():
SearchResult(
path_id=row.path_id,
file_path=row.file_path,
type=row.type,
type=SearchItemType(row.type), # Convert string to enum
score=row.score,
metadata=json.loads(row.metadata)
)
@@ -90,7 +92,7 @@ class SearchRepository():
content: str,
path_id: str,
file_path: str,
type: str,
type: SearchItemType, # Now accepts enum
metadata: dict
):
"""Index or update a single item."""
@@ -114,7 +116,7 @@ class SearchRepository():
"content": content,
"path_id": path_id,
"file_path": file_path,
"type": type,
"type": type.value, # Store the string value
"metadata": json.dumps(metadata)
}
)
+15 -2
View File
@@ -1,16 +1,29 @@
"""Search schemas for Basic Memory."""
from typing import Optional, List
from datetime import datetime
from enum import Enum
from pydantic import BaseModel
class SearchItemType(str, Enum):
"""Types of searchable items."""
DOCUMENT = "document"
ENTITY = "entity"
class SearchQuery(BaseModel):
"""Search query parameters."""
text: str
types: Optional[List[str]] = None
types: Optional[List[SearchItemType]] = None
entity_types: Optional[List[str]] = None
after_date: Optional[datetime] = None
class SearchResult(BaseModel):
"""Search result item."""
path_id: str
file_path: str
type: str
type: SearchItemType
score: float
metadata: dict
+31 -4
View File
@@ -3,7 +3,7 @@
from typing import List, Optional
from basic_memory.repository.search_repository import SearchRepository
from basic_memory.schemas.search import SearchQuery, SearchResult
from basic_memory.schemas.search import SearchQuery, SearchResult, SearchItemType
class SearchService:
@@ -52,7 +52,7 @@ class SearchService:
content=content,
path_id=entity.path_id,
file_path=entity.file_path,
type="entity",
type=SearchItemType.ENTITY,
metadata=metadata,
)
else:
@@ -60,10 +60,37 @@ class SearchService:
content=content,
path_id=entity.path_id,
file_path=entity.file_path,
type="entity",
type=SearchItemType.ENTITY,
metadata=metadata,
)
async def index_document(self, document, content: str, background_tasks=None):
"""Index a document and its content."""
metadata = {
**document.doc_metadata,
"created_at": document.created_at.isoformat(),
"updated_at": document.updated_at.isoformat(),
}
# Queue indexing if background_tasks provided
if background_tasks:
background_tasks.add_task(
self._do_index,
content=content,
path_id=document.path_id,
file_path=document.file_path,
type=SearchItemType.DOCUMENT,
metadata=metadata,
)
else:
await self._do_index(
content=content,
path_id=document.path_id,
file_path=document.file_path,
type=SearchItemType.DOCUMENT,
metadata=metadata,
)
async def _do_index(self, **kwargs):
"""Actually perform the indexing."""
await self.repository.index_item(**kwargs)
await self.repository.index_item(**kwargs)