diff --git a/app/protocols/memory_protocol.py b/app/protocols/memory_protocol.py index 4a8bd78..95801d0 100644 --- a/app/protocols/memory_protocol.py +++ b/app/protocols/memory_protocol.py @@ -183,6 +183,65 @@ async def get_memories_for_reembedding(self, limit: int, offset: int) -> list[Me """Fetch memories in batches for re-embedding (all users, ordered by id)""" ... + async def count_memories_for_targeted_rebuild( + self, + user_id: UUID, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> int: + """Count memories eligible for a *user-scoped, targeted* embedding rebuild. + + Used by `rebuild_embeddings` to avoid the destructive global + `re_embed_all` path. Filters are user-scoped; user_id is mandatory and + scopes ownership. Matching memories are rebuilt unconditionally. + + Args: + user_id: Authenticated user id; ownership filter, never optional. + memory_ids: Restrict to explicit ids (already user-scoped server-side). + project_id: Restrict to a single project owned by `user_id`. + + Returns: + Number of non-obsolete memories matching the filters and owned by + `user_id`. + """ + ... + + async def get_memories_for_targeted_rebuild( + self, + user_id: UUID, + limit: int, + after_id: int | None = None, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> list[Memory]: + """Page through memories selected by `count_memories_for_targeted_rebuild`. + + Same semantics and same filters as the count variant. Uses keyset + pagination by memory id so callers can process deterministically. + """ + ... + + async def upsert_targeted_embeddings( + self, + user_id: UUID, + updates: list[tuple[int, list[float]]], + ) -> list[int]: + """Idempotently write a batch of new embeddings for `user_id`. + + Unlike `bulk_update_embeddings`, this method: + - is user-scoped (rejects ids not owned by `user_id`), + - never resets vector storage, + - is idempotent for repeated calls on the same memory id. + + Args: + user_id: Authenticated user id; ownership filter, never optional. + updates: list of (memory_id, embedding) tuples. + + Returns: + list of memory ids whose embedding was actually written. + """ + ... + async def reset_embedding_storage(self) -> None: """Prepare vector storage for new dimensions. SQLite: DROP + CREATE vec_memories with new dimensions diff --git a/app/repositories/postgres/memory_repository.py b/app/repositories/postgres/memory_repository.py index b21f02d..15ab19b 100644 --- a/app/repositories/postgres/memory_repository.py +++ b/app/repositories/postgres/memory_repository.py @@ -964,6 +964,115 @@ async def bulk_update_embeddings(self, updates: list[tuple[int, list[float]]]) - .values(embedding=embedding), ) + # ---- Targeted (user-scoped) rebuild support, never resets vec storage ---- + + def _build_targeted_rebuild_filter( + self, + user_id: UUID, + memory_ids: list[int] | None, + project_id: int | None, + ): + conditions = [ + MemoryTable.user_id == user_id, + MemoryTable.is_obsolete.is_(False), + ] + if memory_ids: + conditions.append(MemoryTable.id.in_(memory_ids)) + stmt = select(MemoryTable).where(*conditions) + if project_id is not None: + stmt = stmt.join(MemoryTable.projects).where( + ProjectsTable.id == project_id, + ProjectsTable.user_id == user_id, + ) + return stmt + + async def count_memories_for_targeted_rebuild( + self, + user_id: UUID, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> int: + from sqlalchemy import func + async with self.db_adapter.system_session() as session: + base = self._build_targeted_rebuild_filter( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + count_stmt = select(func.count()).select_from(base.subquery()) + return (await session.execute(count_stmt)).scalar() or 0 + + async def get_memories_for_targeted_rebuild( + self, + user_id: UUID, + limit: int, + after_id: int | None = None, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> list[Memory]: + async with self.db_adapter.system_session() as session: + stmt = self._build_targeted_rebuild_filter( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + if after_id is not None: + stmt = stmt.where(MemoryTable.id > after_id) + stmt = ( + stmt + .options( + selectinload(MemoryTable.projects), + selectinload(MemoryTable.linked_memories), + selectinload(MemoryTable.linking_memories), + selectinload(MemoryTable.code_artifacts), + selectinload(MemoryTable.documents), + selectinload(MemoryTable.files), + ) + .order_by(MemoryTable.id.asc()) + .limit(limit) + ) + result = await session.execute(stmt) + memories_orm = result.scalars().unique().all() + return [Memory.model_validate(m) for m in memories_orm] + + async def upsert_targeted_embeddings( + self, + user_id: UUID, + updates: list[tuple[int, list[float]]], + ) -> list[int]: + """Idempotently set embeddings for memories owned by `user_id`. + + Postgres stores embeddings inline on `memories.embedding`, so an UPDATE + is naturally idempotent. We still verify ownership server-side instead + of trusting the input list. + """ + if not updates: + return [] + ids = [mid for mid, _ in updates] + async with self.db_adapter.system_session() as session: + owned_rows = await session.execute( + select(MemoryTable.id).where( + MemoryTable.user_id == user_id, + MemoryTable.is_obsolete.is_(False), + MemoryTable.id.in_(ids), + ), + ) + owned_ids = {row[0] for row in owned_rows} + written: list[int] = [] + for memory_id, embedding in updates: + if memory_id not in owned_ids: + continue + await session.execute( + update(MemoryTable) + .where( + MemoryTable.id == memory_id, + MemoryTable.user_id == user_id, + ) + .values(embedding=embedding), + ) + written.append(memory_id) + return written + async def validate_embedding_count(self) -> bool: """Check embedding count matches non-obsolete memory count""" from sqlalchemy import func diff --git a/app/repositories/sqlite/memory_repository.py b/app/repositories/sqlite/memory_repository.py index c8f48e5..4c2af4c 100644 --- a/app/repositories/sqlite/memory_repository.py +++ b/app/repositories/sqlite/memory_repository.py @@ -983,6 +983,127 @@ async def bulk_update_embeddings(self, updates: list[tuple[int, list[float]]]) - {"memory_id": str(memory_id), "embedding": embedding_bytes}, ) + # ---- Targeted (user-scoped) rebuild support, never resets vec storage ---- + + def _build_targeted_rebuild_filter( + self, + user_id: UUID, + memory_ids: list[int] | None, + project_id: int | None, + ): + """Shared SQLAlchemy filter for targeted rebuild count/select queries.""" + conditions = [ + MemoryTable.user_id == str(user_id), + MemoryTable.is_obsolete.is_(False), + ] + if memory_ids: + conditions.append(MemoryTable.id.in_(memory_ids)) + stmt = select(MemoryTable).where(*conditions) + if project_id is not None: + stmt = stmt.join(MemoryTable.projects).where( + ProjectsTable.id == project_id, + ProjectsTable.user_id == str(user_id), + ) + return stmt + + async def count_memories_for_targeted_rebuild( + self, + user_id: UUID, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> int: + from sqlalchemy import func + async with self.db_adapter.system_session() as session: + base = self._build_targeted_rebuild_filter( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + count_stmt = select(func.count()).select_from(base.subquery()) + return (await session.execute(count_stmt)).scalar() or 0 + + async def get_memories_for_targeted_rebuild( + self, + user_id: UUID, + limit: int, + after_id: int | None = None, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> list[Memory]: + async with self.db_adapter.system_session() as session: + stmt = self._build_targeted_rebuild_filter( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + if after_id is not None: + stmt = stmt.where(MemoryTable.id > after_id) + stmt = ( + stmt + .options( + selectinload(MemoryTable.projects), + selectinload(MemoryTable.linked_memories), + selectinload(MemoryTable.linking_memories), + selectinload(MemoryTable.code_artifacts), + selectinload(MemoryTable.documents), + selectinload(MemoryTable.files), + ) + .order_by(MemoryTable.id.asc()) + .limit(limit) + ) + result = await session.execute(stmt) + memories_orm = result.scalars().unique().all() + return [Memory.model_validate(m) for m in memories_orm] + + async def upsert_targeted_embeddings( + self, + user_id: UUID, + updates: list[tuple[int, list[float]]], + ) -> list[int]: + """Idempotently write embeddings for memories owned by `user_id`. + + Implementation note: SQLite-vec's `vec_memories` virtual table is + keyed by `memory_id` and rejects duplicate inserts. We delete the + existing row (if any) before inserting the new embedding. We also + verify ownership of every memory id against the canonical `memories` + table so a malicious caller cannot poison another user's vector. + """ + if not updates: + return [] + ids = [mid for mid, _ in updates] + async with self.db_adapter.system_session() as session: + owned_rows = await session.execute( + select(MemoryTable.id).where( + MemoryTable.user_id == str(user_id), + MemoryTable.is_obsolete.is_(False), + MemoryTable.id.in_(ids), + ), + ) + owned_ids = {row[0] for row in owned_rows} + written: list[int] = [] + for memory_id, embedding in updates: + if memory_id not in owned_ids: + continue + embedding_bytes = sqlite_vec.serialize_float32(embedding) + await session.execute( + text( + "DELETE FROM vec_memories WHERE memory_id = :memory_id " + "AND EXISTS (" + "SELECT 1 FROM memories " + "WHERE memories.id = CAST(:memory_id AS INTEGER) " + "AND memories.user_id = :user_id " + "AND memories.is_obsolete = 0" + ")", + ), + {"memory_id": str(memory_id), "user_id": str(user_id)}, + ) + await session.execute( + text("INSERT INTO vec_memories (memory_id, embedding) VALUES (:memory_id, :embedding)"), + {"memory_id": str(memory_id), "embedding": embedding_bytes}, + ) + written.append(memory_id) + return written + async def validate_embedding_count(self) -> bool: """Check embedding count matches non-obsolete memory count""" async with self.db_adapter.system_session() as session: diff --git a/app/routes/mcp/memory_tools.py b/app/routes/mcp/memory_tools.py index 7674713..0a81b77 100644 --- a/app/routes/mcp/memory_tools.py +++ b/app/routes/mcp/memory_tools.py @@ -673,3 +673,97 @@ async def mark_memory_obsolete( "error_message": str(e), }) raise ToolError(f"INTERNAL_ERROR: Marking memory obsolete failed - {type(e).__name__}: {e!s}") + + @mcp.tool() + async def rebuild_embeddings( + ctx: Context, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> dict: + """Rebuild embeddings for a user-scoped subset of memories (issue #39). + + WHEN: A small set of memories, a project, or the caller's corpus needs + fresh embeddings without running the destructive global --re-embed CLI. + + BEHAVIOR: Re-runs the configured embedding pipeline on memories owned + by the authenticated user, scoped by `memory_ids` and/or `project_id`. + Never resets vector storage. Returns a structured result so the caller + can act on partial success. + + NOT-USE: Switching embedding providers/dimensions across the whole + corpus - use the offline `--re-embed` CLI for that destructive + migration. + + Args: + memory_ids: Explicit ids to rebuild (already user-scoped server-side). + project_id: Restrict to a single project owned by the caller. + + Returns: + Dict with keys: rebuilt_ids (list[int]), skipped_ids (list[int]), + failed (list of {memory_id, reason}). + """ + try: + logger.info("MCP Tool -> rebuild_embeddings", extra={ + "memory_ids": memory_ids, + "project_id": project_id, + }) + + user = await get_user_from_auth(ctx) + memory_service = ctx.fastmcp.memory_service + repo = memory_service.memory_repo + if memory_ids is not None and len(memory_ids) == 0: + raise ToolError( + "VALIDATION_ERROR: memory_ids must be non-empty if provided; " + "pass null for global scope", + ) + if memory_ids is not None and project_id is not None: + requested_count = len(set(memory_ids)) + count_with_project = await repo.count_memories_for_targeted_rebuild( + user_id=user.id, + memory_ids=memory_ids, + project_id=project_id, + ) + count_without_project = await repo.count_memories_for_targeted_rebuild( + user_id=user.id, + memory_ids=memory_ids, + ) + if count_with_project != count_without_project or count_without_project != requested_count: + raise ToolError( + "VALIDATION_ERROR: memory_ids do not all belong to " + f"project_id={project_id} (or some are not owned/obsolete). " + f"expected={requested_count}, matched_in_project={count_with_project}, " + f"owned_total={count_without_project}", + ) + + # Lazy import to avoid pulling re-embedding dependencies in modules + # that only need read-only memory access. + from app.services.re_embedding_service import ReEmbeddingService + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=repo.embedding_adapter, + ) + result = await service.rebuild_targeted( + user_id=user.id, + memory_ids=memory_ids, + project_id=project_id, + ) + return { + "rebuilt_ids": result.rebuilt_ids, + "skipped_ids": result.skipped_ids, + "failed": result.failed, + } + except ToolError: + raise + except ValidationError as e: + logger.debug("MCP Tool - rebuild_embeddings validation error", extra={ + "error_type": "ValidationError", + "error_message": str(e), + }) + raise ToolError(f"VALIDATION_ERROR: {e!s}") + except Exception as e: + logger.error("MCP Tool -> rebuild_embeddings failed", exc_info=True, extra={ + "error_type": type(e).__name__, + "error_message": str(e), + }) + raise ToolError(f"INTERNAL_ERROR: Rebuilding embeddings failed - {type(e).__name__}: {e!s}") diff --git a/app/routes/mcp/tool_adapters.py b/app/routes/mcp/tool_adapters.py index b79ab9d..a6bdc7d 100644 --- a/app/routes/mcp/tool_adapters.py +++ b/app/routes/mcp/tool_adapters.py @@ -7,6 +7,7 @@ from typing import Any from fastmcp import Context +from fastmcp.exceptions import ToolError from app.config.logging_config import logging from app.config.settings import settings @@ -47,6 +48,7 @@ from app.services.entity_service import EntityService from app.services.memory_service import MemoryService from app.services.project_service import ProjectService +from app.services.re_embedding_service import ReEmbeddingService from app.services.user_service import UserService from app.utils.pydantic_helper import filter_none_values @@ -108,6 +110,53 @@ def __init__(self, memory_service: MemoryService, user_service: UserService): self.memory_service = memory_service self.user_service = user_service + def _build_re_embedding_service(self) -> ReEmbeddingService: + """Build a ReEmbeddingService bound to the current memory repository. + + Lazy-built per-call so we always pick up the current embedding adapter + and avoid keeping process-wide migration state in this thin adapter. + """ + repo = self.memory_service.memory_repo + return ReEmbeddingService( + memory_repository=repo, + embedding_adapter=repo.embedding_adapter, + ) + + async def _validate_rebuild_scope( + self, + user_id, + memory_ids: list[int] | None, + project_id: int | None, + ) -> None: + """Validate rebuild scope before running embedding work.""" + if memory_ids is not None and len(memory_ids) == 0: + raise ToolError( + "VALIDATION_ERROR: memory_ids must be non-empty if provided; " + "pass null for global scope", + ) + + if memory_ids is None or project_id is None: + return + + requested_count = len(set(memory_ids)) + repo = self.memory_service.memory_repo + count_with_project = await repo.count_memories_for_targeted_rebuild( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + count_without_project = await repo.count_memories_for_targeted_rebuild( + user_id=user_id, + memory_ids=memory_ids, + ) + if count_with_project != count_without_project or count_without_project != requested_count: + raise ToolError( + "VALIDATION_ERROR: memory_ids do not all belong to " + f"project_id={project_id} (or some are not owned/obsolete). " + f"expected={requested_count}, matched_in_project={count_with_project}, " + f"owned_total={count_without_project}", + ) + async def create_memory( self, title: str, @@ -420,6 +469,50 @@ async def get_recent_memories( return memories + async def rebuild_embeddings( + self, + ctx: Context, + memory_ids: list[int] | None = None, + project_id: int | None = None, + ) -> dict[str, Any]: + """Adapter for `rebuild_embeddings` MCP tool (issue #39). + + Re-runs the embedding pipeline on a *user-scoped* subset of memories + without ever resetting global vector storage. Designed to refresh + explicit ids, a single project, or all caller-owned memories instead + of forcing the destructive `--re-embed` full-migration path used by + `re_embed_all`. + + Returns a structured result so callers can act on partial success: + `rebuilt_ids`, `skipped_ids`, `failed`. + """ + logger.info( + "MCP Tool -> rebuild_embeddings", + extra={ + "memory_ids": memory_ids, + "project_id": project_id, + }, + ) + + user = await get_user_from_auth(ctx) + await self._validate_rebuild_scope( + user_id=user.id, + memory_ids=memory_ids, + project_id=project_id, + ) + + service = self._build_re_embedding_service() + result = await service.rebuild_targeted( + user_id=user.id, + memory_ids=memory_ids, + project_id=project_id, + ) + return { + "rebuilt_ids": result.rebuilt_ids, + "skipped_ids": result.skipped_ids, + "failed": result.failed, + } + def create_memory_adapters( memory_service: MemoryService, user_service: UserService, @@ -435,6 +528,7 @@ def create_memory_adapters( "get_memory": adapters.get_memory, "mark_memory_obsolete": adapters.mark_memory_obsolete, "get_recent_memories": adapters.get_recent_memories, + "rebuild_embeddings": adapters.rebuild_embeddings, } diff --git a/app/routes/mcp/tool_metadata_registry.py b/app/routes/mcp/tool_metadata_registry.py index 0680b9b..c47fc97 100644 --- a/app/routes/mcp/tool_metadata_registry.py +++ b/app/routes/mcp/tool_metadata_registry.py @@ -722,6 +722,49 @@ def register_memory_tools_metadata( ], "tags": ["memory", "query", "recency", "timeline"], }, + { + "name": "rebuild_embeddings", + "mutates": True, + "description": ( + "Rebuild embeddings for a user-scoped subset of memories without resetting " + "global vector storage. Targets explicit memory ids, a single project, or " + "all caller-owned memories." + ), + "parameters": [ + { + "name": "ctx", + "type": "Context", + "description": "FastMCP Context (automatically injected)", + "required": True, + }, + { + "name": "memory_ids", + "type": "Optional[List[int]]", + "description": "Explicit memory ids owned by the caller; leave null to use project_id or global user scope", + "required": False, + "default": None, + "example": [123, 124], + }, + { + "name": "project_id", + "type": "Optional[int]", + "description": "Restrict the rebuild to memories of a single project owned by the caller", + "required": False, + "default": None, + "example": 1, + }, + ], + "returns": ( + "Dict with rebuilt_ids (List[int]), skipped_ids (List[int]), " + "failed (List[Dict[memory_id, reason]])" + ), + "examples": [ + 'execute_forgetful_tool("rebuild_embeddings", {"memory_ids": [123, 124]})', + 'execute_forgetful_tool("rebuild_embeddings", {"project_id": 1})', + 'execute_forgetful_tool("rebuild_embeddings", {})', + ], + "tags": ["memory", "embeddings", "admin", "repair"], + }, ] for tool_def in tools: diff --git a/app/services/re_embedding_service.py b/app/services/re_embedding_service.py index b49e8fd..f1750ac 100644 --- a/app/services/re_embedding_service.py +++ b/app/services/re_embedding_service.py @@ -5,8 +5,10 @@ """ import logging from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field +from uuid import UUID +from app.config.settings import settings from app.protocols.memory_protocol import MemoryRepository, ValidationResult from app.repositories.embeddings.embedding_adapter import EmbeddingsAdapter from app.repositories.helpers import build_embedding_text @@ -22,6 +24,19 @@ class ReEmbedResult: validation: ValidationResult | None = None +@dataclass +class TargetedRebuildResult: + """Outcome of a user-scoped, targeted `rebuild_embeddings` call. + + Distinguishes successful rebuilds from skipped (not owned / no-op) and + failed (embedding generation or write error) memories so callers can + surface partial success without re-running the whole rebuild. + """ + rebuilt_ids: list[int] = field(default_factory=list) + skipped_ids: list[int] = field(default_factory=list) + failed: list[dict] = field(default_factory=list) + + class ReEmbeddingService: """Orchestrates batch re-embedding of all memories""" @@ -99,6 +114,179 @@ async def re_embed_all( validation=validation, ) + async def rebuild_targeted( + self, + user_id: UUID, + memory_ids: list[int] | None = None, + project_id: int | None = None, + progress_callback: Callable[[int, int], None] | None = None, + ) -> TargetedRebuildResult: + """User-scoped, *non-destructive* embedding rebuild for a memory subset. + + Unlike `re_embed_all`, this method: + - never calls `reset_embedding_storage()`, + - operates only on memories owned by `user_id`, + - reuses the existing embedding pipeline (`build_embedding_text` + + `EmbeddingsAdapter.generate_embedding`) so embeddings stay + dimensionally consistent with the rest of the corpus. + + Args: + user_id: Authenticated user id; ownership filter, never optional. + memory_ids: Restrict to explicit memory ids (already user-scoped). + project_id: Restrict to a single project owned by `user_id`. + progress_callback: Optional `(processed, total)` callback per batch. + + Returns: + `TargetedRebuildResult` with per-memory outcome lists. + """ + total = await self.memory_repository.count_memories_for_targeted_rebuild( + user_id=user_id, + memory_ids=memory_ids, + project_id=project_id, + ) + result = TargetedRebuildResult() + if total == 0: + self._record_unresolved_memory_ids( + memory_ids=memory_ids, + project_id=project_id, + result=result, + ) + logger.info( + "rebuild_embeddings: no candidates", + extra={ + "user_id": str(user_id), + "memory_ids": memory_ids, + "project_id": project_id, + }, + ) + return result + + processed = 0 + last_seen_id: int | None = None + while True: + memories = await self.memory_repository.get_memories_for_targeted_rebuild( + user_id=user_id, + limit=self.batch_size, + after_id=last_seen_id, + memory_ids=memory_ids, + project_id=project_id, + ) + if not memories: + break + + updates: list[tuple[int, list[float]]] = [] + for memory in memories: + try: + embedding_text = build_embedding_text(memory) + embedding = await self.embedding_adapter.generate_embedding(embedding_text) + updates.append((memory.id, embedding)) + except Exception as exc: # noqa: BLE001 - intentional broad capture + logger.exception("rebuild_embeddings: embedding failed", extra={ + "memory_id": memory.id, + "user_id": str(user_id), + }) + result.failed.append({"memory_id": memory.id, "reason": str(exc)}) + + if updates: + try: + written_ids = await self.memory_repository.upsert_targeted_embeddings( + user_id=user_id, + updates=updates, + ) + written_set = set(written_ids) + result.rebuilt_ids.extend(written_ids) + for memory_id, _ in updates: + if memory_id not in written_set: + result.skipped_ids.append(memory_id) + await self._recompute_auto_links(user_id=user_id, memory_ids=written_ids, result=result) + except Exception as exc: # noqa: BLE001 + logger.exception("rebuild_embeddings: persistence failed", extra={ + "user_id": str(user_id), + "batch_size": len(updates), + }) + for memory_id, _ in updates: + result.failed.append({"memory_id": memory_id, "reason": str(exc)}) + + processed += len(memories) + if progress_callback: + progress_callback(processed, total) + last_seen_id = memories[-1].id + + self._record_unresolved_memory_ids( + memory_ids=memory_ids, + project_id=project_id, + result=result, + ) + logger.info( + "rebuild_embeddings: done", + extra={ + "user_id": str(user_id), + "rebuilt": len(result.rebuilt_ids), + "skipped": len(result.skipped_ids), + "failed": len(result.failed), + "total_candidates": total, + }, + ) + return result + + def _record_unresolved_memory_ids( + self, + memory_ids: list[int] | None, + project_id: int | None, + result: TargetedRebuildResult, + ) -> None: + """Surface explicit IDs hidden by user/obsolete filters as failures.""" + if memory_ids is None or project_id is not None: + return + + resolved_ids = set(result.rebuilt_ids) | set(result.skipped_ids) + resolved_ids.update( + failure["memory_id"] + for failure in result.failed + if "memory_id" in failure + ) + for memory_id in dict.fromkeys(memory_ids): + if memory_id not in resolved_ids: + result.failed.append({ + "memory_id": memory_id, + "reason": "not_found_or_not_owned", + }) + + async def _recompute_auto_links( + self, + user_id: UUID, + memory_ids: list[int], + result: TargetedRebuildResult, + ) -> None: + """Refresh additive auto-links after new vectors have been written.""" + if settings.MEMORY_NUM_AUTO_LINK <= 0: + return + + for memory_id in memory_ids: + try: + similar_memories = await self.memory_repository.find_similar_memories( + user_id=user_id, + memory_id=memory_id, + max_links=settings.MEMORY_NUM_AUTO_LINK, + ) + if not similar_memories: + continue + await self.memory_repository.create_links_batch( + user_id=user_id, + source_id=memory_id, + target_ids=[memory.id for memory in similar_memories], + ) + except Exception as exc: # noqa: BLE001 - link failures are per-memory partial failures + logger.exception("rebuild_embeddings: auto-link failed", extra={ + "memory_id": memory_id, + "user_id": str(user_id), + }) + result.failed.append({ + "memory_id": memory_id, + "phase": "auto_link", + "reason": str(exc), + }) + async def validate(self) -> ValidationResult: """Post-migration validation checks delegated to repository""" count_ok = await self.memory_repository.validate_embedding_count() diff --git a/tests/e2e_sqlite/test_re_embedding_sqlite.py b/tests/e2e_sqlite/test_re_embedding_sqlite.py index 07932b8..f401ffe 100644 --- a/tests/e2e_sqlite/test_re_embedding_sqlite.py +++ b/tests/e2e_sqlite/test_re_embedding_sqlite.py @@ -21,7 +21,8 @@ def embedding_adapter(): """Module-scoped embedding adapter (expensive to load). Forces FastEmbed-compatible model regardless of env config (e.g. docker/.env - may set EMBEDDING_MODEL to an OpenAI model). + may set EMBEDDING_MODEL to an OpenAI model). Skips the whole module when + the model cannot be downloaded (offline runners). """ original_model = settings.EMBEDDING_MODEL original_dims = settings.EMBEDDING_DIMENSIONS @@ -29,6 +30,8 @@ def embedding_adapter(): settings.EMBEDDING_DIMENSIONS = 384 try: return FastEmbeddingAdapter() + except Exception as exc: # noqa: BLE001 - HF/httpx network errors vary + pytest.skip(f"FastEmbed model unavailable (offline?): {exc}") finally: settings.EMBEDDING_MODEL = original_model settings.EMBEDDING_DIMENSIONS = original_dims @@ -220,3 +223,80 @@ async def test_re_embed_validation_checks(sqlite_repo, embedding_adapter): assert result.validation.dimensions_ok assert result.validation.search_ok assert result.validation.all_passed + + +# --------------------------------------------------------------------------- +# rebuild_targeted (issue #39) — user-scoped, non-destructive rebuild +# --------------------------------------------------------------------------- + + +async def _count_vec_memories(db_adapter) -> int: + async with db_adapter.system_session() as session: + result = await session.execute(text("SELECT COUNT(*) FROM vec_memories")) + return result.scalar() or 0 + + +@pytest.mark.asyncio +async def test_rebuild_targeted_is_user_scoped( + sqlite_repo, embedding_adapter, +): + """A user can never trigger a rebuild on another user's memories.""" + repo, db_adapter = sqlite_repo + user_a, memories_a = await _create_test_memories(repo, db_adapter, count=2) + user_b, memories_b = await _create_test_memories(repo, db_adapter, count=2) + + # Drop a vec row for user A's memory; user B should not be able to rebuild it. + target = memories_a[0] + async with db_adapter.system_session() as session: + await session.execute( + text("DELETE FROM vec_memories WHERE memory_id = :memory_id"), + {"memory_id": str(target.id)}, + ) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=10, + ) + # User B passes A's id explicitly: candidate count must be zero. + result_b = await service.rebuild_targeted( + user_id=user_b, + memory_ids=[target.id], + ) + + assert result_b.rebuilt_ids == [] + + # User A can repair their own memory. + result_a = await service.rebuild_targeted( + user_id=user_a, + ) + assert sorted(result_a.rebuilt_ids) == sorted(memory.id for memory in memories_a) + + +@pytest.mark.asyncio +async def test_rebuild_targeted_does_not_call_reset( + sqlite_repo, embedding_adapter, monkeypatch, +): + """Targeted rebuild must never invoke reset_embedding_storage.""" + repo, db_adapter = sqlite_repo + user_id, memories = await _create_test_memories(repo, db_adapter, count=2) + + reset_calls = {"n": 0} + original_reset = repo.reset_embedding_storage + + async def patched_reset(): + reset_calls["n"] += 1 + await original_reset() + + monkeypatch.setattr(repo, "reset_embedding_storage", patched_reset) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=5, + ) + await service.rebuild_targeted(user_id=user_id) + + assert reset_calls["n"] == 0 + # Vec storage retained for both memories + assert await _count_vec_memories(db_adapter) == 2 diff --git a/tests/integration/test_re_embedding_service.py b/tests/integration/test_re_embedding_service.py index e124793..4279595 100644 --- a/tests/integration/test_re_embedding_service.py +++ b/tests/integration/test_re_embedding_service.py @@ -7,6 +7,7 @@ import pytest +from app.config.settings import settings from app.models.memory_models import Memory from app.services.re_embedding_service import ReEmbeddingService @@ -202,3 +203,151 @@ async def test_re_embed_validation_failure(): assert result.validation.dimensions_ok assert result.validation.search_ok assert not result.validation.all_passed + + +# --------------------------------------------------------------------------- +# rebuild_targeted (issue #39) — user-scoped, non-destructive rebuild +# --------------------------------------------------------------------------- + + +def _make_targeted_repo(memories: list[Memory], owned_by_user: bool = True): + """Mock repo wired for rebuild_targeted (no reset, user-scoped upsert).""" + repo = AsyncMock() + repo.count_memories_for_targeted_rebuild.return_value = len(memories) + + async def get_batch(user_id, limit, after_id=None, memory_ids=None, + project_id=None): + eligible = memories + if after_id is not None: + eligible = [memory for memory in eligible if memory.id > after_id] + return eligible[:limit] + + async def upsert_targeted(user_id, updates): + return [mid for mid, _ in updates] if owned_by_user else [] + + repo.get_memories_for_targeted_rebuild.side_effect = get_batch + repo.upsert_targeted_embeddings.side_effect = upsert_targeted + repo.find_similar_memories.return_value = [] + repo.create_links_batch.return_value = [] + return repo + + +@pytest.mark.asyncio +async def test_rebuild_targeted_does_not_reset_storage(): + """rebuild_targeted must never call reset_embedding_storage().""" + from uuid import uuid4 + + memories = [_make_memory(i) for i in range(1, 4)] + repo = _make_targeted_repo(memories) + adapter = _make_mock_adapter() + + service = ReEmbeddingService(repo, adapter, batch_size=20) + result = await service.rebuild_targeted(user_id=uuid4()) + + repo.reset_embedding_storage.assert_not_called() + assert sorted(result.rebuilt_ids) == [1, 2, 3] + assert result.skipped_ids == [] + assert result.failed == [] + + +@pytest.mark.asyncio +async def test_rebuild_targeted_no_candidates_short_circuits(): + """Empty candidate set -> no embedding work, empty result.""" + from uuid import uuid4 + + repo = _make_targeted_repo([]) + adapter = _make_mock_adapter() + + service = ReEmbeddingService(repo, adapter, batch_size=20) + result = await service.rebuild_targeted(user_id=uuid4()) + + assert result.rebuilt_ids == [] + adapter.generate_embedding.assert_not_called() + repo.upsert_targeted_embeddings.assert_not_called() + + +@pytest.mark.asyncio +async def test_rebuild_targeted_skips_unowned_ids(): + """Memories rejected by upsert (not owned) land in skipped_ids.""" + from uuid import uuid4 + + memories = [_make_memory(i) for i in range(1, 4)] + repo = _make_mock_repo([]) # not used directly + repo.count_memories_for_targeted_rebuild.return_value = len(memories) + + async def get_batch(user_id, limit, after_id=None, memory_ids=None, + project_id=None): + eligible = memories + if after_id is not None: + eligible = [memory for memory in eligible if memory.id > after_id] + return eligible[:limit] + + async def partial_upsert(user_id, updates): + return [updates[0][0]] # Only the first id is owned + + repo.get_memories_for_targeted_rebuild.side_effect = get_batch + repo.upsert_targeted_embeddings.side_effect = partial_upsert + repo.find_similar_memories.return_value = [] + repo.create_links_batch.return_value = [] + adapter = _make_mock_adapter() + + service = ReEmbeddingService(repo, adapter, batch_size=20) + result = await service.rebuild_targeted(user_id=uuid4()) + + assert result.rebuilt_ids == [1] + assert sorted(result.skipped_ids) == [2, 3] + assert result.failed == [] + + +@pytest.mark.asyncio +async def test_rebuild_targeted_records_embedding_failures(): + """Per-memory embedding error -> entry in result.failed, others succeed.""" + from uuid import uuid4 + + memories = [_make_memory(i) for i in range(1, 4)] + repo = _make_targeted_repo(memories) + adapter = _make_mock_adapter() + + call_count = {"n": 0} + + async def flaky_embedding(_text): + call_count["n"] += 1 + if call_count["n"] == 2: + raise RuntimeError("transient API failure") + return [0.1] * 384 + + adapter.generate_embedding.side_effect = flaky_embedding + + service = ReEmbeddingService(repo, adapter, batch_size=20) + result = await service.rebuild_targeted(user_id=uuid4()) + + assert sorted(result.rebuilt_ids) == [1, 3] + assert len(result.failed) == 1 + assert result.failed[0]["memory_id"] == 2 + assert "transient API failure" in result.failed[0]["reason"] + + +@pytest.mark.asyncio +async def test_rebuild_targeted_recomputes_auto_links(monkeypatch): + """After writing fresh vectors, targeted rebuild refreshes additive auto-links.""" + from uuid import uuid4 + + user_id = uuid4() + memories = [_make_memory(i) for i in range(1, 3)] + repo = _make_targeted_repo(memories) + repo.find_similar_memories.return_value = [memories[1]] + repo.create_links_batch.return_value = [memories[1].id] + adapter = _make_mock_adapter() + monkeypatch.setattr(settings, "MEMORY_NUM_AUTO_LINK", 1) + + service = ReEmbeddingService(repo, adapter, batch_size=20) + result = await service.rebuild_targeted(user_id=user_id) + + assert sorted(result.rebuilt_ids) == [1, 2] + assert repo.find_similar_memories.call_count == 2 + assert repo.create_links_batch.call_count == 2 + repo.find_similar_memories.assert_any_call( + user_id=user_id, + memory_id=1, + max_links=1, + ) diff --git a/tests/integration/test_targeted_rebuild_sqlite.py b/tests/integration/test_targeted_rebuild_sqlite.py new file mode 100644 index 0000000..39ada0c --- /dev/null +++ b/tests/integration/test_targeted_rebuild_sqlite.py @@ -0,0 +1,418 @@ +"""Integration tests for targeted rebuild on a real (in-memory) SQLite stack. + +Unlike `tests/e2e_sqlite/test_re_embedding_sqlite.py`, this module never +downloads a FastEmbed model; it runs offline using a deterministic mock +embedding adapter so CI and offline laptops can exercise the persistence +layer of the new `rebuild_embeddings` MCP tool (issue #39). +""" +import hashlib +import random +from datetime import UTC, datetime +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastmcp.exceptions import ToolError +from sqlalchemy import text + +from app.config.settings import settings +from app.models.memory_models import MemoryCreate +from app.models.user_models import User +from app.repositories.sqlite.memory_repository import SqliteMemoryRepository +from app.repositories.sqlite.sqlite_adapter import SqliteDatabaseAdapter +from app.routes.mcp.tool_adapters import MemoryToolAdapters +from app.services.re_embedding_service import ReEmbeddingService + + +class StaticEmbeddingAdapter: + """Deterministic, fully offline embedding adapter for SQLite-vec tests.""" + + def __init__(self, dimensions: int = 384): + self.dimensions = dimensions + + async def generate_embedding(self, text: str) -> list[float]: + seed = int(hashlib.md5(text.encode()).hexdigest()[:8], 16) + rng = random.Random(seed) + vec = [rng.random() for _ in range(self.dimensions)] + magnitude = sum(x * x for x in vec) ** 0.5 or 1.0 + return [x / magnitude for x in vec] + + +@pytest.fixture +async def sqlite_repo(): + original_sqlite_memory = settings.SQLITE_MEMORY + original_database = settings.DATABASE + original_dims = settings.EMBEDDING_DIMENSIONS + + settings.DATABASE = "SQLite" + settings.SQLITE_MEMORY = True + settings.EMBEDDING_DIMENSIONS = 384 + + db_adapter = SqliteDatabaseAdapter() + await db_adapter.init_db() + + embedding_adapter = StaticEmbeddingAdapter(dimensions=384) + repo = SqliteMemoryRepository( + db_adapter=db_adapter, + embedding_adapter=embedding_adapter, + rerank_adapter=None, + ) + + try: + yield repo, db_adapter, embedding_adapter + finally: + try: + await db_adapter.dispose() + except Exception: # noqa: BLE001 + pass + settings.DATABASE = original_database + settings.SQLITE_MEMORY = original_sqlite_memory + settings.EMBEDDING_DIMENSIONS = original_dims + + +async def _create_user(db_adapter, user_id): + now = datetime.now(UTC).isoformat() + async with db_adapter.system_session() as session: + await session.execute( + text( + "INSERT INTO users (id, external_id, name, email, created_at, updated_at) " + "VALUES (:id, :external_id, :name, :email, :created_at, :updated_at)", + ), + { + "id": str(user_id), + "external_id": f"ext-{user_id}", + "name": "Test User", + "email": "test@example.com", + "created_at": now, + "updated_at": now, + }, + ) + + +async def _create_project(db_adapter, user_id, name="Test Project"): + now = datetime.now(UTC).isoformat() + async with db_adapter.system_session() as session: + result = await session.execute( + text( + "INSERT INTO projects (user_id, name, description, project_type, status, created_at, updated_at) " + "VALUES (:user_id, :name, :description, :project_type, :status, :created_at, :updated_at) " + "RETURNING id", + ), + { + "user_id": str(user_id), + "name": name, + "description": "Test project", + "project_type": "development", + "status": "active", + "created_at": now, + "updated_at": now, + }, + ) + return result.scalar_one() + + +def _make_ctx(user_id): + class FakeUserService: + async def get_or_create_user(self, user): + return User( + id=user_id, + external_id=user.external_id, + name=user.name, + email=user.email, + ) + + user_service = FakeUserService() + return SimpleNamespace( + fastmcp=SimpleNamespace( + auth=None, + user_service=user_service, + ), + ), user_service + + +async def _seed_memories(repo, db_adapter, count=3): + user_id = uuid4() + await _create_user(db_adapter, user_id) + memories = [] + for i in range(count): + memory = await repo.create_memory( + user_id=user_id, + memory=MemoryCreate( + title=f"M{i}", + content=f"content {i}", + context=f"ctx {i}", + keywords=[f"k{i}"], + tags=[f"t{i}"], + importance=5, + ), + ) + memories.append(memory) + return user_id, memories + + +async def _vec_count(db_adapter) -> int: + async with db_adapter.system_session() as session: + return (await session.execute(text("SELECT COUNT(*) FROM vec_memories"))).scalar() or 0 + + +async def _vec_blob(db_adapter, memory_id: int): + async with db_adapter.system_session() as session: + return ( + await session.execute( + text("SELECT embedding FROM vec_memories WHERE memory_id = :memory_id"), + {"memory_id": str(memory_id)}, + ) + ).scalar_one_or_none() + + +@pytest.mark.asyncio +async def test_rebuild_targeted_is_user_scoped(sqlite_repo): + repo, db_adapter, embedding_adapter = sqlite_repo + user_a, memories_a = await _seed_memories(repo, db_adapter, count=2) + user_b, _memories_b = await _seed_memories(repo, db_adapter, count=2) + + target = memories_a[0] + async with db_adapter.system_session() as session: + await session.execute( + text("DELETE FROM vec_memories WHERE memory_id = :memory_id"), + {"memory_id": str(target.id)}, + ) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=10, + ) + + result_b = await service.rebuild_targeted( + user_id=user_b, + memory_ids=[target.id], + ) + assert result_b.rebuilt_ids == [] + + result_a = await service.rebuild_targeted(user_id=user_a) + assert sorted(result_a.rebuilt_ids) == sorted(memory.id for memory in memories_a) + + +@pytest.mark.asyncio +async def test_rebuild_targeted_does_not_call_reset(sqlite_repo, monkeypatch): + repo, db_adapter, embedding_adapter = sqlite_repo + user_id, _memories = await _seed_memories(repo, db_adapter, count=2) + + reset_calls = {"n": 0} + original_reset = repo.reset_embedding_storage + + async def patched_reset(): + reset_calls["n"] += 1 + await original_reset() + + monkeypatch.setattr(repo, "reset_embedding_storage", patched_reset) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=5, + ) + await service.rebuild_targeted(user_id=user_id) + + assert reset_calls["n"] == 0 + assert await _vec_count(db_adapter) == 2 + + +@pytest.mark.asyncio +async def test_upsert_targeted_embeddings_rejects_unowned_ids(sqlite_repo): + repo, db_adapter, embedding_adapter = sqlite_repo + user_a, memories_a = await _seed_memories(repo, db_adapter, count=1) + user_b, memories_b = await _seed_memories(repo, db_adapter, count=1) + + fake_vector = await embedding_adapter.generate_embedding("anything") + + written = await repo.upsert_targeted_embeddings( + user_id=user_a, + updates=[(memories_b[0].id, fake_vector)], + ) + assert written == [] + + written = await repo.upsert_targeted_embeddings( + user_id=user_a, + updates=[(memories_a[0].id, fake_vector)], + ) + assert written == [memories_a[0].id] + + +@pytest.mark.asyncio +async def test_rebuild_targeted_recomputes_auto_links(sqlite_repo, monkeypatch): + repo, db_adapter, embedding_adapter = sqlite_repo + user_id, memories = await _seed_memories(repo, db_adapter, count=2) + monkeypatch.setattr(settings, "MEMORY_NUM_AUTO_LINK", 1) + + find_calls = [] + create_calls = [] + original_find = repo.find_similar_memories + original_create = repo.create_links_batch + + async def tracked_find_similar_memories(*args, **kwargs): + find_calls.append((args, kwargs)) + return await original_find(*args, **kwargs) + + async def tracked_create_links_batch(*args, **kwargs): + create_calls.append((args, kwargs)) + return await original_create(*args, **kwargs) + + monkeypatch.setattr(repo, "find_similar_memories", tracked_find_similar_memories) + monkeypatch.setattr(repo, "create_links_batch", tracked_create_links_batch) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=10, + ) + result = await service.rebuild_targeted(user_id=user_id, memory_ids=[memories[0].id]) + + assert result.rebuilt_ids == [memories[0].id] + assert len(find_calls) == 1 + assert find_calls[0][1]["memory_id"] == memories[0].id + assert create_calls + + +@pytest.mark.asyncio +async def test_rebuild_targeted_multibatch_rebuilds_all_explicit_ids(sqlite_repo): + repo, db_adapter, embedding_adapter = sqlite_repo + user_id, memories = await _seed_memories(repo, db_adapter, count=5) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=2, + ) + result = await service.rebuild_targeted( + user_id=user_id, + memory_ids=[memory.id for memory in memories], + ) + + assert sorted(result.rebuilt_ids) == sorted(memory.id for memory in memories) + assert result.failed == [] + + +@pytest.mark.asyncio +async def test_upsert_targeted_embeddings_leaves_foreign_vector_unchanged(sqlite_repo): + repo, db_adapter, embedding_adapter = sqlite_repo + user_a, _memories_a = await _seed_memories(repo, db_adapter, count=1) + _user_b, memories_b = await _seed_memories(repo, db_adapter, count=1) + foreign_memory = memories_b[0] + original_blob = await _vec_blob(db_adapter, foreign_memory.id) + + fake_vector = await embedding_adapter.generate_embedding("foreign overwrite attempt") + written = await repo.upsert_targeted_embeddings( + user_id=user_a, + updates=[(foreign_memory.id, fake_vector)], + ) + + assert written == [] + assert await _vec_blob(db_adapter, foreign_memory.id) == original_blob + + +@pytest.mark.asyncio +async def test_rebuild_targeted_recomputes_links_after_upsert(sqlite_repo, monkeypatch): + repo, db_adapter, embedding_adapter = sqlite_repo + user_id, memories = await _seed_memories(repo, db_adapter, count=2) + monkeypatch.setattr(settings, "MEMORY_NUM_AUTO_LINK", 1) + + order = [] + original_upsert = repo.upsert_targeted_embeddings + original_find = repo.find_similar_memories + + async def tracked_upsert(*args, **kwargs): + order.append("upsert") + return await original_upsert(*args, **kwargs) + + async def tracked_find(*args, **kwargs): + order.append("find") + assert order[-2] == "upsert" + assert await _vec_blob(db_adapter, kwargs["memory_id"]) is not None + return await original_find(*args, **kwargs) + + monkeypatch.setattr(repo, "upsert_targeted_embeddings", tracked_upsert) + monkeypatch.setattr(repo, "find_similar_memories", tracked_find) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=10, + ) + result = await service.rebuild_targeted(user_id=user_id, memory_ids=[memories[0].id]) + + assert result.rebuilt_ids == [memories[0].id] + assert order == ["upsert", "find"] + + +@pytest.mark.asyncio +async def test_rebuild_embeddings_rejects_empty_memory_ids(sqlite_repo): + repo, _db_adapter, _embedding_adapter = sqlite_repo + user_id = uuid4() + ctx, user_service = _make_ctx(user_id) + adapters = MemoryToolAdapters( + memory_service=SimpleNamespace(memory_repo=repo), + user_service=user_service, + ) + + with pytest.raises(ToolError, match="memory_ids must be non-empty"): + await adapters.rebuild_embeddings(ctx=ctx, memory_ids=[]) + + +@pytest.mark.asyncio +async def test_rebuild_embeddings_rejects_mismatched_memory_ids_and_project_id(sqlite_repo): + repo, db_adapter, _embedding_adapter = sqlite_repo + user_id = uuid4() + await _create_user(db_adapter, user_id) + project_a = await _create_project(db_adapter, user_id, name="A") + project_b = await _create_project(db_adapter, user_id, name="B") + memory = await repo.create_memory( + user_id=user_id, + memory=MemoryCreate( + title="Scoped memory", + content="content", + context="ctx", + keywords=["k"], + tags=["t"], + importance=5, + project_ids=[project_a], + ), + ) + ctx, user_service = _make_ctx(user_id) + adapters = MemoryToolAdapters( + memory_service=SimpleNamespace(memory_repo=repo), + user_service=user_service, + ) + + with pytest.raises(ToolError, match="do not all belong"): + await adapters.rebuild_embeddings( + ctx=ctx, + memory_ids=[memory.id], + project_id=project_b, + ) + + +@pytest.mark.asyncio +async def test_rebuild_targeted_records_unresolved_memory_ids_in_failed(sqlite_repo): + repo, db_adapter, embedding_adapter = sqlite_repo + user_a, memories_a = await _seed_memories(repo, db_adapter, count=1) + _user_b, memories_b = await _seed_memories(repo, db_adapter, count=1) + + service = ReEmbeddingService( + memory_repository=repo, + embedding_adapter=embedding_adapter, + batch_size=10, + ) + result = await service.rebuild_targeted( + user_id=user_a, + memory_ids=[memories_a[0].id, memories_b[0].id], + ) + + assert result.rebuilt_ids == [memories_a[0].id] + assert result.failed == [ + { + "memory_id": memories_b[0].id, + "reason": "not_found_or_not_owned", + }, + ]