mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-18 00:14:38 +00:00
feat(graphrag): durable extract_graph task + ingest-path enqueue + enable route
graph_enabled() sets kind=graphrag + retriever=graphrag. extract_graph_worker fetches the source's pgvector chunks and runs G3 extraction (graphrag_available guard, empty no-op). extract_graph durable+idempotent task; key varies with source updated_at so re-ingest/re-enable re-run incrementally (G3 checkpoint skips done chunks) while concurrent same-state enqueues dedup. The 4 ingest paths enqueue after embed when kind=graphrag (isolated in try/except so a broker hiccup can't fail the ingest). POST /api/sources/<id>/graphrag/enable: pgvector+GRAPHRAG_ENABLED gated, owner/editor write-authz; PATCH config still can't flip kind->graphrag. Unit G4.
This commit is contained in:
1 parent
22582e900c
commit
4742aec4c6
9 files changed
+844
No files matched your search
@@ -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"
|
||||
@@ -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):
|
||||
|
||||
Reference in new issue
Block a user