mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 10:14:32 +00:00
Parse documents on a Celery parsing worker via a read_document tool
Replace the sandbox Docling extractor with read_document, backed by the in-process backend parser (the same one ingestion uses) and offloaded to a dedicated 'parsing' Celery queue so it can run on GPU-capable workers with predictable RAM. The tool resolves the input ref under the run-scoped gate, enqueues the parse, and awaits it with a timeout (degrading to an error rather than hanging); the worker independently re-resolves the artifact through the same gate and never trusts a raw path. Untrusted files get the upload path's safeguards (extension whitelist, size cap, sanitized temp file, cleanup). Options: output (markdown/text/structured/chunks), ocr, pages, engine, max_chars, include_tables, persist, json_schema. The workflow native-file 'extract' fallback now uses the same worker path, so document parsing no longer needs the sandbox and works on every backend. Also fixes the branch's periodic-task test (the sandbox reaper made it 12) and points the dev and e2e Celery workers at the parsing queue.
This commit is contained in:
1 parent
273c830b55
commit
37d93cbd86
25 files changed
+1611
-1275
No files matched your search
@@ -0,0 +1,265 @@
|
||||
"""Unit tests for ReadDocumentTool: run-scoped input gate, enqueue+await, timeout/failure, schema, metadata.
|
||||
|
||||
The parse task is mocked (``parse_document.apply_async(...).get``) so no live
|
||||
worker / DB / storage is touched; these cover the pre-enqueue run-scoped gate
|
||||
(reject cross-tenant before enqueue), the await + degrade behavior, json_schema
|
||||
validation, and the surfaced action params.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
import application.agents.tools.read_document as rd
|
||||
from application.agents.tools.read_document import ReadDocumentTool
|
||||
|
||||
_ART_ID = str(uuid.uuid4())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run-scoped input resolution (mocks the repo gate)
|
||||
# ---------------------------------------------------------------------------
|
||||
def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str]):
|
||||
class _Repo:
|
||||
def __init__(self, conn):
|
||||
pass
|
||||
|
||||
def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None):
|
||||
if not found or n != 1:
|
||||
return None
|
||||
if conv is not None and conversation_id != conv:
|
||||
return None
|
||||
if run is not None and workflow_run_id != run:
|
||||
return None
|
||||
return _ART_ID
|
||||
|
||||
def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
|
||||
if not found:
|
||||
return None
|
||||
if conv is not None and conversation_id != conv:
|
||||
return None
|
||||
if run is not None and workflow_run_id != run:
|
||||
return None
|
||||
return {"id": artifact_id, "current_version": 1, "title": "statement.pdf"}
|
||||
|
||||
class _Conn:
|
||||
def __enter__(self):
|
||||
return object()
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(rd, "db_readonly", lambda: _Conn())
|
||||
monkeypatch.setattr(rd, "ArtifactsRepository", _Repo)
|
||||
|
||||
|
||||
class _FakeAsyncResult:
|
||||
def __init__(self, payload=None, exc=None):
|
||||
self._payload = payload
|
||||
self._exc = exc
|
||||
|
||||
def get(self, timeout=None):
|
||||
if self._exc is not None:
|
||||
raise self._exc
|
||||
return self._payload
|
||||
|
||||
|
||||
def _patch_task(monkeypatch, *, payload=None, exc=None):
|
||||
"""Patch parse_document.apply_async so no broker is touched; capture call args."""
|
||||
import application.api.user.tasks as tasks
|
||||
|
||||
captured: Dict[str, Any] = {}
|
||||
|
||||
def _apply_async(args=None, queue=None, **kw):
|
||||
captured["args"] = args
|
||||
captured["queue"] = queue
|
||||
return _FakeAsyncResult(payload=payload, exc=exc)
|
||||
|
||||
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
||||
return captured
|
||||
|
||||
|
||||
def _tool(**config) -> ReadDocumentTool:
|
||||
base = {"conversation_id": "conv-1", "tool_id": "t-1"}
|
||||
base.update(config)
|
||||
return ReadDocumentTool(tool_config=base, user_id="u-1")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Guards + metadata
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_unknown_action_rejected():
|
||||
out = _tool().execute_action("nope", input="a")
|
||||
assert out["status"] == "error" and "unknown action" in out["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_requires_user_and_parent():
|
||||
no_user = ReadDocumentTool({"conversation_id": "c"}, user_id=None)
|
||||
assert "user_id" in no_user.execute_action("read_document", input="a")["error"]
|
||||
no_parent = ReadDocumentTool({}, user_id="u")
|
||||
assert "conversation_id" in no_parent.execute_action("read_document", input="a")["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_input_required():
|
||||
assert "input artifact id is required" in _tool().execute_action("read_document", input=" ")["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_action_metadata_surfaces_new_params():
|
||||
meta = _tool().get_actions_metadata()[0]
|
||||
assert meta["name"] == "read_document"
|
||||
props = meta["parameters"]["properties"]
|
||||
for key in ("input", "output", "ocr", "pages", "engine", "max_chars", "include_tables", "persist", "json_schema"):
|
||||
assert key in props, key
|
||||
assert meta["parameters"]["required"] == ["input"]
|
||||
# No sandbox/Docling wording in the action description.
|
||||
assert "sandbox" not in meta["description"].lower()
|
||||
assert "docling" not in meta["description"].lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Enqueue + await happy path
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_resolves_input_then_enqueues_and_returns_payload(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
payload = {"status": "ok", "output": "markdown", "content": "# Hi", "truncated": False}
|
||||
captured = _patch_task(monkeypatch, payload=payload)
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
|
||||
|
||||
assert out["status"] == "ok"
|
||||
assert out["content"] == "# Hi"
|
||||
# The task got the resolved id, the run-scoped parent, the user, and the parsing queue.
|
||||
assert captured["args"][0] == _ART_ID
|
||||
assert captured["args"][1] == {"conversation_id": "conv-1"}
|
||||
assert captured["args"][2] == "u-1"
|
||||
assert captured["queue"] == rd.settings.DOCUMENT_PARSE_QUEUE
|
||||
options = captured["args"][3]
|
||||
assert options["output"] == "markdown" and options["persist"] is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_short_ref_input_resolves(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "x", "truncated": False})
|
||||
|
||||
out = _tool().execute_action("read_document", input="A1", persist=False)
|
||||
assert out["status"] == "ok"
|
||||
# The short ref was resolved to the real id BEFORE enqueue.
|
||||
assert captured["args"][0] == _ART_ID
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_artifact_ref_sets_last_artifact_id(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
payload = {"status": "ok", "content": "x", "truncated": False,
|
||||
"artifact": {"artifact_id": "new-art", "version": 1}}
|
||||
_patch_task(monkeypatch, payload=payload)
|
||||
|
||||
tool = _tool()
|
||||
out = tool.execute_action("read_document", input=_ART_ID)
|
||||
assert out["artifact"]["artifact_id"] == "new-art"
|
||||
assert tool.get_artifact_id("read_document") == "new-art"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-tenant: rejected BEFORE enqueue
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_cross_tenant_rejected_before_enqueue(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-OTHER", run=None)
|
||||
enqueued = {"called": False}
|
||||
|
||||
import application.api.user.tasks as tasks
|
||||
|
||||
def _apply_async(*a, **k):
|
||||
enqueued["called"] = True
|
||||
raise AssertionError("must not enqueue a cross-tenant input")
|
||||
|
||||
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
|
||||
assert out["status"] == "error" and "not found in this conversation/run" in out["error"]
|
||||
assert enqueued["called"] is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_missing_input_rejected_before_enqueue(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=False, conv="conv-1", run=None)
|
||||
import application.api.user.tasks as tasks
|
||||
monkeypatch.setattr(
|
||||
tasks.parse_document, "apply_async",
|
||||
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue")),
|
||||
)
|
||||
out = _tool().execute_action("read_document", input="ghost", persist=False)
|
||||
assert out["status"] == "error" and "not found" in out["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timeout + task failure degrade to an error result (never hang/raise)
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_timeout_degrades_to_error(monkeypatch):
|
||||
from celery.exceptions import TimeoutError as CeleryTimeoutError
|
||||
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
_patch_task(monkeypatch, exc=CeleryTimeoutError("timed out"))
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
|
||||
assert out["status"] == "error" and "timed out" in out["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_task_failure_degrades_to_error(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
_patch_task(monkeypatch, exc=RuntimeError("worker blew up"))
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
|
||||
assert out["status"] == "error" and "document parsing failed" in out["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# json_schema validation
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_json_schema_validation_passes(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
payload = {"status": "ok", "content": "x", "truncated": False,
|
||||
"structured": {"texts": [{}], "schema_name": "DoclingDocument"}}
|
||||
_patch_task(monkeypatch, payload=payload)
|
||||
schema = {"type": "object", "required": ["texts"], "properties": {"texts": {"type": "array"}}}
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, output="structured",
|
||||
json_schema=schema, persist=False)
|
||||
assert out["status"] == "ok"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_json_schema_validation_fails_cleanly(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
payload = {"status": "ok", "content": "x", "truncated": False, "structured": {"texts": [{}]}}
|
||||
_patch_task(monkeypatch, payload=payload)
|
||||
schema = {"type": "object", "required": ["amount"], "properties": {"amount": {"type": "number"}}}
|
||||
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, output="structured",
|
||||
json_schema=schema, persist=False)
|
||||
assert out["status"] == "error" and "did not match json_schema" in out["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_malformed_json_schema_rejected_before_enqueue(monkeypatch):
|
||||
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
|
||||
import application.api.user.tasks as tasks
|
||||
monkeypatch.setattr(
|
||||
tasks.parse_document, "apply_async",
|
||||
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue")),
|
||||
)
|
||||
out = _tool().execute_action("read_document", input=_ART_ID, json_schema={"properties": {}}, persist=False)
|
||||
assert out["status"] == "error" and "invalid json_schema" in out["error"]
|
||||
Reference in new issue
Block a user