mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
arc53-machine
5d0992eef8
run_agent_headless records each unattended run under its endpoint; the scheduler passes its run id and the webhook worker its task id so Logs rows can find their trace, while the LLM's own request id stays untouched for quota counts. /api/search and MCP search_docs record their retrieval, and a graph build records every extraction call under one step.
223 lines
7.7 KiB
Python
223 lines
7.7 KiB
Python
"""Tests for ``docsgpt.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 docsgpt.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(
|
|
"docsgpt.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 docsgpt 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(
|
|
"docsgpt.graphrag.graphrag_available", lambda: True
|
|
)
|
|
extract = MagicMock(
|
|
name="extract_graph_for_source",
|
|
return_value={"nodes": 3, "edges": 2, "chunks_processed": 2},
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.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 docsgpt import worker
|
|
|
|
store = _patch_store(monkeypatch, [])
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.graphrag_available", lambda: False
|
|
)
|
|
extract = MagicMock(name="extract_graph_for_source")
|
|
monkeypatch.setattr(
|
|
"docsgpt.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 docsgpt import worker
|
|
|
|
source_id = _seed_source(pg_conn)
|
|
_patch_store(monkeypatch, [])
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.graphrag_available", lambda: True
|
|
)
|
|
extract = MagicMock(
|
|
name="extract_graph_for_source",
|
|
return_value={"nodes": 0, "edges": 0, "chunks_processed": 0},
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.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
|
|
|
|
def test_publishes_completed_event(
|
|
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
|
):
|
|
from docsgpt import worker
|
|
|
|
source_id = _seed_source(pg_conn)
|
|
_patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}])
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.graphrag_available", lambda: True
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.extraction.extract_graph_for_source",
|
|
MagicMock(return_value={"nodes": 1, "edges": 0, "chunks_processed": 1}),
|
|
)
|
|
events = []
|
|
monkeypatch.setattr(
|
|
worker, "publish_user_event",
|
|
lambda user, etype, payload, **kw: events.append((etype, payload)),
|
|
)
|
|
|
|
worker.extract_graph_worker(task_self, source_id, "alice")
|
|
|
|
types = [e[0] for e in events]
|
|
assert "graph.extract.progress" in types
|
|
assert types[-1] == "graph.extract.completed"
|
|
assert events[-1][1]["nodes"] == 1
|
|
|
|
def test_publishes_failed_event_on_error(
|
|
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
|
):
|
|
from docsgpt import worker
|
|
|
|
source_id = _seed_source(pg_conn)
|
|
_patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}])
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.graphrag_available", lambda: True
|
|
)
|
|
|
|
def _boom(*a, **kw):
|
|
raise RuntimeError("extraction blew up")
|
|
|
|
monkeypatch.setattr(
|
|
"docsgpt.graphrag.extraction.extract_graph_for_source", _boom
|
|
)
|
|
events = []
|
|
monkeypatch.setattr(
|
|
worker, "publish_user_event",
|
|
lambda user, etype, payload, **kw: events.append((etype, payload)),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError):
|
|
worker.extract_graph_worker(task_self, source_id, "alice")
|
|
|
|
assert "graph.extract.failed" in [e[0] for e in events]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExtractGraphTrace:
|
|
"""A graph build is one execution trace holding every extraction call."""
|
|
|
|
def _run(self, pg_conn, monkeypatch, task_self, extract):
|
|
from docsgpt import tracing, worker
|
|
from docsgpt.core.settings import settings
|
|
|
|
monkeypatch.setattr(settings, "TRACES_ENABLED", True)
|
|
monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False)
|
|
source_id = _seed_source(pg_conn)
|
|
_patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}])
|
|
monkeypatch.setattr("docsgpt.graphrag.graphrag_available", lambda: True)
|
|
monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract)
|
|
flushed = []
|
|
|
|
def _flush(trace, status=None):
|
|
trace.flushed = True
|
|
trace.finish(status)
|
|
flushed.append(trace)
|
|
|
|
monkeypatch.setattr(tracing, "flush", _flush)
|
|
return worker, source_id, flushed
|
|
|
|
def test_successful_build_is_traced(self, pg_conn, patch_worker_db, task_self, monkeypatch):
|
|
from docsgpt import tracing
|
|
|
|
def _extract(*_a, **_kw):
|
|
with tracing.span(tracing.KIND_LLM, "chat m"):
|
|
pass
|
|
return {"nodes": 3, "edges": 2, "chunks_processed": 1}
|
|
|
|
worker, source_id, flushed = self._run(pg_conn, monkeypatch, task_self, _extract)
|
|
worker.extract_graph_worker(task_self, source_id, "alice")
|
|
(trace,) = flushed
|
|
assert trace.source == "graph_extraction"
|
|
assert trace.user_id == "alice"
|
|
step, llm = trace.spans
|
|
assert step.attributes["docsgpt.graph.nodes"] == 3
|
|
assert llm.parent_id == step.id
|
|
|
|
def test_failed_build_is_an_error_trace(self, pg_conn, patch_worker_db, task_self, monkeypatch):
|
|
worker, source_id, flushed = self._run(
|
|
pg_conn, monkeypatch, task_self, MagicMock(side_effect=RuntimeError("llm down"))
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
worker.extract_graph_worker(task_self, source_id, "alice")
|
|
assert flushed[0].status == "error"
|