From d68be86244ca3264e7fe917e2e5dba9a3b8e06b0 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 22 Jun 2026 17:49:04 +0100 Subject: [PATCH] feat(wiki): reembed_wiki_page durable task Per-page re-embed (Unit 2 of F-Wiki): targeted delete of the page's old chunks, re-chunk via the source's chunking config, add_chunk with reingest-matching metadata (source=path), set embed_status embedded/failed. Durable + idempotent (key=content_hash), mirroring reingest_source_task. Missing page => purge only. --- application/api/user/tasks.py | 10 ++ application/worker.py | 80 +++++++++++ tests/worker/test_reembed_wiki_page.py | 183 +++++++++++++++++++++++++ 3 files changed, 273 insertions(+) create mode 100644 tests/worker/test_reembed_wiki_page.py diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 990aef70..9eed3202 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -7,6 +7,7 @@ from application.worker import ( attachment_worker, ingest_worker, mcp_oauth, + reembed_wiki_page_worker, remote_worker, sync, sync_worker, @@ -118,6 +119,15 @@ def reingest_source_task(self, source_id, user, idempotency_key=None): return resp +@celery.task(**DURABLE_TASK) +@with_idempotency(task_name="reembed_wiki_page") +def reembed_wiki_page( + self, source_id, path, content_hash, user, idempotency_key=None, +): + resp = reembed_wiki_page_worker(self, source_id, path, content_hash, user) + return resp + + # Beat-driven dispatch tasks default to ``acks_late=False``: a SIGKILL # of a beat tick is harmless to redeliver only if the dispatch itself is # idempotent. We keep these early-ACK so the broker doesn't replay a diff --git a/application/worker.py b/application/worker.py index ada4e35f..7e61bc9c 100755 --- a/application/worker.py +++ b/application/worker.py @@ -39,6 +39,7 @@ from application.storage.db.repositories.ingest_chunk_progress import ( IngestChunkProgressRepository, ) from application.storage.db.repositories.sources import SourcesRepository +from application.storage.db.repositories.wiki_pages import WikiPagesRepository from application.storage.db.session import db_readonly, db_session from application.storage.db.source_config import SourceConfig from application.storage.storage_creator import StorageCreator @@ -1978,3 +1979,82 @@ def mcp_oauth(self, config: Dict[str, Any], user_id: str = None) -> Dict[str, An logging.error("MCP OAuth init failed: %s", error_msg, exc_info=True) publish_oauth("mcp.oauth.failed", {"error": error_msg[:1024]}) return {"success": False, "error": error_msg} + + +def reembed_wiki_page_worker(self, source_id, path, content_hash, user): + """Re-embed one wiki page after an edit, or purge its chunks on delete. + + Targeted delete of the page's existing chunks runs first, so a deleted + page becomes a pure purge and an edited page is re-embedded cleanly. The + chunk metadata shape matches the reingest path (``source``/``title``/ + ``filename``) so retrieval and future targeted deletes line up. + + Args: + self: Celery task instance. + source_id: Source the page belongs to. + path: Page path; doubles as the chunk ``source`` metadata key. + content_hash: The edit's content hash (idempotency anchor). + user: Owner identifier used to load the source row. + + Returns: + A status dict: ``{"status", "added", "deleted"}`` on re-embed, or + ``{"status": "deleted", "deleted": n}`` when the page is gone. + """ + from application.vectorstore.vector_creator import VectorCreator + + source_id = str(source_id) + + with db_readonly() as conn: + source = SourcesRepository(conn).get_any(source_id, user) + if not source: + raise ValueError(f"Source {source_id} not found or access denied") + source_id = str(source["id"]) + cfg = SourceConfig.parse(source.get("config")) + + store = VectorCreator.create_vectorstore(settings.VECTOR_STORE, source_id) + deleted = store.delete_chunks_by_source_path(path) + + with db_readonly() as conn: + page = WikiPagesRepository(conn).get_by_path(source_id, path) + + if page is None: + return {"status": "deleted", "deleted": deleted} + + try: + title = page.get("title") or path + chunker = ChunkerCreator.create_chunker( + cfg.chunking.strategy, + chunking_strategy=cfg.chunking.strategy, + max_tokens=cfg.chunking.max_tokens, + min_tokens=cfg.chunking.min_tokens, + duplicate_headers=cfg.chunking.duplicate_headers, + ) + chunks = chunker.chunk( + documents=[ + Document( + text=page["content"], + extra_info={ + "source": path, + "title": title, + "filename": path, + }, + ) + ] + ) + + added = 0 + for chunk in chunks: + store.add_chunk( + chunk.text, + metadata={"source": path, "title": title, "filename": path}, + ) + added += 1 + + with db_session() as conn: + WikiPagesRepository(conn).set_embed_status(source_id, path, "embedded") + except Exception: + with db_session() as conn: + WikiPagesRepository(conn).set_embed_status(source_id, path, "failed") + raise + + return {"status": "embedded", "added": added, "deleted": deleted} diff --git a/tests/worker/test_reembed_wiki_page.py b/tests/worker/test_reembed_wiki_page.py new file mode 100644 index 00000000..36027734 --- /dev/null +++ b/tests/worker/test_reembed_wiki_page.py @@ -0,0 +1,183 @@ +"""Tests for the per-page wiki re-embed worker and Celery task. + +The worker loads the source row from the ephemeral DB, then re-embeds (or +purges) one wiki page. The vector store, ``ChunkerCreator``, and the page +lookup are mocked so no live embeddings run; the ``sources`` row is real so +``SourcesRepository.get_any`` resolves. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock + +import pytest + +from application.parser.schema.base import Document +from application.storage.db.repositories.sources import SourcesRepository + + +def _seed_source(pg_conn) -> str: + src = SourcesRepository(pg_conn).create( + "wiki-set", + user_id="alice", + type="wiki", + config={"kind": "wiki"}, + ) + return str(src["id"]) + + +def _patch_store(monkeypatch, store): + monkeypatch.setattr( + "application.vectorstore.vector_creator.VectorCreator.create_vectorstore", + lambda *a, **kw: store, + ) + + +def _patch_repo(monkeypatch, page): + """Replace ``WikiPagesRepository`` with one returning ``page`` and a stub setter.""" + repo = MagicMock(name="wiki_repo") + repo.get_by_path.return_value = page + repo.set_embed_status.return_value = True + monkeypatch.setattr( + "application.worker.WikiPagesRepository", lambda conn: repo + ) + return repo + + +def _patch_chunker(monkeypatch, chunks): + calls: list[dict] = [] + + def _create_chunker(strategy, **kwargs): + calls.append({"strategy": strategy, **kwargs}) + chunker = MagicMock(name="chunker") + chunker.chunk.return_value = chunks + return chunker + + monkeypatch.setattr( + "application.worker.ChunkerCreator.create_chunker", + staticmethod(_create_chunker), + ) + return calls + + +@pytest.mark.unit +class TestReembedWikiPageWorker: + def test_page_exists_reembeds( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + source_id = _seed_source(pg_conn) + + store = MagicMock(name="vector_store") + store.delete_chunks_by_source_path.return_value = 3 + _patch_store(monkeypatch, store) + + repo = _patch_repo( + monkeypatch, + {"content": "page body", "title": "Page Title"}, + ) + _patch_chunker( + monkeypatch, + [Document(text="c1"), Document(text="c2")], + ) + + result = worker.reembed_wiki_page_worker( + task_self, source_id, "guide/intro.md", "hash-1", "alice" + ) + + store.delete_chunks_by_source_path.assert_called_once_with( + "guide/intro.md" + ) + assert store.add_chunk.call_count == 2 + for call in store.add_chunk.call_args_list: + assert call.kwargs["metadata"]["source"] == "guide/intro.md" + assert call.kwargs["metadata"]["title"] == "Page Title" + assert call.kwargs["metadata"]["filename"] == "guide/intro.md" + repo.set_embed_status.assert_called_once_with( + source_id, "guide/intro.md", "embedded" + ) + assert result == {"status": "embedded", "added": 2, "deleted": 3} + + def test_page_missing_purges( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + source_id = _seed_source(pg_conn) + + store = MagicMock(name="vector_store") + store.delete_chunks_by_source_path.return_value = 4 + _patch_store(monkeypatch, store) + + repo = _patch_repo(monkeypatch, None) + + result = worker.reembed_wiki_page_worker( + task_self, source_id, "old/page.md", "hash-x", "alice" + ) + + store.delete_chunks_by_source_path.assert_called_once_with("old/page.md") + store.add_chunk.assert_not_called() + repo.set_embed_status.assert_not_called() + assert result == {"status": "deleted", "deleted": 4} + + def test_embed_failure_sets_failed_and_reraises( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + source_id = _seed_source(pg_conn) + + store = MagicMock(name="vector_store") + store.delete_chunks_by_source_path.return_value = 0 + store.add_chunk.side_effect = RuntimeError("embed boom") + _patch_store(monkeypatch, store) + + repo = _patch_repo( + monkeypatch, {"content": "page body", "title": "T"} + ) + _patch_chunker(monkeypatch, [Document(text="c1")]) + + with pytest.raises(RuntimeError, match="embed boom"): + worker.reembed_wiki_page_worker( + task_self, source_id, "p.md", "hash-2", "alice" + ) + + repo.set_embed_status.assert_called_once_with(source_id, "p.md", "failed") + + +@pytest.mark.unit +class TestReembedWikiPageTask: + def test_idempotency_key_is_content_hash(self, pg_conn, monkeypatch): + """A redelivery with the same content_hash key short-circuits the worker.""" + from application.api.user import tasks + + @contextmanager + def _yield(): + yield pg_conn + + monkeypatch.setattr( + "application.api.user.idempotency.db_session", _yield + ) + monkeypatch.setattr( + "application.api.user.idempotency.db_readonly", _yield + ) + + calls: list[str] = [] + + def _fake_worker(self, source_id, path, content_hash, user): + calls.append(content_hash) + return {"status": "embedded", "added": 1, "deleted": 0} + + monkeypatch.setattr(tasks, "reembed_wiki_page_worker", _fake_worker) + + first = tasks.reembed_wiki_page( + "src-1", "p.md", "abc123", "alice", idempotency_key="abc123", + ) + second = tasks.reembed_wiki_page( + "src-1", "p.md", "abc123", "alice", idempotency_key="abc123", + ) + + assert first == second + assert len(calls) == 1