mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
perf(sync): batch file indexing in core (#726)
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
"""Database management commands."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
@@ -12,6 +13,7 @@ from basic_memory import db
|
||||
from basic_memory.cli.app import app
|
||||
from basic_memory.cli.commands.command_utils import run_with_cleanup
|
||||
from basic_memory.config import ConfigManager, ProjectMode
|
||||
from basic_memory.indexing import IndexProgress
|
||||
from basic_memory.repository import ProjectRepository
|
||||
from basic_memory.services.initialization import reconcile_projects_with_config
|
||||
from basic_memory.sync.sync_service import get_sync_service
|
||||
@@ -19,6 +21,39 @@ from basic_memory.sync.sync_service import get_sync_service
|
||||
console = Console()
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class EmbeddingProgress:
|
||||
"""Typed CLI progress payload for embedding backfills."""
|
||||
|
||||
entity_id: int
|
||||
index: int
|
||||
total: int
|
||||
|
||||
|
||||
def _format_eta(seconds: float | None) -> str:
|
||||
"""Render a compact ETA string for CLI progress descriptions."""
|
||||
if seconds is None:
|
||||
return "--:--"
|
||||
|
||||
whole_seconds = max(int(seconds), 0)
|
||||
minutes, remaining_seconds = divmod(whole_seconds, 60)
|
||||
hours, remaining_minutes = divmod(minutes, 60)
|
||||
if hours:
|
||||
return f"{hours:d}:{remaining_minutes:02d}:{remaining_seconds:02d}"
|
||||
return f"{remaining_minutes:02d}:{remaining_seconds:02d}"
|
||||
|
||||
|
||||
def _format_index_progress(progress: IndexProgress) -> str:
|
||||
"""Render typed index progress as a compact Rich task description."""
|
||||
files_per_minute = int(progress.files_per_minute) if progress.files_per_minute else 0
|
||||
return (
|
||||
" Indexing files... "
|
||||
f"{progress.files_processed}/{progress.files_total} files | "
|
||||
f"{progress.batches_completed}/{progress.batches_total} batches | "
|
||||
f"{files_per_minute}/min | ETA {_format_eta(progress.eta_seconds)}"
|
||||
)
|
||||
|
||||
|
||||
async def _reindex_projects(app_config):
|
||||
"""Reindex all projects in a single async context.
|
||||
|
||||
@@ -185,10 +220,34 @@ async def _reindex(app_config, search: bool, embeddings: bool, project: str | No
|
||||
console.print(f"\n[bold]Project: [cyan]{proj.name}[/cyan][/bold]")
|
||||
|
||||
if search:
|
||||
console.print(" Rebuilding full-text search index...")
|
||||
sync_service = await get_sync_service(proj)
|
||||
sync_dir = Path(proj.path)
|
||||
await sync_service.sync(sync_dir, project_name=proj.name)
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(" Indexing files... scanning changes", total=1)
|
||||
|
||||
async def on_index_progress(update: IndexProgress) -> None:
|
||||
total = update.files_total or 1
|
||||
completed = update.files_processed if update.files_total else 1
|
||||
progress.update(
|
||||
task,
|
||||
description=_format_index_progress(update),
|
||||
total=total,
|
||||
completed=min(completed, total),
|
||||
)
|
||||
|
||||
await sync_service.sync(
|
||||
sync_dir,
|
||||
project_name=proj.name,
|
||||
progress_callback=on_index_progress,
|
||||
)
|
||||
progress.update(task, completed=progress.tasks[task].total or 1)
|
||||
|
||||
console.print(" [green]✓[/green] Full-text search index rebuilt")
|
||||
|
||||
if embeddings:
|
||||
@@ -213,7 +272,16 @@ async def _reindex(app_config, search: bool, embeddings: bool, project: str | No
|
||||
task = progress.add_task(" Embedding entities...", total=None)
|
||||
|
||||
def on_progress(entity_id, index, total):
|
||||
progress.update(task, total=total, completed=index)
|
||||
embedding_progress = EmbeddingProgress(
|
||||
entity_id=entity_id,
|
||||
index=index,
|
||||
total=total,
|
||||
)
|
||||
progress.update(
|
||||
task,
|
||||
total=embedding_progress.total,
|
||||
completed=embedding_progress.index,
|
||||
)
|
||||
|
||||
stats = await search_service.reindex_vectors(progress_callback=on_progress)
|
||||
progress.update(task, completed=stats["total_entities"])
|
||||
|
||||
@@ -193,6 +193,11 @@ class BasicMemoryConfig(BaseSettings):
|
||||
description="Batch size for embedding generation.",
|
||||
gt=0,
|
||||
)
|
||||
semantic_embedding_request_concurrency: int = Field(
|
||||
default=4,
|
||||
description="Maximum number of concurrent provider requests for batched embedding generation when the active provider supports request-level concurrency.",
|
||||
gt=0,
|
||||
)
|
||||
semantic_embedding_sync_batch_size: int = Field(
|
||||
default=64,
|
||||
description="Batch size for vector sync orchestration flushes.",
|
||||
@@ -286,6 +291,31 @@ class BasicMemoryConfig(BaseSettings):
|
||||
description="Maximum number of files to process concurrently during sync. Limits memory usage on large projects (2000+ files). Lower values reduce memory consumption.",
|
||||
gt=0,
|
||||
)
|
||||
index_batch_size: int = Field(
|
||||
default=32,
|
||||
description="Maximum number of changed files to load into one indexing batch.",
|
||||
gt=0,
|
||||
)
|
||||
index_batch_max_bytes: int = Field(
|
||||
default=8 * 1024 * 1024,
|
||||
description="Maximum total bytes to load into one indexing batch. Large files still run as single-file batches.",
|
||||
gt=0,
|
||||
)
|
||||
index_parse_max_concurrent: int = Field(
|
||||
default=8,
|
||||
description="Maximum number of markdown parse tasks to run concurrently inside one indexing batch.",
|
||||
gt=0,
|
||||
)
|
||||
index_entity_max_concurrent: int = Field(
|
||||
default=4,
|
||||
description="Maximum number of entity create/update tasks to run concurrently inside one indexing batch.",
|
||||
gt=0,
|
||||
)
|
||||
index_metadata_update_max_concurrent: int = Field(
|
||||
default=4,
|
||||
description="Maximum number of metadata/search refresh tasks to run concurrently inside one indexing batch.",
|
||||
gt=0,
|
||||
)
|
||||
|
||||
kebab_filenames: bool = Field(
|
||||
default=False,
|
||||
|
||||
@@ -44,101 +44,6 @@ _engine: Optional[AsyncEngine] = None
|
||||
_session_maker: Optional[async_sessionmaker[AsyncSession]] = None
|
||||
|
||||
|
||||
async def _needs_semantic_embedding_backfill(
|
||||
app_config: BasicMemoryConfig,
|
||||
session_maker: async_sessionmaker[AsyncSession],
|
||||
) -> bool:
|
||||
"""Check if entities exist but vector embeddings are empty.
|
||||
|
||||
This is the reliable way to detect that embeddings need to be generated,
|
||||
regardless of how migrations were applied (fresh DB, upgrade, reset, etc.).
|
||||
"""
|
||||
if not app_config.semantic_search_enabled:
|
||||
return False
|
||||
|
||||
try:
|
||||
async with scoped_session(session_maker) as session:
|
||||
entity_count = (
|
||||
await session.execute(text("SELECT COUNT(*) FROM entity"))
|
||||
).scalar() or 0
|
||||
if entity_count == 0:
|
||||
return False
|
||||
|
||||
# Check if vector chunks table exists and is empty
|
||||
embedding_count = (
|
||||
await session.execute(text("SELECT COUNT(*) FROM search_vector_chunks"))
|
||||
).scalar() or 0
|
||||
|
||||
return embedding_count == 0
|
||||
except Exception as exc:
|
||||
# Table might not exist yet (pre-migration)
|
||||
logger.debug(f"Could not check embedding status: {exc}")
|
||||
return False
|
||||
|
||||
|
||||
async def _run_semantic_embedding_backfill(
|
||||
app_config: BasicMemoryConfig,
|
||||
session_maker: async_sessionmaker[AsyncSession],
|
||||
) -> None:
|
||||
"""Backfill semantic embeddings for all active projects/entities."""
|
||||
if not app_config.semantic_search_enabled:
|
||||
logger.info("Skipping automatic semantic embedding backfill: semantic search is disabled.")
|
||||
return
|
||||
|
||||
async with scoped_session(session_maker) as session:
|
||||
project_result = await session.execute(
|
||||
text("SELECT id, name FROM project WHERE is_active = :is_active ORDER BY id"),
|
||||
{"is_active": True},
|
||||
)
|
||||
projects = [(int(row[0]), str(row[1])) for row in project_result.fetchall()]
|
||||
|
||||
if not projects:
|
||||
logger.info("Skipping automatic semantic embedding backfill: no active projects found.")
|
||||
return
|
||||
|
||||
repository_class = (
|
||||
PostgresSearchRepository
|
||||
if app_config.database_backend == DatabaseBackend.POSTGRES
|
||||
else SQLiteSearchRepository
|
||||
)
|
||||
|
||||
total_entities = 0
|
||||
for project_id, project_name in projects:
|
||||
async with scoped_session(session_maker) as session:
|
||||
entity_result = await session.execute(
|
||||
text("SELECT id FROM entity WHERE project_id = :project_id ORDER BY id"),
|
||||
{"project_id": project_id},
|
||||
)
|
||||
entity_ids = [int(row[0]) for row in entity_result.fetchall()]
|
||||
|
||||
if not entity_ids:
|
||||
continue
|
||||
|
||||
total_entities += len(entity_ids)
|
||||
logger.info(
|
||||
"Automatic semantic embedding backfill: "
|
||||
f"project={project_name}, entities={len(entity_ids)}"
|
||||
)
|
||||
|
||||
search_repository = repository_class(
|
||||
session_maker,
|
||||
project_id=project_id,
|
||||
app_config=app_config,
|
||||
)
|
||||
batch_result = await search_repository.sync_entity_vectors_batch(entity_ids)
|
||||
if batch_result.entities_failed > 0:
|
||||
logger.warning(
|
||||
"Automatic semantic embedding backfill encountered entity failures: "
|
||||
f"project={project_name}, failed={batch_result.entities_failed}, "
|
||||
f"failed_entity_ids={batch_result.failed_entity_ids}"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Automatic semantic embedding backfill complete: "
|
||||
f"projects={len(projects)}, entities={total_entities}"
|
||||
)
|
||||
|
||||
|
||||
class DatabaseType(Enum):
|
||||
"""Types of supported databases."""
|
||||
|
||||
@@ -521,14 +426,6 @@ async def run_migrations(
|
||||
else:
|
||||
await SQLiteSearchRepository(session_maker, 1).init_search_index()
|
||||
|
||||
# Check if backfill is needed — actual backfill runs in background
|
||||
# from the MCP server lifespan to avoid blocking startup.
|
||||
if await _needs_semantic_embedding_backfill(app_config, session_maker):
|
||||
logger.info(
|
||||
"Semantic embeddings missing — backfill will run in background after startup"
|
||||
)
|
||||
else:
|
||||
logger.info("Semantic embeddings: up to date")
|
||||
except Exception as e: # pragma: no cover
|
||||
logger.error(f"Error running migrations: {e}")
|
||||
raise
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Reusable indexing primitives shared by local sync and future remote callers."""
|
||||
|
||||
from basic_memory.indexing.batch_indexer import BatchIndexer
|
||||
from basic_memory.indexing.batching import build_index_batches
|
||||
from basic_memory.indexing.models import (
|
||||
IndexedEntity,
|
||||
IndexBatch,
|
||||
IndexFileMetadata,
|
||||
IndexFileWriter,
|
||||
IndexFrontmatterUpdate,
|
||||
IndexFrontmatterWriteResult,
|
||||
IndexingBatchResult,
|
||||
IndexInputFile,
|
||||
IndexProgress,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"BatchIndexer",
|
||||
"IndexedEntity",
|
||||
"IndexBatch",
|
||||
"IndexFileMetadata",
|
||||
"IndexFileWriter",
|
||||
"IndexFrontmatterUpdate",
|
||||
"IndexFrontmatterWriteResult",
|
||||
"IndexingBatchResult",
|
||||
"IndexInputFile",
|
||||
"IndexProgress",
|
||||
"build_index_batches",
|
||||
]
|
||||
@@ -0,0 +1,556 @@
|
||||
"""Reusable batch executor for bounded-parallel file indexing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Awaitable, Callable, Mapping, TypeVar
|
||||
|
||||
from loguru import logger
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from basic_memory.config import BasicMemoryConfig
|
||||
from basic_memory.file_utils import compute_checksum, has_frontmatter
|
||||
from basic_memory.markdown.schemas import EntityMarkdown
|
||||
from basic_memory.indexing.models import (
|
||||
IndexedEntity,
|
||||
IndexFileWriter,
|
||||
IndexFrontmatterUpdate,
|
||||
IndexingBatchResult,
|
||||
IndexInputFile,
|
||||
)
|
||||
from basic_memory.models import Entity, Relation
|
||||
from basic_memory.services import EntityService
|
||||
from basic_memory.services.exceptions import SyncFatalError
|
||||
from basic_memory.services.search_service import SearchService
|
||||
from basic_memory.repository import EntityRepository, RelationRepository
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _PreparedMarkdownFile:
|
||||
file: IndexInputFile
|
||||
content: str
|
||||
final_checksum: str
|
||||
markdown: EntityMarkdown
|
||||
file_contains_frontmatter: bool
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _PreparedEntity:
|
||||
path: str
|
||||
entity_id: int
|
||||
checksum: str
|
||||
content_type: str | None
|
||||
search_content: str | None
|
||||
markdown_content: str | None = None
|
||||
|
||||
|
||||
class BatchIndexer:
|
||||
"""Index already-loaded files without assuming where they came from."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
app_config: BasicMemoryConfig,
|
||||
entity_service: EntityService,
|
||||
entity_repository: EntityRepository,
|
||||
relation_repository: RelationRepository,
|
||||
search_service: SearchService,
|
||||
file_writer: IndexFileWriter,
|
||||
) -> None:
|
||||
self.app_config = app_config
|
||||
self.entity_service = entity_service
|
||||
self.entity_repository = entity_repository
|
||||
self.relation_repository = relation_repository
|
||||
self.search_service = search_service
|
||||
self.file_writer = file_writer
|
||||
|
||||
async def index_files(
|
||||
self,
|
||||
files: Mapping[str, IndexInputFile],
|
||||
*,
|
||||
max_concurrent: int,
|
||||
parse_max_concurrent: int | None = None,
|
||||
existing_permalink_by_path: dict[str, str | None] | None = None,
|
||||
) -> IndexingBatchResult:
|
||||
"""Index one batch of loaded files with bounded concurrency."""
|
||||
if max_concurrent <= 0:
|
||||
raise ValueError("max_concurrent must be greater than zero")
|
||||
|
||||
ordered_paths = sorted(files)
|
||||
if not ordered_paths:
|
||||
return IndexingBatchResult()
|
||||
|
||||
parse_limit = parse_max_concurrent or max_concurrent
|
||||
error_by_path: dict[str, str] = {}
|
||||
|
||||
markdown_paths = [path for path in ordered_paths if self._is_markdown(files[path])]
|
||||
regular_paths = [path for path in ordered_paths if path not in markdown_paths]
|
||||
|
||||
prepared_markdown, parse_errors = await self._run_bounded(
|
||||
markdown_paths,
|
||||
limit=parse_limit,
|
||||
worker=lambda path: self._prepare_markdown_file(files[path]),
|
||||
)
|
||||
error_by_path.update(parse_errors)
|
||||
|
||||
prepared_markdown, normalization_errors = await self._normalize_markdown_batch(
|
||||
prepared_markdown,
|
||||
existing_permalink_by_path=existing_permalink_by_path,
|
||||
)
|
||||
error_by_path.update(normalization_errors)
|
||||
|
||||
indexed_entities: list[IndexedEntity] = []
|
||||
resolved_count = 0
|
||||
unresolved_count = 0
|
||||
search_indexed = 0
|
||||
|
||||
prepared_entities: dict[str, _PreparedEntity] = {}
|
||||
|
||||
markdown_upserts, markdown_errors = await self._run_bounded(
|
||||
[path for path in markdown_paths if path not in error_by_path],
|
||||
limit=max_concurrent,
|
||||
worker=lambda path: self._upsert_markdown_file(prepared_markdown[path]),
|
||||
)
|
||||
error_by_path.update(markdown_errors)
|
||||
prepared_entities.update(markdown_upserts)
|
||||
|
||||
regular_upserts, regular_errors = await self._run_bounded(
|
||||
regular_paths,
|
||||
limit=max_concurrent,
|
||||
worker=lambda path: self._upsert_regular_file(files[path]),
|
||||
)
|
||||
error_by_path.update(regular_errors)
|
||||
prepared_entities.update(regular_upserts)
|
||||
|
||||
markdown_entity_ids = [
|
||||
prepared_entities[path].entity_id
|
||||
for path in markdown_paths
|
||||
if path in prepared_entities
|
||||
]
|
||||
if markdown_entity_ids:
|
||||
resolved_count, unresolved_count = await self._resolve_batch_relations(
|
||||
markdown_entity_ids,
|
||||
max_concurrent=max_concurrent,
|
||||
)
|
||||
|
||||
refreshed_entities = await self.entity_repository.find_by_ids(
|
||||
[prepared.entity_id for prepared in prepared_entities.values()]
|
||||
)
|
||||
entities_by_id = {entity.id: entity for entity in refreshed_entities}
|
||||
|
||||
refreshed, refresh_errors = await self._run_bounded(
|
||||
[path for path in ordered_paths if path in prepared_entities],
|
||||
limit=self.app_config.index_metadata_update_max_concurrent,
|
||||
worker=lambda path: self._refresh_search_index(
|
||||
prepared_entities[path],
|
||||
entities_by_id[prepared_entities[path].entity_id],
|
||||
),
|
||||
)
|
||||
error_by_path.update(refresh_errors)
|
||||
|
||||
for path in ordered_paths:
|
||||
indexed = refreshed.get(path)
|
||||
if indexed is not None:
|
||||
indexed_entities.append(indexed)
|
||||
|
||||
search_indexed = len(indexed_entities)
|
||||
|
||||
return IndexingBatchResult(
|
||||
indexed=indexed_entities,
|
||||
errors=[(path, error_by_path[path]) for path in ordered_paths if path in error_by_path],
|
||||
relations_resolved=resolved_count,
|
||||
relations_unresolved=unresolved_count,
|
||||
search_indexed=search_indexed,
|
||||
)
|
||||
|
||||
# --- Preparation ---
|
||||
|
||||
async def _prepare_markdown_file(self, file: IndexInputFile) -> _PreparedMarkdownFile:
|
||||
if file.content is None:
|
||||
raise ValueError(f"Missing content for markdown file: {file.path}")
|
||||
|
||||
content = file.content.decode("utf-8")
|
||||
file_contains_frontmatter = has_frontmatter(content)
|
||||
final_checksum = await self._resolve_checksum(file)
|
||||
entity_markdown = await self.entity_service.entity_parser.parse_markdown_content(
|
||||
file_path=Path(file.path),
|
||||
content=content,
|
||||
mtime=file.last_modified.timestamp() if file.last_modified else None,
|
||||
ctime=file.created_at.timestamp() if file.created_at else None,
|
||||
)
|
||||
|
||||
return _PreparedMarkdownFile(
|
||||
file=file,
|
||||
content=content,
|
||||
final_checksum=final_checksum,
|
||||
markdown=entity_markdown,
|
||||
file_contains_frontmatter=file_contains_frontmatter,
|
||||
)
|
||||
|
||||
async def _normalize_markdown_batch(
|
||||
self,
|
||||
prepared_markdown: dict[str, _PreparedMarkdownFile],
|
||||
*,
|
||||
existing_permalink_by_path: dict[str, str | None] | None = None,
|
||||
) -> tuple[dict[str, _PreparedMarkdownFile], dict[str, str]]:
|
||||
if not prepared_markdown:
|
||||
return {}, {}
|
||||
|
||||
if existing_permalink_by_path is None:
|
||||
existing_permalink_by_path = {
|
||||
path: permalink
|
||||
for path, permalink in (
|
||||
await self.entity_repository.get_file_path_to_permalink_map()
|
||||
).items()
|
||||
}
|
||||
|
||||
batch_paths = set(prepared_markdown)
|
||||
reserved_permalinks = {
|
||||
permalink
|
||||
for path, permalink in existing_permalink_by_path.items()
|
||||
if path not in batch_paths and permalink
|
||||
}
|
||||
|
||||
normalized: dict[str, _PreparedMarkdownFile] = {}
|
||||
errors: dict[str, str] = {}
|
||||
|
||||
for path in sorted(prepared_markdown):
|
||||
try:
|
||||
normalized[path] = await self._normalize_markdown_file(
|
||||
prepared_markdown[path],
|
||||
reserved_permalinks,
|
||||
)
|
||||
existing_permalink_by_path[path] = normalized[path].markdown.frontmatter.permalink
|
||||
except Exception as exc:
|
||||
errors[path] = str(exc)
|
||||
logger.warning("Batch markdown normalization failed", path=path, error=str(exc))
|
||||
|
||||
return normalized, errors
|
||||
|
||||
async def _normalize_markdown_file(
|
||||
self,
|
||||
prepared: _PreparedMarkdownFile,
|
||||
reserved_permalinks: set[str],
|
||||
) -> _PreparedMarkdownFile:
|
||||
final_checksum = prepared.final_checksum
|
||||
final_content = prepared.content
|
||||
final_permalink = await self._resolve_batch_permalink(prepared, reserved_permalinks)
|
||||
|
||||
# Trigger: markdown file has no frontmatter and sync enforcement is enabled.
|
||||
# Why: downstream indexing relies on normalized metadata and stable permalinks.
|
||||
# Outcome: write derived metadata back through the storage-agnostic writer.
|
||||
if not prepared.file_contains_frontmatter and self.app_config.ensure_frontmatter_on_sync:
|
||||
frontmatter_updates = {
|
||||
"title": prepared.markdown.frontmatter.title,
|
||||
"type": prepared.markdown.frontmatter.type,
|
||||
"permalink": final_permalink,
|
||||
}
|
||||
write_result = await self.file_writer.write_frontmatter(
|
||||
IndexFrontmatterUpdate(path=prepared.file.path, metadata=frontmatter_updates)
|
||||
)
|
||||
final_checksum = write_result.checksum
|
||||
final_content = write_result.content
|
||||
prepared.markdown.frontmatter.metadata.update(frontmatter_updates)
|
||||
|
||||
# Trigger: existing markdown frontmatter may lack the canonical permalink.
|
||||
# Why: batch sync keeps permalinks stable without forcing a full rewrite when unchanged.
|
||||
# Outcome: only the permalink field is updated when it actually differs.
|
||||
elif (
|
||||
prepared.file_contains_frontmatter
|
||||
and not self.app_config.disable_permalinks
|
||||
and final_permalink != prepared.markdown.frontmatter.permalink
|
||||
):
|
||||
prepared.markdown.frontmatter.metadata["permalink"] = final_permalink
|
||||
write_result = await self.file_writer.write_frontmatter(
|
||||
IndexFrontmatterUpdate(
|
||||
path=prepared.file.path,
|
||||
metadata={"permalink": final_permalink},
|
||||
)
|
||||
)
|
||||
final_checksum = write_result.checksum
|
||||
final_content = write_result.content
|
||||
|
||||
return _PreparedMarkdownFile(
|
||||
file=prepared.file,
|
||||
content=final_content,
|
||||
final_checksum=final_checksum,
|
||||
markdown=prepared.markdown,
|
||||
file_contains_frontmatter=prepared.file_contains_frontmatter,
|
||||
)
|
||||
|
||||
async def _resolve_batch_permalink(
|
||||
self,
|
||||
prepared: _PreparedMarkdownFile,
|
||||
reserved_permalinks: set[str],
|
||||
) -> str | None:
|
||||
should_resolve_permalink = (
|
||||
not prepared.file_contains_frontmatter and self.app_config.ensure_frontmatter_on_sync
|
||||
) or (prepared.file_contains_frontmatter and not self.app_config.disable_permalinks)
|
||||
if not should_resolve_permalink:
|
||||
permalink = prepared.markdown.frontmatter.permalink
|
||||
if permalink:
|
||||
reserved_permalinks.add(permalink)
|
||||
return permalink
|
||||
|
||||
desired_permalink = await self.entity_service.resolve_permalink(
|
||||
prepared.file.path,
|
||||
markdown=prepared.markdown,
|
||||
skip_conflict_check=True,
|
||||
)
|
||||
return self._reserve_batch_permalink(desired_permalink, reserved_permalinks)
|
||||
|
||||
def _reserve_batch_permalink(
|
||||
self,
|
||||
desired_permalink: str,
|
||||
reserved_permalinks: set[str],
|
||||
) -> str:
|
||||
permalink = desired_permalink
|
||||
suffix = 1
|
||||
while permalink in reserved_permalinks:
|
||||
permalink = f"{desired_permalink}-{suffix}"
|
||||
suffix += 1
|
||||
reserved_permalinks.add(permalink)
|
||||
return permalink
|
||||
|
||||
# --- Persistence ---
|
||||
|
||||
async def _upsert_markdown_file(self, prepared: _PreparedMarkdownFile) -> _PreparedEntity:
|
||||
existing = await self.entity_repository.get_by_file_path(
|
||||
prepared.file.path,
|
||||
load_relations=False,
|
||||
)
|
||||
entity = await self.entity_service.upsert_entity_from_markdown(
|
||||
Path(prepared.file.path),
|
||||
prepared.markdown,
|
||||
is_new=existing is None,
|
||||
)
|
||||
updated = await self.entity_repository.update(
|
||||
entity.id,
|
||||
self._entity_metadata_updates(prepared.file, prepared.final_checksum),
|
||||
)
|
||||
if updated is None:
|
||||
raise ValueError(f"Failed to update markdown entity metadata for {prepared.file.path}")
|
||||
|
||||
return _PreparedEntity(
|
||||
path=prepared.file.path,
|
||||
entity_id=updated.id,
|
||||
checksum=prepared.final_checksum,
|
||||
content_type=prepared.file.content_type,
|
||||
search_content=(
|
||||
prepared.markdown.content
|
||||
if prepared.markdown.content is not None
|
||||
else prepared.content
|
||||
),
|
||||
markdown_content=prepared.content,
|
||||
)
|
||||
|
||||
async def _upsert_regular_file(self, file: IndexInputFile) -> _PreparedEntity:
|
||||
checksum = await self._resolve_checksum(file)
|
||||
existing = await self.entity_repository.get_by_file_path(file.path, load_relations=False)
|
||||
is_new_entity = existing is None
|
||||
|
||||
if existing is None:
|
||||
await self.entity_service.resolve_permalink(file.path, skip_conflict_check=True)
|
||||
entity = Entity(
|
||||
note_type="file",
|
||||
file_path=file.path,
|
||||
checksum=checksum,
|
||||
title=Path(file.path).name,
|
||||
created_at=file.created_at or datetime.now().astimezone(),
|
||||
updated_at=file.last_modified or datetime.now().astimezone(),
|
||||
content_type=file.content_type or "text/plain",
|
||||
mtime=file.last_modified.timestamp() if file.last_modified else None,
|
||||
size=file.size,
|
||||
)
|
||||
|
||||
try:
|
||||
created = await self.entity_repository.add(entity)
|
||||
entity_id = created.id
|
||||
except IntegrityError as exc:
|
||||
message = str(exc)
|
||||
if (
|
||||
"UNIQUE constraint failed: entity.file_path" in message
|
||||
or "uix_entity_file_path_project" in message
|
||||
or (
|
||||
"duplicate key value violates unique constraint" in message
|
||||
and "file_path" in message
|
||||
)
|
||||
):
|
||||
existing = await self.entity_repository.get_by_file_path(
|
||||
file.path,
|
||||
load_relations=False,
|
||||
)
|
||||
if existing is None:
|
||||
raise ValueError(
|
||||
f"Entity not found after file_path conflict: {file.path}"
|
||||
) from exc
|
||||
entity_id = existing.id
|
||||
else:
|
||||
raise
|
||||
else:
|
||||
entity_id = existing.id
|
||||
|
||||
updated = await self.entity_repository.update(
|
||||
entity_id,
|
||||
self._entity_metadata_updates(file, checksum, include_created_at=is_new_entity),
|
||||
)
|
||||
if updated is None:
|
||||
raise ValueError(f"Failed to update file entity metadata for {file.path}")
|
||||
|
||||
return _PreparedEntity(
|
||||
path=file.path,
|
||||
entity_id=updated.id,
|
||||
checksum=checksum,
|
||||
content_type=file.content_type,
|
||||
search_content=None,
|
||||
markdown_content=None,
|
||||
)
|
||||
|
||||
# --- Relations ---
|
||||
|
||||
async def _resolve_batch_relations(
|
||||
self,
|
||||
entity_ids: list[int],
|
||||
*,
|
||||
max_concurrent: int,
|
||||
) -> tuple[int, int]:
|
||||
unresolved_relation_lists = await asyncio.gather(
|
||||
*(
|
||||
self.relation_repository.find_unresolved_relations_for_entity(entity_id)
|
||||
for entity_id in entity_ids
|
||||
)
|
||||
)
|
||||
unresolved_relations = [
|
||||
relation for relation_list in unresolved_relation_lists for relation in relation_list
|
||||
]
|
||||
|
||||
if not unresolved_relations:
|
||||
return 0, 0
|
||||
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
|
||||
async def resolve_relation(relation: Relation) -> int:
|
||||
async with semaphore:
|
||||
try:
|
||||
resolved_entity = await self.entity_service.link_resolver.resolve_link(
|
||||
relation.to_name
|
||||
)
|
||||
if resolved_entity is None or resolved_entity.id == relation.from_id:
|
||||
return 0
|
||||
|
||||
try:
|
||||
await self.relation_repository.update(
|
||||
relation.id,
|
||||
{
|
||||
"to_id": resolved_entity.id,
|
||||
"to_name": resolved_entity.title,
|
||||
},
|
||||
)
|
||||
except IntegrityError:
|
||||
await self.relation_repository.delete(relation.id)
|
||||
return 1
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.warning(
|
||||
"Batch relation resolution failed",
|
||||
relation_id=relation.id,
|
||||
from_id=relation.from_id,
|
||||
to_name=relation.to_name,
|
||||
error=str(exc),
|
||||
)
|
||||
return 0
|
||||
|
||||
resolved_counts = await asyncio.gather(
|
||||
*(resolve_relation(relation) for relation in unresolved_relations)
|
||||
)
|
||||
|
||||
remaining_relation_lists = await asyncio.gather(
|
||||
*(
|
||||
self.relation_repository.find_unresolved_relations_for_entity(entity_id)
|
||||
for entity_id in entity_ids
|
||||
)
|
||||
)
|
||||
remaining_unresolved = sum(len(relations) for relations in remaining_relation_lists)
|
||||
|
||||
return sum(resolved_counts), remaining_unresolved
|
||||
|
||||
# --- Search refresh ---
|
||||
|
||||
async def _refresh_search_index(
|
||||
self, prepared: _PreparedEntity, entity: Entity
|
||||
) -> IndexedEntity:
|
||||
await self.search_service.index_entity_data(entity, content=prepared.search_content)
|
||||
return IndexedEntity(
|
||||
path=prepared.path,
|
||||
entity_id=entity.id,
|
||||
permalink=entity.permalink,
|
||||
checksum=prepared.checksum,
|
||||
content_type=prepared.content_type,
|
||||
markdown_content=prepared.markdown_content,
|
||||
)
|
||||
|
||||
# --- Helpers ---
|
||||
|
||||
async def _resolve_checksum(self, file: IndexInputFile) -> str:
|
||||
if file.checksum is not None:
|
||||
return file.checksum
|
||||
if file.content is None:
|
||||
raise ValueError(f"Missing checksum and content for file: {file.path}")
|
||||
return await compute_checksum(file.content)
|
||||
|
||||
def _entity_metadata_updates(
|
||||
self,
|
||||
file: IndexInputFile,
|
||||
checksum: str,
|
||||
*,
|
||||
include_created_at: bool = True,
|
||||
) -> dict[str, object]:
|
||||
updates: dict[str, object] = {
|
||||
"file_path": file.path,
|
||||
"checksum": checksum,
|
||||
"size": file.size,
|
||||
}
|
||||
if include_created_at and file.created_at is not None:
|
||||
updates["created_at"] = file.created_at
|
||||
if file.last_modified is not None:
|
||||
updates["updated_at"] = file.last_modified
|
||||
updates["mtime"] = file.last_modified.timestamp()
|
||||
if file.content_type is not None:
|
||||
updates["content_type"] = file.content_type
|
||||
return updates
|
||||
|
||||
def _is_markdown(self, file: IndexInputFile) -> bool:
|
||||
if file.content_type is not None:
|
||||
return file.content_type == "text/markdown"
|
||||
return Path(file.path).suffix.lower() in {".md", ".markdown"}
|
||||
|
||||
async def _run_bounded(
|
||||
self,
|
||||
paths: list[str],
|
||||
*,
|
||||
limit: int,
|
||||
worker: Callable[[str], Awaitable[T]],
|
||||
) -> tuple[dict[str, T], dict[str, str]]:
|
||||
if not paths:
|
||||
return {}, {}
|
||||
|
||||
semaphore = asyncio.Semaphore(limit)
|
||||
results: dict[str, T] = {}
|
||||
errors: dict[str, str] = {}
|
||||
|
||||
async def run(path: str) -> None:
|
||||
async with semaphore:
|
||||
try:
|
||||
results[path] = await worker(path)
|
||||
except Exception as exc:
|
||||
if isinstance(exc, SyncFatalError) or isinstance(exc.__cause__, SyncFatalError):
|
||||
raise
|
||||
errors[path] = str(exc)
|
||||
logger.warning("Batch indexing failed", path=path, error=str(exc))
|
||||
|
||||
await asyncio.gather(*(run(path) for path in paths))
|
||||
return results, errors
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Deterministic helpers for planning bounded indexing batches."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
from basic_memory.indexing.models import IndexBatch, IndexFileMetadata
|
||||
|
||||
|
||||
def build_index_batches(
|
||||
paths: Sequence[str],
|
||||
metadata_by_path: Mapping[str, IndexFileMetadata],
|
||||
*,
|
||||
max_files: int,
|
||||
max_bytes: int,
|
||||
) -> list[IndexBatch]:
|
||||
"""Build deterministic batches bounded by file count and total bytes."""
|
||||
if max_files <= 0:
|
||||
raise ValueError("max_files must be greater than zero")
|
||||
if max_bytes <= 0:
|
||||
raise ValueError("max_bytes must be greater than zero")
|
||||
|
||||
ordered_paths = sorted(paths)
|
||||
batches: list[IndexBatch] = []
|
||||
current_paths: list[str] = []
|
||||
current_bytes = 0
|
||||
|
||||
for path in ordered_paths:
|
||||
metadata = metadata_by_path.get(path)
|
||||
if metadata is None:
|
||||
raise KeyError(f"Missing metadata for path: {path}")
|
||||
|
||||
file_bytes = max(metadata.size, 0)
|
||||
|
||||
# Trigger: the next file would overflow the active batch.
|
||||
# Why: keep batches memory-bounded and predictable for both local and remote callers.
|
||||
# Outcome: flush the current batch before placing the next file.
|
||||
if current_paths and (
|
||||
len(current_paths) >= max_files or current_bytes + file_bytes > max_bytes
|
||||
):
|
||||
batches.append(IndexBatch(paths=current_paths, total_bytes=current_bytes))
|
||||
current_paths = []
|
||||
current_bytes = 0
|
||||
|
||||
# Trigger: one file is larger than the configured byte budget.
|
||||
# Why: we still need to index it, but splitting a single file is out of scope.
|
||||
# Outcome: emit a dedicated single-file batch that may exceed max_bytes.
|
||||
if file_bytes > max_bytes:
|
||||
batches.append(IndexBatch(paths=[path], total_bytes=file_bytes))
|
||||
continue
|
||||
|
||||
current_paths.append(path)
|
||||
current_bytes += file_bytes
|
||||
|
||||
if len(current_paths) >= max_files or current_bytes == max_bytes:
|
||||
batches.append(IndexBatch(paths=current_paths, total_bytes=current_bytes))
|
||||
current_paths = []
|
||||
current_bytes = 0
|
||||
|
||||
if current_paths:
|
||||
batches.append(IndexBatch(paths=current_paths, total_bytes=current_bytes))
|
||||
|
||||
return batches
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Typed models for the reusable indexing execution path."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexFileMetadata:
|
||||
"""Storage-agnostic metadata for a file queued for indexing."""
|
||||
|
||||
path: str
|
||||
size: int
|
||||
checksum: str | None = None
|
||||
content_type: str | None = None
|
||||
last_modified: datetime | None = None
|
||||
created_at: datetime | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexInputFile(IndexFileMetadata):
|
||||
"""Fully loaded file payload consumed by the batch executor."""
|
||||
|
||||
content: bytes | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexBatch:
|
||||
"""A deterministic batch of files bounded by count and total bytes."""
|
||||
|
||||
paths: list[str]
|
||||
total_bytes: int
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexProgress:
|
||||
"""Batch indexing progress emitted to callers such as the CLI."""
|
||||
|
||||
files_total: int
|
||||
files_processed: int
|
||||
batches_total: int
|
||||
batches_completed: int
|
||||
current_batch_bytes: int = 0
|
||||
files_per_minute: float = 0.0
|
||||
eta_seconds: float | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexFrontmatterUpdate:
|
||||
"""A typed frontmatter write request for a single file."""
|
||||
|
||||
path: str
|
||||
metadata: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexFrontmatterWriteResult:
|
||||
"""Typed result for a frontmatter write performed during indexing."""
|
||||
|
||||
checksum: str
|
||||
content: str
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexedEntity:
|
||||
"""Stable output describing one file that finished indexing successfully."""
|
||||
|
||||
path: str
|
||||
entity_id: int
|
||||
permalink: str | None
|
||||
checksum: str
|
||||
content_type: str | None = None
|
||||
markdown_content: str | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class IndexingBatchResult:
|
||||
"""Outcome for one batch execution."""
|
||||
|
||||
indexed: list[IndexedEntity] = field(default_factory=list)
|
||||
errors: list[tuple[str, str]] = field(default_factory=list)
|
||||
relations_resolved: int = 0
|
||||
relations_unresolved: int = 0
|
||||
search_indexed: int = 0
|
||||
|
||||
|
||||
class IndexFileWriter(Protocol):
|
||||
"""Narrow protocol for frontmatter writes during indexing."""
|
||||
|
||||
async def write_frontmatter(
|
||||
self, update: IndexFrontmatterUpdate
|
||||
) -> IndexFrontmatterWriteResult: ...
|
||||
@@ -2,7 +2,6 @@
|
||||
Basic Memory FastMCP server.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
@@ -13,12 +12,7 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.cli.auth import CLIAuth
|
||||
from basic_memory.config import BasicMemoryConfig
|
||||
from basic_memory.db import (
|
||||
scoped_session,
|
||||
_needs_semantic_embedding_backfill,
|
||||
_run_semantic_embedding_backfill,
|
||||
)
|
||||
from basic_memory.db import scoped_session
|
||||
from basic_memory.mcp.container import McpContainer, set_container
|
||||
from basic_memory.services.initialization import initialize_app
|
||||
from basic_memory import telemetry
|
||||
@@ -43,7 +37,7 @@ async def _log_embedding_status(session_maker: async_sessionmaker[AsyncSession])
|
||||
elif embedding_count == 0:
|
||||
logger.warning(
|
||||
f"Semantic embeddings: EMPTY — {entity_count} entities have no embeddings. "
|
||||
"Backfill running in background..."
|
||||
"Run 'bm reindex --embeddings' to build them."
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
@@ -54,20 +48,6 @@ async def _log_embedding_status(session_maker: async_sessionmaker[AsyncSession])
|
||||
logger.debug(f"Could not check embedding status at startup: {exc}")
|
||||
|
||||
|
||||
async def _background_embedding_backfill(
|
||||
config: BasicMemoryConfig,
|
||||
session_maker: async_sessionmaker[AsyncSession],
|
||||
) -> None:
|
||||
"""Run semantic embedding backfill in the background without blocking startup."""
|
||||
try:
|
||||
if await _needs_semantic_embedding_backfill(config, session_maker):
|
||||
logger.info("Background embedding backfill starting...")
|
||||
await _run_semantic_embedding_backfill(config, session_maker)
|
||||
await _log_embedding_status(session_maker)
|
||||
except Exception as exc:
|
||||
logger.error(f"Background embedding backfill failed: {exc}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastMCP):
|
||||
"""Lifecycle manager for the MCP server.
|
||||
@@ -133,14 +113,8 @@ async def lifespan(app: FastMCP):
|
||||
await initialize_app(container.config)
|
||||
|
||||
# Log embedding status so it's easy to spot in the logs
|
||||
backfill_task: asyncio.Task | None = None # type: ignore[type-arg]
|
||||
if config.semantic_search_enabled and db._session_maker is not None:
|
||||
await _log_embedding_status(db._session_maker)
|
||||
# Launch backfill in background so MCP server is ready immediately
|
||||
backfill_task = asyncio.create_task(
|
||||
_background_embedding_backfill(config, db._session_maker),
|
||||
name="embedding-backfill",
|
||||
)
|
||||
|
||||
# Create and start sync coordinator (lifecycle centralized in coordinator)
|
||||
sync_coordinator = container.create_sync_coordinator()
|
||||
@@ -157,14 +131,6 @@ async def lifespan(app: FastMCP):
|
||||
):
|
||||
logger.debug("Shutting down Basic Memory MCP server")
|
||||
|
||||
# Cancel embedding backfill if still running
|
||||
if backfill_task is not None and not backfill_task.done():
|
||||
backfill_task.cancel()
|
||||
try:
|
||||
await backfill_task
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Background embedding backfill cancelled during shutdown")
|
||||
|
||||
await sync_coordinator.stop()
|
||||
|
||||
# Only shutdown DB if we created it (not if test fixture provided it)
|
||||
|
||||
@@ -5,7 +5,16 @@ from threading import Lock
|
||||
from basic_memory.config import BasicMemoryConfig
|
||||
from basic_memory.repository.embedding_provider import EmbeddingProvider
|
||||
|
||||
type ProviderCacheKey = tuple[str, str, int | None, int, str | None, int | None, int | None]
|
||||
type ProviderCacheKey = tuple[
|
||||
str,
|
||||
str,
|
||||
int | None,
|
||||
int,
|
||||
int,
|
||||
str | None,
|
||||
int | None,
|
||||
int | None,
|
||||
]
|
||||
|
||||
_EMBEDDING_PROVIDER_CACHE: dict[ProviderCacheKey, EmbeddingProvider] = {}
|
||||
_EMBEDDING_PROVIDER_CACHE_LOCK = Lock()
|
||||
@@ -18,6 +27,7 @@ def _provider_cache_key(app_config: BasicMemoryConfig) -> ProviderCacheKey:
|
||||
app_config.semantic_embedding_model,
|
||||
app_config.semantic_embedding_dimensions,
|
||||
app_config.semantic_embedding_batch_size,
|
||||
app_config.semantic_embedding_request_concurrency,
|
||||
app_config.semantic_embedding_cache_dir,
|
||||
app_config.semantic_embedding_threads,
|
||||
app_config.semantic_embedding_parallel,
|
||||
@@ -73,6 +83,7 @@ def create_embedding_provider(app_config: BasicMemoryConfig) -> EmbeddingProvide
|
||||
provider = OpenAIEmbeddingProvider(
|
||||
model_name=model_name,
|
||||
batch_size=app_config.semantic_embedding_batch_size,
|
||||
request_concurrency=app_config.semantic_embedding_request_concurrency,
|
||||
**extra_kwargs,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -18,6 +18,7 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
|
||||
model_name: str = "text-embedding-3-small",
|
||||
*,
|
||||
batch_size: int = 64,
|
||||
request_concurrency: int = 4,
|
||||
dimensions: int = 1536,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
@@ -26,6 +27,7 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
|
||||
self.model_name = model_name
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.request_concurrency = request_concurrency
|
||||
self._api_key = api_key
|
||||
self._base_url = base_url
|
||||
self._timeout = timeout
|
||||
@@ -67,25 +69,49 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
|
||||
return []
|
||||
|
||||
client = await self._get_client()
|
||||
all_vectors: list[list[float]] = []
|
||||
batches = [
|
||||
texts[start : start + self.batch_size]
|
||||
for start in range(0, len(texts), self.batch_size)
|
||||
]
|
||||
batch_vectors: list[list[list[float]] | None] = [None] * len(batches)
|
||||
semaphore = asyncio.Semaphore(self.request_concurrency)
|
||||
|
||||
for start in range(0, len(texts), self.batch_size):
|
||||
batch = texts[start : start + self.batch_size]
|
||||
response = await client.embeddings.create(
|
||||
model=self.model_name,
|
||||
input=batch,
|
||||
)
|
||||
vectors_by_index: dict[int, list[float]] = {
|
||||
int(item.index): [float(value) for value in item.embedding]
|
||||
for item in response.data
|
||||
}
|
||||
async def embed_batch(batch_index: int, batch: list[str]) -> None:
|
||||
async with semaphore:
|
||||
response = await client.embeddings.create(
|
||||
model=self.model_name,
|
||||
input=batch,
|
||||
)
|
||||
|
||||
vectors_by_index: dict[int, list[float]] = {}
|
||||
for item in response.data:
|
||||
response_index = int(item.index)
|
||||
if response_index in vectors_by_index:
|
||||
raise RuntimeError(
|
||||
"OpenAI embedding response returned duplicate vector indexes."
|
||||
)
|
||||
vectors_by_index[response_index] = [float(value) for value in item.embedding]
|
||||
|
||||
ordered_vectors: list[list[float]] = []
|
||||
for index in range(len(batch)):
|
||||
vector = vectors_by_index.get(index)
|
||||
if vector is None:
|
||||
raise RuntimeError(
|
||||
"OpenAI embedding response is missing expected vector index."
|
||||
)
|
||||
all_vectors.append(vector)
|
||||
ordered_vectors.append(vector)
|
||||
|
||||
batch_vectors[batch_index] = ordered_vectors
|
||||
|
||||
await asyncio.gather(
|
||||
*(embed_batch(batch_index, batch) for batch_index, batch in enumerate(batches))
|
||||
)
|
||||
|
||||
all_vectors: list[list[float]] = []
|
||||
for vectors in batch_vectors:
|
||||
if vectors is None:
|
||||
raise RuntimeError("OpenAI embedding batch did not produce vectors.")
|
||||
all_vectors.extend(vectors)
|
||||
|
||||
if all_vectors and len(all_vectors[0]) != self.dimensions:
|
||||
raise RuntimeError(
|
||||
|
||||
@@ -475,12 +475,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
|
||||
"""Prepare chunk mutations with Postgres-specific bulk upserts."""
|
||||
sync_start = time.perf_counter()
|
||||
|
||||
logger.info(
|
||||
"Vector sync start: project_id={project_id} entity_id={entity_id}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
)
|
||||
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
await self._prepare_vector_session(session)
|
||||
|
||||
@@ -509,13 +503,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
|
||||
source_rows_count = len(rows)
|
||||
|
||||
if not rows:
|
||||
logger.info(
|
||||
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
|
||||
"source_rows_count={source_rows_count} built_chunk_records_count=0",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
source_rows_count=source_rows_count,
|
||||
)
|
||||
await self._delete_entity_chunks(session, entity_id)
|
||||
await session.commit()
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
@@ -531,15 +518,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
|
||||
built_chunk_records_count = len(chunk_records)
|
||||
current_entity_fingerprint = self._build_entity_fingerprint(chunk_records)
|
||||
current_embedding_model = self._embedding_model_key()
|
||||
logger.info(
|
||||
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
|
||||
"source_rows_count={source_rows_count} "
|
||||
"built_chunk_records_count={built_chunk_records_count}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
source_rows_count=source_rows_count,
|
||||
built_chunk_records_count=built_chunk_records_count,
|
||||
)
|
||||
if not chunk_records:
|
||||
await self._delete_entity_chunks(session, entity_id)
|
||||
await session.commit()
|
||||
@@ -592,16 +570,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
|
||||
)
|
||||
)
|
||||
if skip_unchanged_entity:
|
||||
logger.info(
|
||||
"Vector sync skipped unchanged entity: project_id={project_id} "
|
||||
"entity_id={entity_id} chunks_skipped={chunks_skipped} "
|
||||
"entity_fingerprint={entity_fingerprint} embedding_model={embedding_model}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
chunks_skipped=built_chunk_records_count,
|
||||
entity_fingerprint=current_entity_fingerprint,
|
||||
embedding_model=current_embedding_model,
|
||||
)
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
return _PreparedEntityVectorSync(
|
||||
entity_id=entity_id,
|
||||
@@ -715,31 +683,6 @@ class PostgresSearchRepository(SearchRepositoryBase):
|
||||
row_id = upserted_ids_by_key[record["chunk_key"]]
|
||||
embedding_jobs.append((row_id, record["chunk_text"]))
|
||||
|
||||
logger.info(
|
||||
"Vector sync diff complete: project_id={project_id} entity_id={entity_id} "
|
||||
"existing_chunks_count={existing_chunks_count} "
|
||||
"stale_chunks_count={stale_chunks_count} "
|
||||
"orphan_chunks_count={orphan_chunks_count} "
|
||||
"chunks_skipped={chunks_skipped} "
|
||||
"embedding_jobs_count={embedding_jobs_count} "
|
||||
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
|
||||
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard} "
|
||||
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
existing_chunks_count=existing_chunks_count,
|
||||
stale_chunks_count=stale_chunks_count,
|
||||
orphan_chunks_count=orphan_chunks_count,
|
||||
chunks_skipped=skipped_chunks_count,
|
||||
embedding_jobs_count=len(embedding_jobs),
|
||||
pending_jobs_total=shard_plan.pending_jobs_total,
|
||||
shard_index=shard_plan.shard_index,
|
||||
shard_count=shard_plan.shard_count,
|
||||
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
|
||||
oversized_entity=shard_plan.oversized_entity,
|
||||
entity_complete=shard_plan.entity_complete,
|
||||
)
|
||||
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
return _PreparedEntityVectorSync(
|
||||
entity_id=entity_id,
|
||||
|
||||
@@ -590,7 +590,6 @@ class SearchRepositoryBase(ABC):
|
||||
if shard_plan.pending_jobs_total == 0:
|
||||
return
|
||||
|
||||
scheduled_jobs_count = shard_plan.pending_jobs_total - shard_plan.remaining_jobs_after_shard
|
||||
if shard_plan.oversized_entity:
|
||||
logger.warning(
|
||||
"Vector sync oversized entity detected: project_id={project_id} "
|
||||
@@ -603,23 +602,6 @@ class SearchRepositoryBase(ABC):
|
||||
shard_count=shard_plan.shard_count,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Vector sync shard planned: project_id={project_id} entity_id={entity_id} "
|
||||
"pending_jobs_total={pending_jobs_total} scheduled_jobs_count={scheduled_jobs_count} "
|
||||
"shard_index={shard_index} shard_count={shard_count} "
|
||||
"remaining_jobs_after_shard={remaining_jobs_after_shard} "
|
||||
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
pending_jobs_total=shard_plan.pending_jobs_total,
|
||||
scheduled_jobs_count=scheduled_jobs_count,
|
||||
shard_index=shard_plan.shard_index,
|
||||
shard_count=shard_plan.shard_count,
|
||||
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
|
||||
oversized_entity=shard_plan.oversized_entity,
|
||||
entity_complete=shard_plan.entity_complete,
|
||||
)
|
||||
|
||||
# --- Text splitting ---
|
||||
|
||||
def _split_text_into_chunks(self, text_value: str) -> list[str]:
|
||||
@@ -1044,12 +1026,6 @@ class SearchRepositoryBase(ABC):
|
||||
"""Prepare chunk mutations and embedding jobs for one entity."""
|
||||
sync_start = time.perf_counter()
|
||||
|
||||
logger.info(
|
||||
"Vector sync start: project_id={project_id} entity_id={entity_id}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
)
|
||||
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
await self._prepare_vector_session(session)
|
||||
|
||||
@@ -1080,15 +1056,6 @@ class SearchRepositoryBase(ABC):
|
||||
|
||||
# No search_index rows → delete all chunk/embedding data for this entity.
|
||||
if not rows:
|
||||
logger.info(
|
||||
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
|
||||
"source_rows_count={source_rows_count} "
|
||||
"built_chunk_records_count={built_chunk_records_count}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
source_rows_count=source_rows_count,
|
||||
built_chunk_records_count=built_chunk_records_count,
|
||||
)
|
||||
await self._delete_entity_chunks(session, entity_id)
|
||||
await session.commit()
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
@@ -1104,15 +1071,6 @@ class SearchRepositoryBase(ABC):
|
||||
built_chunk_records_count = len(chunk_records)
|
||||
current_entity_fingerprint = self._build_entity_fingerprint(chunk_records)
|
||||
current_embedding_model = self._embedding_model_key()
|
||||
logger.info(
|
||||
"Vector sync source prepared: project_id={project_id} entity_id={entity_id} "
|
||||
"source_rows_count={source_rows_count} "
|
||||
"built_chunk_records_count={built_chunk_records_count}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
source_rows_count=source_rows_count,
|
||||
built_chunk_records_count=built_chunk_records_count,
|
||||
)
|
||||
if not chunk_records:
|
||||
await self._delete_entity_chunks(session, entity_id)
|
||||
await session.commit()
|
||||
@@ -1178,16 +1136,6 @@ class SearchRepositoryBase(ABC):
|
||||
)
|
||||
)
|
||||
if skip_unchanged_entity:
|
||||
logger.info(
|
||||
"Vector sync skipped unchanged entity: project_id={project_id} "
|
||||
"entity_id={entity_id} chunks_skipped={chunks_skipped} "
|
||||
"entity_fingerprint={entity_fingerprint} embedding_model={embedding_model}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
chunks_skipped=built_chunk_records_count,
|
||||
entity_fingerprint=current_entity_fingerprint,
|
||||
embedding_model=current_embedding_model,
|
||||
)
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
return _PreparedEntityVectorSync(
|
||||
entity_id=entity_id,
|
||||
@@ -1305,31 +1253,6 @@ class SearchRepositoryBase(ABC):
|
||||
)
|
||||
row_id = int(inserted.scalar_one())
|
||||
embedding_jobs.append((row_id, record["chunk_text"]))
|
||||
|
||||
logger.info(
|
||||
"Vector sync diff complete: project_id={project_id} entity_id={entity_id} "
|
||||
"existing_chunks_count={existing_chunks_count} "
|
||||
"stale_chunks_count={stale_chunks_count} "
|
||||
"orphan_chunks_count={orphan_chunks_count} "
|
||||
"chunks_skipped={chunks_skipped} "
|
||||
"embedding_jobs_count={embedding_jobs_count} "
|
||||
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
|
||||
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard} "
|
||||
"oversized_entity={oversized_entity} entity_complete={entity_complete}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
existing_chunks_count=existing_chunks_count,
|
||||
stale_chunks_count=stale_chunks_count,
|
||||
orphan_chunks_count=orphan_chunks_count,
|
||||
chunks_skipped=skipped_chunks_count,
|
||||
embedding_jobs_count=len(embedding_jobs),
|
||||
pending_jobs_total=shard_plan.pending_jobs_total,
|
||||
shard_index=shard_plan.shard_index,
|
||||
shard_count=shard_plan.shard_count,
|
||||
remaining_jobs_after_shard=shard_plan.remaining_jobs_after_shard,
|
||||
oversized_entity=shard_plan.oversized_entity,
|
||||
entity_complete=shard_plan.entity_complete,
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
prepare_seconds = time.perf_counter() - sync_start
|
||||
@@ -1364,15 +1287,6 @@ class SearchRepositoryBase(ABC):
|
||||
texts = [job.chunk_text for job in flush_jobs]
|
||||
embeddings = await self._embedding_provider.embed_documents(texts)
|
||||
embed_seconds = time.perf_counter() - embed_start
|
||||
embed_rate = (len(flush_jobs) / embed_seconds) if embed_seconds > 0 else 0.0
|
||||
logger.info(
|
||||
"Vector batch embed flush: project_id={project_id} chunk_count={chunk_count} "
|
||||
"embed_seconds={embed_seconds:.3f} embed_rate_chunks_per_second={embed_rate:.2f}",
|
||||
project_id=self.project_id,
|
||||
chunk_count=len(flush_jobs),
|
||||
embed_seconds=embed_seconds,
|
||||
embed_rate=embed_rate,
|
||||
)
|
||||
if len(embeddings) != len(flush_jobs):
|
||||
raise RuntimeError("Embedding provider returned an unexpected number of vectors.")
|
||||
|
||||
@@ -1383,15 +1297,6 @@ class SearchRepositoryBase(ABC):
|
||||
await self._write_embeddings(session, write_jobs, embeddings)
|
||||
await session.commit()
|
||||
write_seconds = time.perf_counter() - write_start
|
||||
write_rate = (len(flush_jobs) / write_seconds) if write_seconds > 0 else 0.0
|
||||
logger.info(
|
||||
"Vector batch write flush: project_id={project_id} row_count={row_count} "
|
||||
"write_seconds={write_seconds:.3f} write_rate_rows_per_second={write_rate:.2f}",
|
||||
project_id=self.project_id,
|
||||
row_count=len(flush_jobs),
|
||||
write_seconds=write_seconds,
|
||||
write_rate=write_rate,
|
||||
)
|
||||
|
||||
flush_size = len(flush_jobs)
|
||||
entity_job_counts: dict[int, int] = {}
|
||||
@@ -1485,35 +1390,6 @@ class SearchRepositoryBase(ABC):
|
||||
remaining_jobs_after_shard: int,
|
||||
) -> None:
|
||||
"""Log completion and slow-entity warnings with a consistent format."""
|
||||
logger.info(
|
||||
"Vector sync complete: project_id={project_id} entity_id={entity_id} "
|
||||
"total_seconds={total_seconds:.3f} prepare_seconds={prepare_seconds:.3f} "
|
||||
"queue_wait_seconds={queue_wait_seconds:.3f} embed_seconds={embed_seconds:.3f} "
|
||||
"write_seconds={write_seconds:.3f} source_rows_count={source_rows_count} "
|
||||
"chunks_total={chunks_total} chunks_skipped={chunks_skipped} "
|
||||
"embedding_jobs_count={embedding_jobs_count} entity_skipped={entity_skipped} "
|
||||
"entity_complete={entity_complete} oversized_entity={oversized_entity} "
|
||||
"pending_jobs_total={pending_jobs_total} shard_index={shard_index} "
|
||||
"shard_count={shard_count} remaining_jobs_after_shard={remaining_jobs_after_shard}",
|
||||
project_id=self.project_id,
|
||||
entity_id=entity_id,
|
||||
total_seconds=total_seconds,
|
||||
prepare_seconds=prepare_seconds,
|
||||
queue_wait_seconds=queue_wait_seconds,
|
||||
embed_seconds=embed_seconds,
|
||||
write_seconds=write_seconds,
|
||||
source_rows_count=source_rows_count,
|
||||
chunks_total=chunks_total,
|
||||
chunks_skipped=chunks_skipped,
|
||||
embedding_jobs_count=embedding_jobs_count,
|
||||
entity_skipped=entity_skipped,
|
||||
entity_complete=entity_complete,
|
||||
oversized_entity=oversized_entity,
|
||||
pending_jobs_total=pending_jobs_total,
|
||||
shard_index=shard_index,
|
||||
shard_count=shard_count,
|
||||
remaining_jobs_after_shard=remaining_jobs_after_shard,
|
||||
)
|
||||
if total_seconds > 10:
|
||||
logger.warning(
|
||||
"Vector sync slow entity: project_id={project_id} entity_id={entity_id} "
|
||||
@@ -1719,22 +1595,6 @@ class SearchRepositoryBase(ABC):
|
||||
return
|
||||
|
||||
total_ms = (time.perf_counter() - query_start) * 1000
|
||||
logger.info(
|
||||
"Semantic query timing: project_id={project_id} retrieval_mode={retrieval_mode} "
|
||||
"query_length={query_length} candidate_limit={candidate_limit} "
|
||||
"vector_row_count={vector_row_count} embed_ms={embed_ms:.2f} "
|
||||
"vector_query_ms={vector_query_ms:.2f} hydrate_ms={hydrate_ms:.2f} "
|
||||
"total_ms={total_ms:.2f}",
|
||||
project_id=self.project_id,
|
||||
retrieval_mode="vector",
|
||||
query_length=len(query_text),
|
||||
candidate_limit=candidate_limit,
|
||||
vector_row_count=vector_row_count,
|
||||
embed_ms=embed_ms,
|
||||
vector_query_ms=vector_query_ms,
|
||||
hydrate_ms=hydrate_ms,
|
||||
total_ms=total_ms,
|
||||
)
|
||||
if total_ms > 2000:
|
||||
logger.warning(
|
||||
"[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} "
|
||||
@@ -2073,22 +1933,6 @@ class SearchRepositoryBase(ABC):
|
||||
output.append(replace(row, score=fused_score))
|
||||
fusion_ms = (time.perf_counter() - fusion_start) * 1000
|
||||
total_ms = (time.perf_counter() - query_start) * 1000
|
||||
logger.info(
|
||||
"Semantic query timing: project_id={project_id} retrieval_mode={retrieval_mode} "
|
||||
"query_length={query_length} candidate_limit={candidate_limit} "
|
||||
"fts_count={fts_count} vector_count={vector_count} fts_ms={fts_ms:.2f} "
|
||||
"vector_ms={vector_ms:.2f} fusion_ms={fusion_ms:.2f} total_ms={total_ms:.2f}",
|
||||
project_id=self.project_id,
|
||||
retrieval_mode="hybrid",
|
||||
query_length=len(query_text),
|
||||
candidate_limit=candidate_limit,
|
||||
fts_count=len(fts_results),
|
||||
vector_count=len(vector_results),
|
||||
fts_ms=fts_ms,
|
||||
vector_ms=vector_ms,
|
||||
fusion_ms=fusion_ms,
|
||||
total_ms=total_ms,
|
||||
)
|
||||
if total_ms > 2500:
|
||||
logger.warning(
|
||||
"[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} "
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import mimetypes
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
|
||||
@@ -25,6 +26,14 @@ from basic_memory.utils import FilePath
|
||||
from loguru import logger
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FrontmatterUpdateResult:
|
||||
"""Final content emitted by a frontmatter rewrite without a follow-up reread."""
|
||||
|
||||
checksum: str
|
||||
content: str
|
||||
|
||||
|
||||
class FileService:
|
||||
"""Service for handling file operations with concurrency control.
|
||||
|
||||
@@ -301,7 +310,7 @@ class FileService:
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("File read error", path=str(full_path), error=str(e))
|
||||
raise FileOperationError(f"Failed to read file: {e}")
|
||||
raise FileOperationError(f"Failed to read file: {e}") from e
|
||||
|
||||
async def read_file(self, path: FilePath) -> Tuple[str, str]:
|
||||
"""Read file and compute checksum using true async I/O.
|
||||
@@ -401,12 +410,14 @@ class FileService:
|
||||
)
|
||||
raise FileOperationError(f"Failed to move file {source} -> {destination}: {e}")
|
||||
|
||||
async def update_frontmatter(self, path: FilePath, updates: Dict[str, Any]) -> str:
|
||||
"""Update frontmatter fields in a file while preserving all content.
|
||||
async def update_frontmatter_with_result(
|
||||
self, path: FilePath, updates: Dict[str, Any]
|
||||
) -> FrontmatterUpdateResult:
|
||||
"""Update frontmatter and return the exact final written markdown content.
|
||||
|
||||
Only modifies the frontmatter section, leaving all content untouched.
|
||||
Creates frontmatter section if none exists.
|
||||
Returns checksum of updated file.
|
||||
Returns both checksum and final content so callers do not need a reread.
|
||||
|
||||
Uses aiofiles for true async I/O (non-blocking).
|
||||
|
||||
@@ -415,7 +426,7 @@ class FileService:
|
||||
updates: Dict of frontmatter fields to update
|
||||
|
||||
Returns:
|
||||
Checksum of updated file
|
||||
Typed result containing checksum and final content
|
||||
|
||||
Raises:
|
||||
FileOperationError: If file operations fail
|
||||
@@ -467,7 +478,10 @@ class FileService:
|
||||
if formatted_content is not None:
|
||||
content_for_checksum = formatted_content # pragma: no cover
|
||||
|
||||
return await file_utils.compute_checksum(content_for_checksum)
|
||||
return FrontmatterUpdateResult(
|
||||
checksum=await file_utils.compute_checksum(content_for_checksum),
|
||||
content=content_for_checksum,
|
||||
)
|
||||
|
||||
except Exception as e: # pragma: no cover
|
||||
# Only log real errors (not YAML parsing, which is handled above)
|
||||
@@ -479,6 +493,11 @@ class FileService:
|
||||
)
|
||||
raise FileOperationError(f"Failed to update frontmatter: {e}")
|
||||
|
||||
async def update_frontmatter(self, path: FilePath, updates: Dict[str, Any]) -> str:
|
||||
"""Update frontmatter fields in a file while preserving all content."""
|
||||
result = await self.update_frontmatter_with_result(path, updates)
|
||||
return result.checksum
|
||||
|
||||
async def compute_checksum(self, path: FilePath) -> str:
|
||||
"""Compute checksum for a file using true async I/O.
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Service for search operations."""
|
||||
|
||||
import asyncio
|
||||
import ast
|
||||
import re
|
||||
from datetime import datetime
|
||||
@@ -427,6 +428,15 @@ class SearchService:
|
||||
|
||||
async def sync_entity_vectors(self, entity_id: int) -> None:
|
||||
"""Refresh vector chunks for one entity in repositories that support semantic indexing."""
|
||||
entity = await self.entity_repository.find_by_id(entity_id)
|
||||
if entity is None:
|
||||
await self._clear_entity_vectors(entity_id)
|
||||
return
|
||||
|
||||
if not self._entity_embeddings_enabled(entity):
|
||||
await self._clear_entity_vectors(entity_id)
|
||||
return
|
||||
|
||||
await self.repository.sync_entity_vectors(entity_id)
|
||||
|
||||
async def sync_entity_vectors_batch(
|
||||
@@ -435,10 +445,90 @@ class SearchService:
|
||||
progress_callback=None,
|
||||
) -> VectorSyncBatchResult:
|
||||
"""Refresh vector chunks for a batch of entities."""
|
||||
return await self.repository.sync_entity_vectors_batch(
|
||||
entity_ids,
|
||||
progress_callback=progress_callback,
|
||||
if not entity_ids:
|
||||
return VectorSyncBatchResult(
|
||||
entities_total=0,
|
||||
entities_synced=0,
|
||||
entities_failed=0,
|
||||
)
|
||||
|
||||
entities_by_id = {
|
||||
entity.id: entity for entity in await self.entity_repository.find_by_ids(entity_ids)
|
||||
}
|
||||
unknown_ids = [entity_id for entity_id in entity_ids if entity_id not in entities_by_id]
|
||||
opted_out_ids = [
|
||||
entity_id
|
||||
for entity_id in entity_ids
|
||||
if (
|
||||
(entity := entities_by_id.get(entity_id)) is not None
|
||||
and not self._entity_embeddings_enabled(entity)
|
||||
)
|
||||
]
|
||||
if opted_out_ids:
|
||||
await asyncio.gather(
|
||||
*(self._clear_entity_vectors(entity_id) for entity_id in opted_out_ids)
|
||||
)
|
||||
|
||||
eligible_entity_ids = [
|
||||
entity_id
|
||||
for entity_id in entity_ids
|
||||
if entity_id in entities_by_id and entity_id not in opted_out_ids
|
||||
]
|
||||
|
||||
cleanup_task = (
|
||||
self.repository.sync_entity_vectors_batch(unknown_ids) if unknown_ids else None
|
||||
)
|
||||
eligible_task = (
|
||||
self.repository.sync_entity_vectors_batch(
|
||||
eligible_entity_ids,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
if eligible_entity_ids
|
||||
else None
|
||||
)
|
||||
repository_results = [
|
||||
result
|
||||
for result in await asyncio.gather(
|
||||
cleanup_task if cleanup_task is not None else asyncio.sleep(0, result=None),
|
||||
eligible_task if eligible_task is not None else asyncio.sleep(0, result=None),
|
||||
)
|
||||
if result is not None
|
||||
]
|
||||
|
||||
if not repository_results:
|
||||
return VectorSyncBatchResult(
|
||||
entities_total=len(entity_ids),
|
||||
entities_synced=0,
|
||||
entities_failed=0,
|
||||
entities_skipped=len(opted_out_ids),
|
||||
)
|
||||
|
||||
batch_result = VectorSyncBatchResult(
|
||||
entities_total=len(entity_ids),
|
||||
entities_synced=sum(result.entities_synced for result in repository_results),
|
||||
entities_failed=sum(result.entities_failed for result in repository_results),
|
||||
entities_deferred=sum(result.entities_deferred for result in repository_results),
|
||||
entities_skipped=(
|
||||
len(opted_out_ids)
|
||||
+ sum(result.entities_skipped for result in repository_results)
|
||||
- len(unknown_ids)
|
||||
),
|
||||
failed_entity_ids=[
|
||||
failed_entity_id
|
||||
for result in repository_results
|
||||
for failed_entity_id in result.failed_entity_ids
|
||||
],
|
||||
chunks_total=sum(result.chunks_total for result in repository_results),
|
||||
chunks_skipped=sum(result.chunks_skipped for result in repository_results),
|
||||
embedding_jobs_total=sum(result.embedding_jobs_total for result in repository_results),
|
||||
prepare_seconds_total=sum(result.prepare_seconds_total for result in repository_results),
|
||||
queue_wait_seconds_total=sum(
|
||||
result.queue_wait_seconds_total for result in repository_results
|
||||
),
|
||||
embed_seconds_total=sum(result.embed_seconds_total for result in repository_results),
|
||||
write_seconds_total=sum(result.write_seconds_total for result in repository_results),
|
||||
)
|
||||
return batch_result
|
||||
|
||||
async def reindex_vectors(self, progress_callback=None) -> dict:
|
||||
"""Rebuild vector embeddings for all entities.
|
||||
@@ -456,14 +546,14 @@ class SearchService:
|
||||
# that reference entity_ids no longer in the entity table
|
||||
await self._purge_stale_search_rows()
|
||||
|
||||
batch_result = await self.repository.sync_entity_vectors_batch(
|
||||
batch_result = await self.sync_entity_vectors_batch(
|
||||
entity_ids,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
stats = {
|
||||
"total_entities": batch_result.entities_total,
|
||||
"embedded": batch_result.entities_synced,
|
||||
"skipped": 0,
|
||||
"skipped": batch_result.entities_skipped,
|
||||
"errors": batch_result.entities_failed,
|
||||
}
|
||||
|
||||
@@ -518,6 +608,60 @@ class SearchService:
|
||||
|
||||
logger.info("Purged stale search rows for deleted entities", project_id=project_id)
|
||||
|
||||
@staticmethod
|
||||
def _entity_embeddings_enabled(entity: Entity) -> bool:
|
||||
"""Return whether semantic embeddings should be generated for this entity."""
|
||||
if not entity.entity_metadata:
|
||||
return True
|
||||
|
||||
embed_value = entity.entity_metadata.get("embed")
|
||||
if embed_value is None:
|
||||
return True
|
||||
if isinstance(embed_value, bool):
|
||||
return embed_value
|
||||
if isinstance(embed_value, str):
|
||||
normalized = embed_value.strip().lower()
|
||||
if normalized in {"false", "0", "no", "off"}:
|
||||
return False
|
||||
if normalized in {"true", "1", "yes", "on"}:
|
||||
return True
|
||||
if isinstance(embed_value, (int, float)):
|
||||
return bool(embed_value)
|
||||
|
||||
# Default unknown values to enabled so malformed metadata does not silently
|
||||
# remove notes from semantic search.
|
||||
return True
|
||||
|
||||
async def _clear_entity_vectors(self, entity_id: int) -> None:
|
||||
"""Delete derived vector rows for one entity."""
|
||||
from basic_memory.repository.search_repository_base import SearchRepositoryBase
|
||||
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
|
||||
|
||||
# Trigger: semantic indexing is disabled for this repository instance.
|
||||
# Why: repositories only create vector tables when semantic search is enabled.
|
||||
# Outcome: skip cleanup because there are no active derived vector rows to maintain.
|
||||
if isinstance(self.repository, SearchRepositoryBase) and not self.repository._semantic_enabled:
|
||||
return
|
||||
|
||||
params = {"project_id": self.repository.project_id, "entity_id": entity_id}
|
||||
if isinstance(self.repository, SQLiteSearchRepository):
|
||||
await self.repository.execute_query(
|
||||
text(
|
||||
"DELETE FROM search_vector_embeddings WHERE rowid IN ("
|
||||
"SELECT id FROM search_vector_chunks "
|
||||
"WHERE project_id = :project_id AND entity_id = :entity_id)"
|
||||
),
|
||||
params,
|
||||
)
|
||||
|
||||
await self.repository.execute_query(
|
||||
text(
|
||||
"DELETE FROM search_vector_chunks "
|
||||
"WHERE project_id = :project_id AND entity_id = :entity_id"
|
||||
),
|
||||
params,
|
||||
)
|
||||
|
||||
async def index_entity_file(
|
||||
self,
|
||||
entity: Entity,
|
||||
|
||||
@@ -8,7 +8,7 @@ from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import AsyncIterator, Dict, List, Optional, Set, Tuple
|
||||
from typing import AsyncIterator, Awaitable, Callable, Dict, List, Optional, Set, Tuple
|
||||
|
||||
import aiofiles.os
|
||||
|
||||
@@ -18,7 +18,15 @@ from sqlalchemy.exc import IntegrityError
|
||||
from basic_memory import telemetry
|
||||
from basic_memory import db
|
||||
from basic_memory.config import BasicMemoryConfig, ConfigManager
|
||||
from basic_memory.file_utils import has_frontmatter
|
||||
from basic_memory.file_utils import compute_checksum, has_frontmatter
|
||||
from basic_memory.indexing import BatchIndexer, IndexFileMetadata, IndexInputFile, IndexProgress
|
||||
from basic_memory.indexing.batching import build_index_batches
|
||||
from basic_memory.indexing.models import (
|
||||
IndexedEntity,
|
||||
IndexFileWriter,
|
||||
IndexFrontmatterUpdate,
|
||||
IndexFrontmatterWriteResult,
|
||||
)
|
||||
from basic_memory.ignore_utils import load_bmignore_patterns, should_ignore_path
|
||||
from basic_memory.markdown import EntityParser, MarkdownProcessor
|
||||
from basic_memory.models import Entity, Project
|
||||
@@ -31,7 +39,7 @@ from basic_memory.repository import (
|
||||
from basic_memory.repository.search_repository import create_search_repository
|
||||
from basic_memory.services import EntityService, FileService
|
||||
from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
|
||||
from basic_memory.services.exceptions import SyncFatalError
|
||||
from basic_memory.services.exceptions import FileOperationError, SyncFatalError
|
||||
from basic_memory.services.link_resolver import LinkResolver
|
||||
from basic_memory.services.search_service import SearchService
|
||||
|
||||
@@ -118,6 +126,23 @@ class ScanResult:
|
||||
errors: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
class _FileServiceIndexWriter(IndexFileWriter):
|
||||
"""Adapt FileService frontmatter updates to the indexing writer protocol."""
|
||||
|
||||
def __init__(self, file_service: FileService) -> None:
|
||||
self.file_service = file_service
|
||||
|
||||
async def write_frontmatter(
|
||||
self, update: IndexFrontmatterUpdate
|
||||
) -> IndexFrontmatterWriteResult:
|
||||
# Why: IndexFrontmatterWriteResult lives in indexing/models.py so the indexing
|
||||
# layer does not need to import FileService. This adapter keeps that boundary intact.
|
||||
result = await self.file_service.update_frontmatter_with_result(
|
||||
update.path, update.metadata
|
||||
)
|
||||
return IndexFrontmatterWriteResult(checksum=result.checksum, content=result.content)
|
||||
|
||||
|
||||
class SyncService:
|
||||
"""Syncs documents and knowledge files with database."""
|
||||
|
||||
@@ -146,6 +171,14 @@ class SyncService:
|
||||
# Use OrderedDict for LRU behavior with bounded size to prevent unbounded memory growth
|
||||
self._file_failures: OrderedDict[str, FileFailureInfo] = OrderedDict()
|
||||
self._max_tracked_failures = 100 # Limit failure cache size
|
||||
self.batch_indexer = BatchIndexer(
|
||||
app_config=app_config,
|
||||
entity_service=entity_service,
|
||||
entity_repository=entity_repository,
|
||||
relation_repository=relation_repository,
|
||||
search_service=search_service,
|
||||
file_writer=_FileServiceIndexWriter(file_service),
|
||||
)
|
||||
|
||||
async def _should_skip_file(self, path: str) -> bool:
|
||||
"""Check if file should be skipped due to repeated failures.
|
||||
@@ -255,7 +288,11 @@ class SyncService:
|
||||
del self._file_failures[path]
|
||||
|
||||
async def sync(
|
||||
self, directory: Path, project_name: Optional[str] = None, force_full: bool = False
|
||||
self,
|
||||
directory: Path,
|
||||
project_name: Optional[str] = None,
|
||||
force_full: bool = False,
|
||||
progress_callback: Callable[[IndexProgress], Awaitable[None]] | None = None,
|
||||
) -> SyncReport:
|
||||
"""Sync all files with database and update scan watermark.
|
||||
|
||||
@@ -263,6 +300,7 @@ class SyncService:
|
||||
directory: Directory to sync
|
||||
project_name: Optional project name
|
||||
force_full: If True, force a full scan bypassing watermark optimization
|
||||
progress_callback: Optional callback for typed indexing progress updates
|
||||
"""
|
||||
|
||||
start_time = time.time()
|
||||
@@ -310,42 +348,14 @@ class SyncService:
|
||||
for path in report.deleted:
|
||||
await self.handle_delete(path)
|
||||
|
||||
# then new and modified — collect entity IDs for batch vector embedding
|
||||
synced_entity_ids: list[int] = []
|
||||
|
||||
for path in report.new:
|
||||
entity, _ = await self.sync_file(path, new=True)
|
||||
|
||||
if entity is not None:
|
||||
synced_entity_ids.append(entity.id)
|
||||
# Track if file was skipped
|
||||
elif await self._should_skip_file(path):
|
||||
failure_info = self._file_failures[path]
|
||||
report.skipped_files.append(
|
||||
SkippedFile(
|
||||
path=path,
|
||||
reason=failure_info.last_error,
|
||||
failure_count=failure_info.count,
|
||||
first_failed=failure_info.first_failure,
|
||||
)
|
||||
)
|
||||
|
||||
for path in report.modified:
|
||||
entity, _ = await self.sync_file(path, new=False)
|
||||
|
||||
if entity is not None:
|
||||
synced_entity_ids.append(entity.id)
|
||||
# Track if file was skipped
|
||||
elif await self._should_skip_file(path):
|
||||
failure_info = self._file_failures[path]
|
||||
report.skipped_files.append(
|
||||
SkippedFile(
|
||||
path=path,
|
||||
reason=failure_info.last_error,
|
||||
failure_count=failure_info.count,
|
||||
first_failed=failure_info.first_failure,
|
||||
)
|
||||
)
|
||||
changed_paths = sorted(report.new | report.modified)
|
||||
indexed_entities, skipped_files = await self._index_changed_files(
|
||||
changed_paths,
|
||||
report.checksums,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
report.skipped_files.extend(skipped_files)
|
||||
synced_entity_ids = [indexed.entity_id for indexed in indexed_entities]
|
||||
|
||||
# Only resolve relations if there were actual changes
|
||||
# If no files changed, no new unresolved relations could have been created
|
||||
@@ -353,11 +363,12 @@ class SyncService:
|
||||
with telemetry.scope(
|
||||
"sync.project.resolve_relations", relation_scope="all_pending"
|
||||
):
|
||||
await self.resolve_relations()
|
||||
synced_entity_ids.extend(await self.resolve_relations())
|
||||
else:
|
||||
logger.info("Skipping relation resolution - no file changes detected")
|
||||
|
||||
# Batch-generate vector embeddings for all synced entities
|
||||
synced_entity_ids = list(dict.fromkeys(synced_entity_ids))
|
||||
if synced_entity_ids and self.app_config.semantic_search_enabled:
|
||||
try:
|
||||
with telemetry.scope(
|
||||
@@ -425,6 +436,277 @@ class SyncService:
|
||||
|
||||
return report
|
||||
|
||||
async def _index_changed_files(
|
||||
self,
|
||||
changed_paths: list[str],
|
||||
checksums_by_path: dict[str, str],
|
||||
*,
|
||||
progress_callback: Callable[[IndexProgress], Awaitable[None]] | None = None,
|
||||
) -> tuple[list[IndexedEntity], list[SkippedFile]]:
|
||||
"""Load, batch, and index changed files without processing them serially."""
|
||||
if not changed_paths:
|
||||
if progress_callback is not None:
|
||||
await progress_callback(
|
||||
IndexProgress(
|
||||
files_total=0,
|
||||
files_processed=0,
|
||||
batches_total=0,
|
||||
batches_completed=0,
|
||||
)
|
||||
)
|
||||
return [], []
|
||||
|
||||
started_at = time.monotonic()
|
||||
files_total = len(changed_paths)
|
||||
files_processed = 0
|
||||
batches_completed = 0
|
||||
skipped_files: list[SkippedFile] = []
|
||||
skipped_paths: set[str] = set()
|
||||
candidate_paths: list[str] = []
|
||||
|
||||
# Trigger: a file exceeded the retry threshold in a previous sync.
|
||||
# Why: repeated retries on unchanged broken files waste the entire batch budget.
|
||||
# Outcome: skip it up front and still count it in progress.
|
||||
for path in changed_paths:
|
||||
if await self._should_skip_file(path):
|
||||
self._append_skipped_file(path, skipped_files, skipped_paths)
|
||||
files_processed += 1
|
||||
else:
|
||||
candidate_paths.append(path)
|
||||
|
||||
(
|
||||
metadata_by_path,
|
||||
metadata_errors,
|
||||
missing_metadata_paths,
|
||||
) = await self._load_index_file_metadata(
|
||||
candidate_paths,
|
||||
checksums_by_path,
|
||||
)
|
||||
files_processed += len(missing_metadata_paths)
|
||||
files_processed += len(metadata_errors)
|
||||
for path, error in metadata_errors:
|
||||
await self._record_index_failure(path, error, skipped_files, skipped_paths)
|
||||
|
||||
batch_paths = sorted(metadata_by_path)
|
||||
batches = build_index_batches(
|
||||
batch_paths,
|
||||
metadata_by_path,
|
||||
max_files=self.app_config.index_batch_size,
|
||||
max_bytes=self.app_config.index_batch_max_bytes,
|
||||
)
|
||||
|
||||
await self._emit_index_progress(
|
||||
progress_callback,
|
||||
files_total=files_total,
|
||||
files_processed=files_processed,
|
||||
batches_total=len(batches),
|
||||
batches_completed=0,
|
||||
current_batch_bytes=0,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
indexed_entities: list[IndexedEntity] = []
|
||||
shared_permalink_by_path: dict[str, str | None] | None = None
|
||||
if any(
|
||||
metadata.content_type == "text/markdown"
|
||||
or (
|
||||
metadata.content_type is None
|
||||
and Path(metadata.path).suffix.lower() in {".md", ".markdown"}
|
||||
)
|
||||
for metadata in metadata_by_path.values()
|
||||
):
|
||||
shared_permalink_by_path = {
|
||||
path: permalink
|
||||
for path, permalink in (
|
||||
await self.entity_repository.get_file_path_to_permalink_map()
|
||||
).items()
|
||||
}
|
||||
|
||||
for batch in batches:
|
||||
loaded_files, load_errors = await self._load_index_batch_files(
|
||||
batch.paths, metadata_by_path
|
||||
)
|
||||
for path, error in load_errors:
|
||||
await self._record_index_failure(path, error, skipped_files, skipped_paths)
|
||||
|
||||
if loaded_files:
|
||||
batch_result = await self.batch_indexer.index_files(
|
||||
loaded_files,
|
||||
max_concurrent=self.app_config.index_entity_max_concurrent,
|
||||
parse_max_concurrent=self.app_config.index_parse_max_concurrent,
|
||||
existing_permalink_by_path=shared_permalink_by_path,
|
||||
)
|
||||
indexed_entities.extend(batch_result.indexed)
|
||||
|
||||
indexed_paths = {indexed.path for indexed in batch_result.indexed}
|
||||
for path in indexed_paths:
|
||||
self._clear_failure(path)
|
||||
|
||||
for path, error in batch_result.errors:
|
||||
await self._record_index_failure(path, error, skipped_files, skipped_paths)
|
||||
|
||||
files_processed += len(batch.paths)
|
||||
batches_completed += 1
|
||||
await self._emit_index_progress(
|
||||
progress_callback,
|
||||
files_total=files_total,
|
||||
files_processed=files_processed,
|
||||
batches_total=len(batches),
|
||||
batches_completed=batches_completed,
|
||||
current_batch_bytes=batch.total_bytes,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
return indexed_entities, skipped_files
|
||||
|
||||
async def _load_index_file_metadata(
|
||||
self,
|
||||
paths: list[str],
|
||||
checksums_by_path: dict[str, str],
|
||||
) -> tuple[dict[str, IndexFileMetadata], list[tuple[str, str]], list[str]]:
|
||||
"""Load typed metadata for batch planning before any file content is read."""
|
||||
if not paths:
|
||||
return {}, [], []
|
||||
|
||||
semaphore = asyncio.Semaphore(self.app_config.sync_max_concurrent_files)
|
||||
metadata_by_path: dict[str, IndexFileMetadata] = {}
|
||||
errors: dict[str, str] = {}
|
||||
missing_paths: list[str] = []
|
||||
|
||||
async def load(path: str) -> None:
|
||||
async with semaphore:
|
||||
try:
|
||||
file_metadata = await self.file_service.get_file_metadata(path)
|
||||
metadata_by_path[path] = IndexFileMetadata(
|
||||
path=path,
|
||||
size=file_metadata.size,
|
||||
checksum=checksums_by_path.get(path),
|
||||
content_type=self.file_service.content_type(path),
|
||||
last_modified=file_metadata.modified_at,
|
||||
created_at=file_metadata.created_at,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
await self.handle_delete(path)
|
||||
missing_paths.append(path)
|
||||
except Exception as exc:
|
||||
errors[path] = str(exc)
|
||||
|
||||
await asyncio.gather(*(load(path) for path in paths))
|
||||
return (
|
||||
metadata_by_path,
|
||||
[(path, errors[path]) for path in sorted(errors)],
|
||||
sorted(missing_paths),
|
||||
)
|
||||
|
||||
async def _load_index_batch_files(
|
||||
self,
|
||||
paths: list[str],
|
||||
metadata_by_path: dict[str, IndexFileMetadata],
|
||||
) -> tuple[dict[str, IndexInputFile], list[tuple[str, str]]]:
|
||||
"""Read one batch of file contents into typed input objects."""
|
||||
if not paths:
|
||||
return {}, []
|
||||
|
||||
semaphore = asyncio.Semaphore(self.app_config.sync_max_concurrent_files)
|
||||
files: dict[str, IndexInputFile] = {}
|
||||
errors: dict[str, str] = {}
|
||||
|
||||
async def load(path: str) -> None:
|
||||
async with semaphore:
|
||||
metadata = metadata_by_path[path]
|
||||
try:
|
||||
content = await self.file_service.read_file_bytes(path)
|
||||
loaded_checksum = await compute_checksum(content)
|
||||
files[path] = IndexInputFile(
|
||||
path=metadata.path,
|
||||
size=metadata.size,
|
||||
checksum=loaded_checksum,
|
||||
content_type=metadata.content_type,
|
||||
last_modified=metadata.last_modified,
|
||||
created_at=metadata.created_at,
|
||||
content=content,
|
||||
)
|
||||
except FileOperationError as exc:
|
||||
# Trigger: FileService wraps binary read failures in FileOperationError.
|
||||
# Why: the service contract should stay consistent for direct callers.
|
||||
# Outcome: sync still treats wrapped missing-file reads as deletions.
|
||||
if isinstance(exc.__cause__, FileNotFoundError):
|
||||
await self.handle_delete(path)
|
||||
else:
|
||||
errors[path] = str(exc)
|
||||
except Exception as exc:
|
||||
errors[path] = str(exc)
|
||||
|
||||
await asyncio.gather(*(load(path) for path in paths))
|
||||
return files, [(path, errors[path]) for path in sorted(errors)]
|
||||
|
||||
async def _record_index_failure(
|
||||
self,
|
||||
path: str,
|
||||
error: str,
|
||||
skipped_files: list[SkippedFile],
|
||||
skipped_paths: set[str],
|
||||
) -> None:
|
||||
"""Record a per-file batch failure and promote it to skipped when threshold is reached."""
|
||||
await self._record_failure(path, error)
|
||||
if await self._should_skip_file(path):
|
||||
self._append_skipped_file(path, skipped_files, skipped_paths)
|
||||
|
||||
def _append_skipped_file(
|
||||
self,
|
||||
path: str,
|
||||
skipped_files: list[SkippedFile],
|
||||
skipped_paths: set[str],
|
||||
) -> None:
|
||||
"""Append one skipped file record once per sync run."""
|
||||
if path in skipped_paths or path not in self._file_failures:
|
||||
return
|
||||
|
||||
failure_info = self._file_failures[path]
|
||||
skipped_files.append(
|
||||
SkippedFile(
|
||||
path=path,
|
||||
reason=failure_info.last_error,
|
||||
failure_count=failure_info.count,
|
||||
first_failed=failure_info.first_failure,
|
||||
)
|
||||
)
|
||||
skipped_paths.add(path)
|
||||
|
||||
async def _emit_index_progress(
|
||||
self,
|
||||
progress_callback: Callable[[IndexProgress], Awaitable[None]] | None,
|
||||
*,
|
||||
files_total: int,
|
||||
files_processed: int,
|
||||
batches_total: int,
|
||||
batches_completed: int,
|
||||
current_batch_bytes: int,
|
||||
started_at: float,
|
||||
) -> None:
|
||||
"""Emit a typed indexing progress update when the caller requested one."""
|
||||
if progress_callback is None:
|
||||
return
|
||||
|
||||
elapsed_seconds = max(time.monotonic() - started_at, 0.001)
|
||||
files_per_minute = files_processed / elapsed_seconds * 60 if files_processed else 0.0
|
||||
eta_seconds = None
|
||||
if files_processed and files_total > files_processed:
|
||||
files_per_second = files_processed / elapsed_seconds
|
||||
eta_seconds = (files_total - files_processed) / files_per_second
|
||||
|
||||
await progress_callback(
|
||||
IndexProgress(
|
||||
files_total=files_total,
|
||||
files_processed=files_processed,
|
||||
batches_total=batches_total,
|
||||
batches_completed=batches_completed,
|
||||
current_batch_bytes=current_batch_bytes,
|
||||
files_per_minute=files_per_minute,
|
||||
eta_seconds=eta_seconds,
|
||||
)
|
||||
)
|
||||
|
||||
async def scan(self, directory, force_full: bool = False):
|
||||
"""Smart scan using watermark and file count for large project optimization.
|
||||
|
||||
@@ -1081,13 +1363,17 @@ class SyncService:
|
||||
# update search index
|
||||
await self.search_service.index_entity(updated)
|
||||
|
||||
async def resolve_relations(self, entity_id: int | None = None):
|
||||
async def resolve_relations(self, entity_id: int | None = None) -> set[int]:
|
||||
"""Try to resolve unresolved relations.
|
||||
|
||||
Args:
|
||||
entity_id: If provided, only resolve relations for this specific entity.
|
||||
Otherwise, resolve all unresolved relations in the database.
|
||||
|
||||
Returns:
|
||||
Set of source entity IDs whose outgoing relations changed.
|
||||
"""
|
||||
affected_entity_ids: set[int] = set()
|
||||
|
||||
if entity_id:
|
||||
# Only get unresolved relations for the specific entity
|
||||
@@ -1131,8 +1417,7 @@ class SyncService:
|
||||
"to_name": resolved_entity.title,
|
||||
},
|
||||
)
|
||||
# update search index only on successful resolution
|
||||
await self.search_service.index_entity(resolved_entity)
|
||||
affected_entity_ids.add(relation.from_id)
|
||||
except IntegrityError:
|
||||
with telemetry.scope(
|
||||
"sync.relation.resolve_conflict",
|
||||
@@ -1164,6 +1449,14 @@ class SyncService:
|
||||
logger.debug(
|
||||
f"Could not delete duplicate relation {relation.id}: {e}"
|
||||
)
|
||||
affected_entity_ids.add(relation.from_id)
|
||||
|
||||
for affected_entity_id in sorted(affected_entity_ids):
|
||||
source_entity = await self.entity_repository.find_by_id(affected_entity_id)
|
||||
if source_entity is not None:
|
||||
await self.search_service.index_entity(source_entity)
|
||||
|
||||
return affected_entity_ids
|
||||
|
||||
async def _quick_count_files(self, directory: Path) -> int:
|
||||
"""Fast file count using find command.
|
||||
|
||||
Reference in New Issue
Block a user