diff --git a/application/api/user/sources/routes.py b/application/api/user/sources/routes.py index b1f6e15b..b45bc603 100644 --- a/application/api/user/sources/routes.py +++ b/application/api/user/sources/routes.py @@ -12,6 +12,7 @@ from application.agents.tools.path_utils import validate_tool_path from application.api import api from application.api.user.tasks import ( convert_source_to_wiki, + extract_graph, reembed_wiki_page, reingest_source_task, sync_source, @@ -22,6 +23,7 @@ from application.api.user.team_sharing import ( visible_with_access, ) from application.core.settings import settings +from application.graphrag import graphrag_available from application.parser.remote.remote_creator import normalize_remote_data from application.storage.db.repositories.ingest_chunk_progress import ( IngestChunkProgressRepository, @@ -1086,3 +1088,84 @@ class ConvertSourceToWiki(Resource): return make_response( jsonify({"success": True, "task_id": task.id}), 200 ) + + +@sources_ns.route("/sources//graphrag/enable") +class EnableSourceGraphRAG(Resource): + @api.doc( + description="Enable GraphRAG on a source: flips its config to graphrag " + "mode and enqueues graph extraction over its embedded chunks. Requires " + "VECTOR_STORE=pgvector and GRAPHRAG_ENABLED. Write access required " + "(owner/editor). The config kind cannot be set to graphrag via the " + "config PATCH endpoint." + ) + def post(self, source_id): + decoded_token = request.decoded_token + if not decoded_token: + return make_response(jsonify({"success": False}), 401) + if not graphrag_available(): + return make_response( + jsonify( + { + "success": False, + "message": ( + "GraphRAG requires VECTOR_STORE=pgvector and " + "GRAPHRAG_ENABLED." + ), + } + ), + 400, + ) + user = decoded_token.get("sub") + try: + with db_session() as conn: + owner = effective_write_owner(conn, "source", source_id, user) + if not owner: + return make_response( + jsonify( + {"success": False, "message": "Source not accessible"} + ), + 403, + ) + doc = SourcesRepository(conn).get_any(source_id, owner) + if doc is None: + return make_response( + jsonify({"success": False, "message": "Source not found"}), + 404, + ) + resolved_source_id = str(doc["id"]) + cfg = SourceConfig.parse(doc.get("config")) + repo = SourcesRepository(conn) + repo.update( + resolved_source_id, owner, {"config": cfg.graph_enabled()} + ) + # Re-read after the config write — it bumps ``updated_at``, + # which keys the extraction so each enable re-runs. + updated = repo.get_any(resolved_source_id, owner) + except Exception as err: + current_app.logger.error( + f"Error enabling GraphRAG for {source_id}: {err}", exc_info=True + ) + return make_response(jsonify({"success": False}), 400) + try: + from application.worker import ( + _source_updated_at, + graph_extraction_key, + ) + + task = extract_graph.delay( + resolved_source_id, + owner, + idempotency_key=graph_extraction_key( + resolved_source_id, _source_updated_at(updated) + ), + ) + except Exception as err: + current_app.logger.error( + f"Error starting GraphRAG extraction for {source_id}: {err}", + exc_info=True, + ) + return make_response(jsonify({"success": False}), 400) + return make_response( + jsonify({"success": True, "task_id": task.id}), 200 + ) diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 4de4bc83..7c45e15c 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -137,6 +137,15 @@ def convert_source_to_wiki(self, source_id, user, idempotency_key=None): return resp +@celery.task(**DURABLE_TASK) +@with_idempotency(task_name="extract_graph") +def extract_graph(self, source_id, user, idempotency_key=None): + from application.worker import extract_graph_worker + + resp = extract_graph_worker(self, source_id, 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/storage/db/source_config.py b/application/storage/db/source_config.py index 6ee354d6..f0f6242c 100644 --- a/application/storage/db/source_config.py +++ b/application/storage/db/source_config.py @@ -161,6 +161,18 @@ class SourceConfig(BaseModel): new_config["retrieval"]["exposure"] = "agentic_tool" return new_config + def graph_enabled(self) -> dict: + """Return a config dict flipped to GraphRAG mode (D28). + + Sets ``kind="graphrag"`` so ingest paths run graph extraction and + ``retrieval.retriever="graphrag"`` so the Dispatcher routes queries to + the graph retriever. All other fields are preserved. + """ + new_config = self.model_dump() + new_config["kind"] = "graphrag" + new_config["retrieval"]["retriever"] = "graphrag" + return new_config + @classmethod def parse(cls, raw: Optional[dict]) -> "SourceConfig": """Lenient read: return all-defaults for ``{}``/``None``. diff --git a/application/worker.py b/application/worker.py index 2e71294c..f5900080 100755 --- a/application/worker.py +++ b/application/worker.py @@ -57,6 +57,56 @@ MAX_TOKENS = 1250 RECURSION_DEPTH = 2 INGEST_HEARTBEAT_INTERVAL_SECONDS = 30 + +def graph_extraction_key(source_id, updated_at) -> str: + """Build the extract_graph idempotency key for a source's current state. + + The key embeds ``updated_at`` (the trigger-maintained column) so it + changes whenever re-extraction is warranted — a re-ingest or re-enable + bumps ``updated_at`` and yields a fresh key, bypassing the 24h completed + cache so the worker re-runs. Two enqueues for the same state share a key + and dedupe. ``updated_at`` is stringified deterministically. + """ + return f"extract-graph:{source_id}:{updated_at}" + + +def _source_updated_at(source) -> str: + """Return the source's ``updated_at`` (falling back to ``date``) as a string.""" + if not source: + return "" + stamp = source.get("updated_at") or source.get("date") + return str(stamp) if stamp is not None else "" + + +def _maybe_enqueue_graph_extraction(cfg, source_id, user): + """Enqueue graph extraction after embed for a graphrag source (D28). + + The graph lights up asynchronously once chunks are embedded, so ClassicRAG + works immediately. A no-op for non-graphrag sources or when GraphRAG is + unavailable. The enqueue is isolated so a broker hiccup can never fail an + otherwise-successful ingest. + """ + if cfg.kind != "graphrag": + return + from application.graphrag import graphrag_available + + if not graphrag_available(): + return + + source_id = str(source_id) + try: + from application.api.user.tasks import extract_graph + + with db_readonly() as conn: + source = SourcesRepository(conn).get_any(source_id, user) + key = graph_extraction_key(source_id, _source_updated_at(source)) + extract_graph.delay(source_id, user, idempotency_key=key) + except Exception as e: + logging.warning( + f"Failed to enqueue graph extraction for {source_id}: {e}", + exc_info=True, + ) + # Re-exported here for backward-compatible imports # (``from application.worker import _derive_source_id`` / # ``DOCSGPT_INGEST_NAMESPACE``) from tests and any other in-tree callers. @@ -654,6 +704,7 @@ def ingest_worker( }, scope={"kind": "source", "id": source_id_for_events}, ) + _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user) except Exception as e: logging.error(f"Error in ingest_worker: {e}", exc_info=True) publish_user_event( @@ -1033,6 +1084,7 @@ def reingest_source_worker(self, source_id, user): completed_payload, scope={"kind": "source", "id": source_id}, ) + _maybe_enqueue_graph_extraction(cfg, source_id, user) return { "source_id": source_id, @@ -1277,6 +1329,7 @@ def remote_worker( }, scope={"kind": "source", "id": source_id_for_events}, ) + _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user) except Exception as e: logging.error("Error in remote_worker task: %s", str(e), exc_info=True) publish_user_event( @@ -1862,6 +1915,7 @@ def ingest_connector( }, scope={"kind": "source", "id": source_id_for_events}, ) + _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user) return { "user": user, @@ -2287,3 +2341,50 @@ def convert_source_to_wiki_worker(self, source_id, user): "pages_created": len(created_pages), "skipped": skipped, } + + +def extract_graph_worker(self, source_id, user): + """Build a graphrag source's knowledge graph from its embedded chunks (D28). + + Loads the source, fetches its chunks from the vector store, and runs the + per-chunk LLM extraction pipeline. The chunks carry ``doc_id`` + ``text``, + matching the retrievable ids the extraction pipeline links against. No-ops + cleanly when GraphRAG is unavailable or the source has no chunks yet. + + Args: + self: Celery task instance. + source_id: Source whose graph is being built. + user: Owner identifier used to load the source and attribute token usage. + + Returns: + ``{"status": "unavailable"}`` when GraphRAG is off, otherwise the + extraction summary ``{nodes, edges, chunks_processed, ...}``. + """ + from application.graphrag import graphrag_available + from application.graphrag.extraction import extract_graph_for_source + from application.vectorstore.vector_creator import VectorCreator + + source_id = str(source_id) + + if not graphrag_available(): + return {"status": "unavailable"} + + 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, settings.EMBEDDINGS_KEY + ) + chunks = store.get_chunks() or [] + + return extract_graph_for_source( + source_id, + user, + chunks, + config=cfg, + request_id=getattr(self.request, "id", None), + ) diff --git a/tests/api/user/sources/test_graphrag_routes.py b/tests/api/user/sources/test_graphrag_routes.py new file mode 100644 index 00000000..427cb0a3 --- /dev/null +++ b/tests/api/user/sources/test_graphrag_routes.py @@ -0,0 +1,196 @@ +"""Tests for the GraphRAG enable route in application/api/user/sources/routes.py. + +``graphrag_available`` and ``extract_graph.delay`` are mocked so no live +vector store, LLM, or model calls run; the ``sources`` row is real so the +authz lookup and config write read back. +""" + +import uuid +from contextlib import contextmanager +from unittest.mock import patch + +import pytest +from flask import Flask + +from application.storage.db.repositories.sources import SourcesRepository +from application.storage.db.source_config import SourceConfig + + +@pytest.fixture +def app(): + return Flask(__name__) + + +@contextmanager +def _patch_db(conn): + @contextmanager + def _yield(): + yield conn + + with patch( + "application.api.user.sources.routes.db_session", _yield + ), patch( + "application.api.user.sources.routes.db_readonly", _yield + ): + yield + + +def _grant_team_access(pg_conn, owner, member, source_id, access_level): + from application.storage.db.repositories.team_members import ( + TeamMembersRepository, + ) + from application.storage.db.repositories.team_resource_grants import ( + TeamResourceGrantsRepository, + ) + from application.storage.db.repositories.teams import TeamsRepository + + team = TeamsRepository(pg_conn).create( + "Acme", f"acme-{uuid.uuid4().hex[:8]}", owner + ) + TeamMembersRepository(pg_conn).add_member( + team["id"], member, role="team_member" + ) + TeamResourceGrantsRepository(pg_conn).grant( + team["id"], "source", source_id, owner_id=owner, granted_by=owner, + access_level=access_level, + ) + + +@pytest.mark.unit +class TestEnableSourceGraphRAG: + def test_returns_401_unauthenticated(self, app): + from application.api.user.sources.routes import EnableSourceGraphRAG + + with app.test_request_context( + "/api/sources/x/graphrag/enable", method="POST" + ): + from flask import request + request.decoded_token = None + response = EnableSourceGraphRAG().post("x") + assert response.status_code == 401 + + def test_unavailable_returns_400(self, app, pg_conn): + from application.api.user.sources.routes import EnableSourceGraphRAG + + user = "u-graph-unavail" + src = SourcesRepository(pg_conn).create( + "files", user_id=user, type="file", + directory_structure={"a.md": {"type": "text/markdown"}}, + ) + sid = str(src["id"]) + + with _patch_db(pg_conn), patch( + "application.api.user.sources.routes.graphrag_available", + return_value=False, + ), patch( + "application.api.user.sources.routes.extract_graph.delay" + ) as mock_extract, app.test_request_context( + f"/api/sources/{sid}/graphrag/enable", method="POST" + ): + from flask import request + request.decoded_token = {"sub": user} + response = EnableSourceGraphRAG().post(sid) + + assert response.status_code == 400 + mock_extract.assert_not_called() + got = SourcesRepository(pg_conn).get_any(sid, user) + assert SourceConfig.parse(got.get("config")).kind == "classic" + + def test_owner_sets_config_and_enqueues(self, app, pg_conn): + from application.api.user.sources.routes import EnableSourceGraphRAG + + user = "u-graph-owner" + src = SourcesRepository(pg_conn).create( + "files", user_id=user, type="file", + directory_structure={"a.md": {"type": "text/markdown"}}, + ) + sid = str(src["id"]) + + fake_task = type("T", (), {"id": "task-g"})() + with _patch_db(pg_conn), patch( + "application.api.user.sources.routes.graphrag_available", + return_value=True, + ), patch( + "application.api.user.sources.routes.extract_graph.delay", + return_value=fake_task, + ) as mock_extract, app.test_request_context( + f"/api/sources/{sid}/graphrag/enable", method="POST" + ): + from flask import request + request.decoded_token = {"sub": user} + response = EnableSourceGraphRAG().post(sid) + + assert response.status_code == 200 + assert response.json["success"] is True + assert response.json["task_id"] == "task-g" + + got = SourcesRepository(pg_conn).get_any(sid, user) + cfg = SourceConfig.parse(got.get("config")) + assert cfg.kind == "graphrag" + assert cfg.retrieval.retriever == "graphrag" + + mock_extract.assert_called_once() + assert mock_extract.call_args.args[0] == sid + assert mock_extract.call_args.args[1] == user + # The key varies with the fresh ``updated_at`` the config write bumped, + # so each enable produces a new key that re-runs the worker. + key = mock_extract.call_args.kwargs["idempotency_key"] + assert key.startswith(f"extract-graph:{sid}:") + assert key != f"extract-graph:{sid}:" + + def test_viewer_rejected_403(self, app, pg_conn): + from application.api.user.sources.routes import EnableSourceGraphRAG + + owner = "alice-graph" + viewer = "bob-graph-viewer" + src = SourcesRepository(pg_conn).create( + "files", user_id=owner, type="file", + directory_structure={"a.md": {"type": "text/markdown"}}, + ) + sid = str(src["id"]) + _grant_team_access(pg_conn, owner, viewer, sid, "viewer") + + with _patch_db(pg_conn), patch( + "application.api.user.sources.routes.graphrag_available", + return_value=True, + ), patch( + "application.api.user.sources.routes.extract_graph.delay" + ) as mock_extract, app.test_request_context( + f"/api/sources/{sid}/graphrag/enable", method="POST" + ): + from flask import request + request.decoded_token = {"sub": viewer} + response = EnableSourceGraphRAG().post(sid) + + assert response.status_code == 403 + mock_extract.assert_not_called() + got = SourcesRepository(pg_conn).get_any(sid, owner) + assert SourceConfig.parse(got.get("config")).kind == "classic" + + +@pytest.mark.unit +class TestConfigPatchCannotSetGraphrag: + """The config PATCH endpoint must not flip kind to graphrag (D28).""" + + def test_patch_kind_graphrag_rejected_400(self, app, pg_conn): + from application.api.user.sources.routes import SourceConfigResource + + user = "u-patch-graph" + src = SourcesRepository(pg_conn).create( + "files", user_id=user, type="file", + directory_structure={"a.md": {"type": "text/markdown"}}, + ) + sid = str(src["id"]) + + with _patch_db(pg_conn), app.test_request_context( + f"/api/sources/{sid}/config", + method="PATCH", + json={"kind": "graphrag"}, + ): + from flask import request + request.decoded_token = {"sub": user} + response = SourceConfigResource().patch(sid) + + assert response.status_code == 400 + got = SourcesRepository(pg_conn).get_any(sid, user) + assert SourceConfig.parse(got.get("config")).kind == "classic" diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index b38c065a..b1170201 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -100,6 +100,43 @@ class TestConvertSourceToWikiTask: assert result == {"status": "converted"} +class TestExtractGraphTask: + @pytest.mark.unit + @patch("application.worker.extract_graph_worker") + def test_calls_extract_graph_worker(self, mock_worker): + from application.api.user.tasks import extract_graph + + mock_worker.return_value = {"nodes": 2, "edges": 1} + + result = extract_graph("source123", "user1") + + mock_worker.assert_called_once_with(ANY, "source123", "user1") + assert result == {"nodes": 2, "edges": 1} + + @pytest.mark.unit + def test_repeat_with_same_key_short_circuits(self, pg_conn): + from application.api.user import tasks + + calls: list[str] = [] + + def _fake_worker(self, source_id, user): + calls.append(source_id) + return {"nodes": 1, "edges": 0} + + with _patch_decorator_db(pg_conn), patch( + "application.worker.extract_graph_worker", _fake_worker + ): + first = tasks.extract_graph( + "src-g", "user1", idempotency_key="extract-graph:src-g", + ) + second = tasks.extract_graph( + "src-g", "user1", idempotency_key="extract-graph:src-g", + ) + + assert first == second + assert len(calls) == 1 + + class TestScheduleSyncsTask: @pytest.mark.unit @patch("application.api.user.tasks.sync_worker") @@ -302,6 +339,7 @@ class TestDurableTaskRetryPolicy: "ingest_connector_task", "reembed_wiki_page", "convert_source_to_wiki", + "extract_graph", ], ) def test_task_has_retry_config(self, task_name): diff --git a/tests/storage/db/test_source_config.py b/tests/storage/db/test_source_config.py index ab0cf28b..3a16e389 100644 --- a/tests/storage/db/test_source_config.py +++ b/tests/storage/db/test_source_config.py @@ -201,3 +201,19 @@ class TestGraphConfig: # Missing graph block still yields the default GraphConfig. cfg = SourceConfig.parse({"kind": "graphrag"}) assert cfg.graph == GraphConfig() + + def test_graph_enabled_sets_kind_and_retriever(self): + out = SourceConfig().graph_enabled() + assert out["kind"] == "graphrag" + assert out["retrieval"]["retriever"] == "graphrag" + # Round-trips through strict validation. + assert SourceConfig.model_validate(out).kind == "graphrag" + + def test_graph_enabled_preserves_other_fields(self): + cfg = SourceConfig.parse( + {"retrieval": {"chunks": 7, "exposure": "agentic_tool"}} + ) + out = cfg.graph_enabled() + assert out["retrieval"]["chunks"] == 7 + assert out["retrieval"]["exposure"] == "agentic_tool" + assert out["retrieval"]["retriever"] == "graphrag" diff --git a/tests/worker/test_extract_graph.py b/tests/worker/test_extract_graph.py new file mode 100644 index 00000000..efca0a07 --- /dev/null +++ b/tests/worker/test_extract_graph.py @@ -0,0 +1,117 @@ +"""Tests for ``application.worker.extract_graph_worker``. + +The worker loads the source row, fetches its chunks from the vector store, and +delegates to ``extract_graph_for_source``. ``graphrag_available``, +``VectorCreator.create_vectorstore``, and the extraction pipeline are mocked so +no live store access or LLM/model calls run; the ``sources`` row is real so +``SourcesRepository.get_any`` resolves. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from application.storage.db.repositories.sources import SourcesRepository + + +def _seed_source(pg_conn, config=None): + src = SourcesRepository(pg_conn).create( + "graph-set", + user_id="alice", + type="file", + retriever="classic", + config=config or {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}}, + ) + return str(src["id"]) + + +def _patch_store(monkeypatch, chunks): + store = MagicMock(name="vectorstore") + store.get_chunks.return_value = chunks + monkeypatch.setattr( + "application.vectorstore.vector_creator.VectorCreator.create_vectorstore", + lambda *a, **kw: store, + ) + return store + + +@pytest.mark.unit +class TestExtractGraphWorker: + def test_fetches_chunks_and_calls_extraction( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + source_id = _seed_source(pg_conn) + chunks = [ + {"doc_id": "c1", "text": "alpha"}, + {"doc_id": "c2", "text": "beta"}, + ] + store = _patch_store(monkeypatch, chunks) + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + extract = MagicMock( + name="extract_graph_for_source", + return_value={"nodes": 3, "edges": 2, "chunks_processed": 2}, + ) + monkeypatch.setattr( + "application.graphrag.extraction.extract_graph_for_source", extract + ) + + result = worker.extract_graph_worker(task_self, source_id, "alice") + + store.get_chunks.assert_called_once() + extract.assert_called_once() + assert extract.call_args.args[0] == source_id + assert extract.call_args.args[1] == "alice" + assert extract.call_args.args[2] == chunks + assert extract.call_args.kwargs["config"].kind == "graphrag" + assert result == {"nodes": 3, "edges": 2, "chunks_processed": 2} + + def test_unavailable_returns_status_no_extraction( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + store = _patch_store(monkeypatch, []) + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: False + ) + extract = MagicMock(name="extract_graph_for_source") + monkeypatch.setattr( + "application.graphrag.extraction.extract_graph_for_source", extract + ) + + result = worker.extract_graph_worker(task_self, "src-x", "alice") + + assert result == {"status": "unavailable"} + store.get_chunks.assert_not_called() + extract.assert_not_called() + + def test_empty_chunks_still_calls_extraction( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + from application import worker + + source_id = _seed_source(pg_conn) + _patch_store(monkeypatch, []) + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + extract = MagicMock( + name="extract_graph_for_source", + return_value={"nodes": 0, "edges": 0, "chunks_processed": 0}, + ) + monkeypatch.setattr( + "application.graphrag.extraction.extract_graph_for_source", extract + ) + + result = worker.extract_graph_worker(task_self, source_id, "alice") + + extract.assert_called_once() + assert extract.call_args.args[2] == [] + assert result["chunks_processed"] == 0 diff --git a/tests/worker/test_graph_extraction_enqueue.py b/tests/worker/test_graph_extraction_enqueue.py new file mode 100644 index 00000000..088eca00 --- /dev/null +++ b/tests/worker/test_graph_extraction_enqueue.py @@ -0,0 +1,272 @@ +"""Ingest paths enqueue graph extraction after embed for graphrag sources (D28). + +Exercises ``remote_worker`` as the representative ingest path: the remote +loader, embedding pipeline, and ``upload_index`` are mocked, so the only thing +under test is the post-embed ``extract_graph.delay`` hook. ``graphrag_available`` +is forced on so the gate doesn't depend on the test env's vector store. The +``sources`` row is real (seeded via ``pg_conn``) so the hook can read the +source's ``updated_at`` for the idempotency key. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from application.parser.schema.base import Document +from application.storage.db.repositories.sources import SourcesRepository + + +@pytest.fixture +def _mock_remote_pipeline(monkeypatch): + from application import worker + + fake_loader = MagicMock(name="remote_loader") + fake_loader.load_data.return_value = [ + Document( + text="page body", + extra_info={"file_path": "guides/setup.md", "title": "setup"}, + doc_id="d1", + ) + ] + monkeypatch.setattr( + worker.RemoteCreator, "create_loader", lambda loader: fake_loader + ) + monkeypatch.setattr( + worker, + "embed_and_store_documents", + lambda docs, full_path, source_id, task, **kw: None, + ) + monkeypatch.setattr( + worker, "assert_index_complete", lambda *a, **kw: None + ) + monkeypatch.setattr( + worker, "upload_index", lambda full_path, file_data: None + ) + + +def _patch_delay(monkeypatch): + delay = MagicMock(name="extract_graph_delay") + monkeypatch.setattr( + "application.api.user.tasks.extract_graph.delay", delay + ) + return delay + + +def _seed_source(pg_conn, user, config): + src = SourcesRepository(pg_conn).create( + "graph-remote", user_id=user, type="crawler", config=config, + ) + return str(src["id"]) + + +@pytest.mark.unit +class TestGraphExtractionKey: + def test_shape_and_state_sensitivity(self): + from application.worker import _source_updated_at, graph_extraction_key + + sid = "11111111-1111-1111-1111-111111111111" + key_a = graph_extraction_key( + sid, _source_updated_at({"updated_at": "2026-06-23T00:00:00+00:00"}) + ) + key_b = graph_extraction_key( + sid, _source_updated_at({"updated_at": "2026-06-23T00:00:01+00:00"}) + ) + key_a_again = graph_extraction_key( + sid, _source_updated_at({"updated_at": "2026-06-23T00:00:00+00:00"}) + ) + + assert key_a.startswith(f"extract-graph:{sid}:") + # A changed ``updated_at`` (re-ingest/re-enable) yields a fresh key. + assert key_a != key_b + # The same state collapses to one key (concurrent dups dedupe). + assert key_a == key_a_again + + def test_falls_back_to_date(self): + from application.worker import _source_updated_at + + assert _source_updated_at({"date": "2026-06-23T00:00:00+00:00"}) == ( + "2026-06-23T00:00:00+00:00" + ) + assert _source_updated_at({"updated_at": None, "date": "x"}) == "x" + assert _source_updated_at(None) == "" + + +@pytest.mark.unit +class TestRemoteWorkerEnqueuesGraphExtraction: + def test_graphrag_source_enqueues_after_embed( + self, task_self, pg_conn, patch_worker_db, monkeypatch, + _mock_remote_pipeline, + ): + from application import worker + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + delay = _patch_delay(monkeypatch) + + config = {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}} + sid = _seed_source(pg_conn, "bob", config) + + worker.remote_worker( + task_self, + {"urls": ["http://example.com"]}, + "graph-remote", + "bob", + "crawler", + directory="temp", + retriever="classic", + operation_mode="upload", + config=config, + source_id=sid, + ) + + delay.assert_called_once() + assert delay.call_args.args[0] == sid + assert delay.call_args.args[1] == "bob" + key = delay.call_args.kwargs["idempotency_key"] + assert key.startswith(f"extract-graph:{sid}:") + + def test_same_state_dedupes_to_one_key( + self, task_self, pg_conn, patch_worker_db, monkeypatch, + _mock_remote_pipeline, + ): + """Two enqueues for the same source state share a key (concurrent dups).""" + from application import worker + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + delay = _patch_delay(monkeypatch) + + config = {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}} + sid = _seed_source(pg_conn, "bob", config) + + def _run(): + worker.remote_worker( + task_self, + {"urls": ["http://example.com"]}, + "graph-remote", + "bob", + "crawler", + directory="temp", + retriever="classic", + operation_mode="upload", + config=config, + source_id=sid, + ) + return delay.call_args.kwargs["idempotency_key"] + + assert _run() == _run() + + def test_classic_source_does_not_enqueue( + self, task_self, pg_conn, patch_worker_db, monkeypatch, + _mock_remote_pipeline, + ): + from application import worker + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + delay = _patch_delay(monkeypatch) + + sid = _seed_source(pg_conn, "bob", {"kind": "classic"}) + + worker.remote_worker( + task_self, + {"urls": ["http://example.com"]}, + "classic-remote", + "bob", + "crawler", + directory="temp", + retriever="classic", + operation_mode="upload", + config={"kind": "classic"}, + source_id=sid, + ) + + delay.assert_not_called() + + def test_graphrag_unavailable_does_not_enqueue( + self, task_self, pg_conn, patch_worker_db, monkeypatch, + _mock_remote_pipeline, + ): + from application import worker + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: False + ) + delay = _patch_delay(monkeypatch) + + config = {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}} + sid = _seed_source(pg_conn, "bob", config) + + worker.remote_worker( + task_self, + {"urls": ["http://example.com"]}, + "graph-remote", + "bob", + "crawler", + directory="temp", + retriever="classic", + operation_mode="upload", + config=config, + source_id=sid, + ) + + delay.assert_not_called() + + +@pytest.mark.unit +class TestEnqueueIsolatesBrokerFailures: + def test_delay_exception_does_not_propagate( + self, pg_conn, patch_worker_db, monkeypatch + ): + """A broker hiccup in ``.delay`` must not fail an otherwise-good ingest.""" + from application import worker + from application.storage.db.source_config import SourceConfig + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + + def _boom(*a, **kw): + raise RuntimeError("broker down") + + monkeypatch.setattr( + "application.api.user.tasks.extract_graph.delay", _boom + ) + + config = {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}} + sid = _seed_source(pg_conn, "bob", config) + + worker._maybe_enqueue_graph_extraction( + SourceConfig.parse(config), sid, "bob" + ) + + def test_updated_at_read_exception_does_not_propagate( + self, pg_conn, patch_worker_db, monkeypatch + ): + """A DB hiccup reading ``updated_at`` must also be swallowed.""" + from application import worker + from application.storage.db.source_config import SourceConfig + + monkeypatch.setattr( + "application.graphrag.graphrag_available", lambda: True + ) + delay = _patch_delay(monkeypatch) + + def _boom(self, source_id, user_id): + raise RuntimeError("pg down") + + monkeypatch.setattr(SourcesRepository, "get_any", _boom) + + config = {"kind": "graphrag", "retrieval": {"retriever": "graphrag"}} + sid = _seed_source(pg_conn, "bob", config) + + worker._maybe_enqueue_graph_extraction( + SourceConfig.parse(config), sid, "bob" + ) + + delay.assert_not_called()