tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-03 07:11:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Parse documents on a Celery parsing worker via a read_document tool

Browse Source 
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:
Alex committed 2026-06-25 13:24:12 +01:00
1 parent 273c830b55
commit 37d93cbd86
25 files changed
+1611 -1275
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.devcontainer/devc-welcome.md
+5 -1
View File
@@ -29,9 +29,13 @@ serves only the WSGI Flask app — it omits `/mcp` and the reconnect reader
### Celery (Task Queue)
```bash
celery -A application.app.celery worker -l INFO
celery -A application.app.celery worker -l INFO -Q docsgpt,parsing
```
The `parsing` queue serves document parsing (the `read_document` tool / workflow
native-file parse); without it those calls hang `DOCUMENT_PARSE_TIMEOUT` then
error. A dedicated `-Q parsing` worker can be GPU-enabled for heavier parsers.
## Github Codespaces Instructions
### 1. Make Ports Public:
AGENTS.md
+6 -2
View File
@@ -57,15 +57,19 @@ full flag set.
Run the Celery worker in a separate terminal (if needed):
```bash
celery -A application.app.celery worker -l INFO
celery -A application.app.celery worker -l INFO -Q docsgpt,parsing
```
On macOS, prefer the solo pool for Celery:
```bash
python -m celery -A application.app.celery worker -l INFO --pool=solo
python -m celery -A application.app.celery worker -l INFO --pool=solo -Q docsgpt,parsing
```
The `parsing` queue serves document parsing (the `read_document` tool / workflow
native-file parse); without it those calls hang `DOCUMENT_PARSE_TIMEOUT` then
error. A dedicated `-Q parsing` worker can be GPU-enabled for heavier parsers.
### Frontend
Install dependencies only when needed, then run the dev server:
application/agents/tools/document_extractor.py
-526
View File
@@ -1,526 +0,0 @@
"""Document Extractor tool: convert an input artifact to schema-validated structured JSON via Docling.
The ``extract_document`` action stages a parent-scoped input artifact into the
sandbox workspace, then runs a FIXED program that uses Docling (MIT) to convert
the document (pdf/docx/pptx/...) into a compact structured payload (markdown +
tables + a structured dict). The program reads its parameters from a JSON DATA
file (``json.load``) and never string-interpolates untrusted content into code,
so a malicious filename/parameter is treated as literal data, not executed. When
a ``json_schema`` is supplied the compact ``structured`` payload is validated
through the existing jsonschema path; the extracted JSON may also be persisted as
a ``data`` artifact by reference.
"""
from __future__ import annotations
import json
import logging
import uuid
from typing import Any, Dict, List, Optional
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.base import Tool
from application.core.json_schema_utils import (
JsonSchemaValidationError,
normalize_json_schema_payload,
)
from application.core.settings import settings
from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact
from application.sandbox.sandbox_creator import SandboxCreator
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
from application.storage.storage_creator import StorageCreator
from application.utils import safe_filename
logger = logging.getLogger(__name__)
try:
import jsonschema
except Exception: # pragma: no cover - jsonschema is a declared dependency
jsonschema = None # type: ignore[assignment]
# Cap the extracted payload returned to the LLM so a huge document can't flood
# context; the full extraction is still persisted as a ``data`` artifact. When the
# markdown exceeds the cap a head+tail window is kept so both the document's
# beginning AND end (e.g. totals/signatures) survive within the byte budget.
_MARKDOWN_MAX_BYTES = 8000
_MAX_TABLES_RETURNED = 20
# Per-table content caps so a single giant table can't bloat context.
_MAX_TABLE_ROWS = 50
_MAX_CELL_CHARS = 200
# Mime + kind for the persisted extraction artifact (JSON by reference).
_EXTRACT_MIME = "application/json"
_EXTRACT_KIND = "data"
# FIXED extraction program. It reads ``params.json`` (the input path + caps) as
# DATA and writes ``result.json`` to the workspace. The params are NEVER
# string-interpolated into the program; ``{params_path}``/``{result_path}`` are
# server-controlled path literals only. A missing Docling install yields a clean
# ``{"error": ...}`` result rather than a traceback.
_EXTRACT_PROGRAM = (
"import json\n"
"params = json.load(open({params_path!r}))\n"
"result_path = {result_path!r}\n"
"def _write(obj):\n"
" with open(result_path, 'w') as fh:\n"
" json.dump(obj, fh)\n"
"try:\n"
" from docling.document_converter import DocumentConverter\n"
"except Exception as exc:\n"
" _write({{'error': 'docling is not available in the sandbox runner: '\n"
" + type(exc).__name__ + ': ' + str(exc)}})\n"
" raise SystemExit(0)\n"
"src = params['input_path']\n"
"md_cap = int(params.get('markdown_max_bytes', 0)) or None\n"
"table_cap = int(params.get('max_tables', 0)) or None\n"
"try:\n"
" converter = DocumentConverter()\n"
" doc = converter.convert(src).document\n"
" markdown = doc.export_to_markdown()\n"
" structured = doc.export_to_dict()\n"
" tables = []\n"
" for tbl in getattr(doc, 'tables', []) or []:\n"
" try:\n"
" df = tbl.export_to_dataframe()\n"
" tables.append({{'columns': [str(c) for c in df.columns],\n"
" 'rows': df.astype(str).values.tolist()}})\n"
" except Exception:\n"
" try:\n"
" tables.append({{'markdown': tbl.export_to_markdown()}})\n"
" except Exception:\n"
" continue\n"
" if table_cap is not None and len(tables) >= table_cap:\n"
" break\n"
" page_count = len(getattr(doc, 'pages', {{}}) or {{}})\n"
" md_truncated = False\n"
" if md_cap is not None and len(markdown) > md_cap:\n"
" head = md_cap // 2\n"
" tail = md_cap - head\n"
" dropped = len(markdown) - head - tail\n"
" markdown = (markdown[:head]\n"
" + '\\n\\n...[truncated ' + str(dropped) + ' chars]...\\n\\n'\n"
" + markdown[-tail:])\n"
" md_truncated = True\n"
" _write({{'ok': True, 'markdown': markdown, 'markdown_truncated': md_truncated,\n"
" 'tables': tables, 'page_count': page_count, 'structured': structured}})\n"
"except Exception as exc:\n"
" _write({{'error': 'docling extraction failed: ' + type(exc).__name__ + ': ' + str(exc)}})\n"
)
def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
"""Bound text to a head+tail byte window so a large file can't flood context."""
cap = int(max_bytes or _MARKDOWN_MAX_BYTES)
if cap <= 0:
return text
encoded = text.encode("utf-8")
if len(encoded) <= cap:
return text
head = cap // 2
tail = cap - head
dropped = len(encoded) - head - tail
head_text = encoded[:head].decode("utf-8", errors="ignore")
tail_text = encoded[-tail:].decode("utf-8", errors="ignore")
return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
def extract_markdown_from_bytes(
data: bytes,
filename: str,
session_id: str,
*,
markdown_max_bytes: Optional[int] = None,
) -> Optional[str]:
"""Run the fixed Docling program in a sandbox session and return the document's markdown, or None.
The bytes ride in as a DATA file the fixed program reads; nothing untrusted is
interpolated into the program. Returns ``None`` when Docling is unavailable or
extraction fails (the caller decides how to degrade).
"""
safe_name = safe_filename(filename) or "document"
token = uuid.uuid4().hex
token_dir = f"extract/{token}"
input_path = f"{token_dir}/inputs/{safe_name}"
params_path = f"{token_dir}/params.json"
result_path = f"{token_dir}/result.json"
params = {
"input_path": input_path,
"markdown_max_bytes": int(markdown_max_bytes or _MARKDOWN_MAX_BYTES),
"max_tables": 0,
}
program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
manager = SandboxCreator.get_manager()
try:
manager.open(session_id, ttl=timeout)
except Exception:
logger.exception("extract_markdown_from_bytes: failed to open sandbox session")
return None
try:
manager.put_file(session_id, input_path, data)
manager.put_file(session_id, params_path, json.dumps(params).encode("utf-8"))
result = manager.exec(session_id, program, timeout=timeout)
if not result.ok:
return None
raw = manager.get_file(session_id, result_path)
except Exception:
logger.exception("extract_markdown_from_bytes: extraction failed")
return None
finally:
manager.remove_path(session_id, token_dir)
try:
manager.close(session_id)
except Exception:
logger.exception("extract_markdown_from_bytes: session close failed")
if not raw:
return None
try:
extracted = json.loads(raw.decode("utf-8"))
except (ValueError, UnicodeDecodeError):
return None
if not isinstance(extracted, dict) or extracted.get("error") or not extracted.get("ok"):
return None
markdown = extracted.get("markdown")
return markdown if isinstance(markdown, str) else None
class DocumentExtractorTool(Tool):
"""Document Extractor
Convert an input document artifact (pdf/docx/pptx/...) to compact, schema-validated structured JSON via Docling.
"""
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
"""Bind the tool to the invoker and its conversation/run-scoped sandbox session."""
self.config: Dict[str, Any] = tool_config or {}
self.user_id: Optional[str] = user_id
self.tool_id: Optional[str] = self.config.get("tool_id")
self.conversation_id: Optional[str] = self.config.get("conversation_id")
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
self._last_artifact_id: Optional[str] = None
# ------------------------------------------------------------------
# Tool ABC
# ------------------------------------------------------------------
def get_actions_metadata(self) -> List[Dict[str, Any]]:
"""Return JSON metadata describing the ``extract_document`` action for tool schemas."""
return [
{
"name": "extract_document",
"description": (
"Extract a document artifact (pdf/docx/pptx/...) into compact structured JSON "
"(markdown + tables + structure) using Docling. Optionally validate the result "
"against a json_schema and persist it as a downloadable data artifact."
),
"active": True,
"parameters": {
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Document (from this conversation/run) to extract; accepts the short "
"ref like `A1` returned by a previous artifact action, or the full artifact id.",
},
"json_schema": {
"type": "object",
"description": "Optional JSON schema the extracted 'structured' payload must satisfy.",
},
"persist": {
"type": "boolean",
"description": "Persist the extracted JSON as a downloadable data artifact (default true).",
},
},
"required": ["input"],
},
}
]
def get_config_requirements(self) -> Dict[str, Any]:
"""Return configuration requirements (none beyond the deployment sandbox backend)."""
return {}
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
"""Return the persisted extraction artifact id so the UI artifact rail lights up."""
return self._last_artifact_id
# ------------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------------
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
"""Dispatch a tool action; only ``extract_document`` is supported."""
self._last_artifact_id = None
if action_name != "extract_document":
return {"status": "error", "error": f"unknown action: {action_name}"}
if not self.user_id:
return {"status": "error", "error": "document_extractor requires a valid user_id."}
if self.conversation_id is None and self.workflow_run_id is None:
return {"status": "error", "error": "document_extractor requires a conversation_id or workflow_run_id."}
return self._extract(**kwargs)
# ------------------------------------------------------------------
# Extraction
# ------------------------------------------------------------------
def _extract(self, **kwargs: Any) -> Dict[str, Any]:
"""Stage the parent-scoped input, run the fixed Docling program, validate, and (optionally) persist."""
input_id = kwargs.get("input")
json_schema = kwargs.get("json_schema")
should_persist = kwargs.get("persist", True)
if not isinstance(input_id, str) or not input_id.strip():
return {"status": "error", "error": "input artifact id is required."}
if json_schema is not None:
schema_err = self._check_schema(json_schema)
if schema_err is not None:
return schema_err
session_id = self._resolve_session_id()
if session_id is None:
return {"status": "error", "error": "document_extractor requires a conversation_id or workflow_run_id."}
loaded = self._load_input(input_id.strip())
if loaded.get("error"):
return {"status": "error", "error": loaded["error"]}
max_input = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
if len(loaded["data"]) > max_input:
return {
"status": "error",
"error": f"input artifact is too large: {len(loaded['data'])} bytes exceeds the "
f"{max_input}-byte sandbox input cap.",
}
token = uuid.uuid4().hex
token_dir = f"extract/{token}"
input_path = f"{token_dir}/inputs/{loaded['filename']}"
params_path = f"{token_dir}/params.json"
result_path = f"{token_dir}/result.json"
params = {
"input_path": input_path,
"markdown_max_bytes": _MARKDOWN_MAX_BYTES,
"max_tables": _MAX_TABLES_RETURNED,
}
program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
manager = SandboxCreator.get_manager()
try:
manager.open(session_id, ttl=timeout)
except Exception as exc:
logger.exception("document_extractor: failed to open sandbox session")
return {"status": "error", "error": f"sandbox unavailable: {type(exc).__name__}: {exc}"}
try:
# The document bytes and the params ride in as DATA files the program
# reads; neither is interpolated into the program, so a hostile
# filename or document content stays inert data.
manager.put_file(session_id, input_path, loaded["data"])
manager.put_file(session_id, params_path, json.dumps(params).encode("utf-8"))
result = manager.exec(session_id, program, timeout=timeout)
if not result.ok:
detail = (
f"{result.error_name}: {result.error_value}"
if result.error_name
else (result.error_value or "extraction failed")
)
return {"status": "error", "error": f"extraction failed: {detail}"}
raw = manager.get_file(session_id, result_path)
except Exception as exc:
logger.exception("document_extractor: extraction failed")
return {"status": "error", "error": f"extraction failed: {type(exc).__name__}: {exc}"}
finally:
# Drop this extraction's scratch dir (staged input + params + result)
# before close so a warm/reused session doesn't accumulate on disk.
manager.remove_path(session_id, token_dir)
try:
manager.close(session_id)
except Exception:
logger.exception("document_extractor: session close failed")
return self._finish(raw, loaded, json_schema, should_persist)
def _finish(
self,
raw: bytes,
loaded: Dict[str, Any],
json_schema: Any,
should_persist: Any,
) -> Dict[str, Any]:
"""Parse the program result, validate against json_schema, and shape the compact payload."""
if not raw:
return {"status": "error", "error": "extractor produced no result."}
try:
extracted = json.loads(raw.decode("utf-8"))
except (ValueError, UnicodeDecodeError):
return {"status": "error", "error": "extractor produced an invalid result."}
if isinstance(extracted, dict) and extracted.get("error"):
return {"status": "error", "error": str(extracted["error"])}
if not isinstance(extracted, dict) or not extracted.get("ok"):
return {"status": "error", "error": "extractor produced an unexpected result."}
structured = extracted.get("structured")
if json_schema is not None:
valid = self._validate(json_schema, structured)
if valid is not None:
return valid
compact = self._compact_payload(extracted)
payload: Dict[str, Any] = {"status": "ok", "structured": compact}
if should_persist:
try:
ref = self._persist(extracted, loaded["title"])
except QuotaExceeded as exc:
# Extraction itself succeeded; surface the quota error alongside
# the in-context structured result rather than failing the call.
payload["artifact_error"] = str(exc)
ref = None
if ref is not None:
self._last_artifact_id = ref["artifact_id"]
payload["artifact"] = ref
return payload
# ------------------------------------------------------------------
# Input / payload helpers
# ------------------------------------------------------------------
def _load_input(self, raw_id: str) -> Dict[str, Any]:
"""Resolve a short ref/uuid, then fetch the parent-scoped input bytes; never cross-tenant."""
artifact_id: Optional[str] = raw_id
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
# A ref (A1/A2/...) resolves to an id within this parent only; the
# resolved id is re-checked through the parent-scoped gate so a ref
# can never reach another tenant.
artifact_id = resolve_artifact_id(
repo,
raw_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
artifact = (
repo.get_artifact_in_parent(
artifact_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact_id is not None
else None
)
if artifact is None:
return {"error": f"input artifact {raw_id} not found in this conversation/run."}
version = repo.get_version(artifact_id, artifact["current_version"])
except Exception:
logger.exception("document_extractor: failed to load input artifact")
return {"error": f"failed to load input artifact {raw_id}."}
if not version or not version.get("storage_path"):
return {"error": f"input artifact {raw_id} has no stored content."}
display_name = version.get("filename") or artifact.get("title") or artifact_id
filename = safe_filename(display_name)
try:
file_obj = StorageCreator.get_storage().get_file(version["storage_path"])
data = file_obj.read()
except Exception:
logger.exception("document_extractor: failed to read input artifact bytes")
return {"error": f"failed to read input artifact {raw_id}."}
return {"data": data, "filename": filename, "title": display_name}
def _compact_payload(self, extracted: Dict[str, Any]) -> Dict[str, Any]:
"""Shape the LLM-facing payload: markdown window + bounded tables + structure summary, no raw bytes."""
structured = extracted.get("structured")
return {
"markdown": extracted.get("markdown", ""),
"markdown_truncated": bool(extracted.get("markdown_truncated")),
"tables": [self._compact_table(tbl) for tbl in extracted.get("tables", [])],
"page_count": extracted.get("page_count", 0),
"summary": self._structure_summary(structured),
}
@staticmethod
def _compact_table(table: Any) -> Any:
"""Bound a single table's rows and cell sizes so one giant table can't bloat context."""
if not isinstance(table, dict):
return table
def _cell(value: Any) -> Any:
if isinstance(value, str) and len(value) > _MAX_CELL_CHARS:
return value[:_MAX_CELL_CHARS] + "...[truncated]"
return value
rows = table.get("rows")
if not isinstance(rows, list):
return table
capped = [[_cell(c) for c in row] if isinstance(row, list) else _cell(row) for row in rows[:_MAX_TABLE_ROWS]]
compact = dict(table)
compact["rows"] = capped
if len(rows) > _MAX_TABLE_ROWS:
compact["rows_truncated"] = True
compact["total_rows"] = len(rows)
return compact
@staticmethod
def _structure_summary(structured: Any) -> Dict[str, Any]:
"""Summarize the Docling structured dict by top-level element counts (keeps context compact)."""
if not isinstance(structured, dict):
return {}
counts: Dict[str, int] = {}
for key in ("texts", "tables", "pictures", "groups", "pages"):
value = structured.get(key)
if isinstance(value, (list, dict)):
counts[key] = len(value)
return counts
def _persist(self, extracted: Dict[str, Any], title: str) -> Optional[Dict[str, Any]]:
"""Persist the full extraction JSON as a ``data`` artifact by reference; return its reference."""
try:
data = json.dumps(extracted).encode("utf-8")
except (TypeError, ValueError):
logger.exception("document_extractor: extraction is not JSON-serializable")
return None
filename = f"{safe_filename(title) or 'extract'}.extract.json"
ref = persist_new_artifact(
user_id=self.user_id,
kind=_EXTRACT_KIND,
data=data,
filename=filename,
mime_type=_EXTRACT_MIME,
title=f"{title} (extracted)",
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
produced_by={"tool": "document_extractor", "action": "extract_document", "tool_id": self.tool_id},
)
return ref
# ------------------------------------------------------------------
# Schema validation
# ------------------------------------------------------------------
@staticmethod
def _check_schema(json_schema: Any) -> Optional[Dict[str, Any]]:
"""Return an error payload when ``json_schema`` itself is malformed, else None."""
try:
normalize_json_schema_payload(json_schema)
except JsonSchemaValidationError as exc:
return {"status": "error", "error": f"invalid json_schema: {exc}"}
return None
@staticmethod
def _validate(json_schema: Any, instance: Any) -> Optional[Dict[str, Any]]:
"""Validate ``instance`` against the (already-normalized) json_schema; error payload on mismatch."""
if jsonschema is None:
return {"status": "error", "error": "jsonschema is required for json_schema validation."}
schema = normalize_json_schema_payload(json_schema)
try:
jsonschema.validate(instance=instance, schema=schema)
except jsonschema.exceptions.ValidationError as exc:
return {"status": "error", "error": f"extracted structure did not match json_schema: {exc.message}"}
return None
# ------------------------------------------------------------------
# Misc helpers
# ------------------------------------------------------------------
def _resolve_session_id(self) -> Optional[str]:
"""Derive the sandbox session id from the bound conversation/run; sanitize to the gateway charset."""
raw = self.conversation_id or self.workflow_run_id
if not raw:
return None
sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in str(raw))
return sanitized or None
application/agents/tools/read_document.py
+268
View File
@@ -0,0 +1,268 @@
"""Read Document tool: parse an input artifact to text/markdown/structured/chunks via the backend parser.
The ``read_document`` action resolves a parent-scoped input artifact, enqueues a
``parse_document`` task on the dedicated ``parsing`` Celery queue, and awaits the
result with a timeout. The run-scoped authz gate is enforced TWICE — here before
enqueue (reject cross-tenant) and again in the worker (re-resolve, never trusting a
raw path). When a ``json_schema`` is supplied the structured payload is validated
through the existing jsonschema path; the full result may also be persisted as a
``data`` artifact by reference (handled in the worker).
"""
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.agents.tools.base import Tool
from application.core.json_schema_utils import (
JsonSchemaValidationError,
normalize_json_schema_payload,
)
from application.core.settings import settings
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
logger = logging.getLogger(__name__)
try:
import jsonschema
except Exception: # pragma: no cover - jsonschema is a declared dependency
jsonschema = None # type: ignore[assignment]
class ReadDocumentTool(Tool):
"""Read Document
Parse an input document artifact (pdf/docx/pptx/...) to text/markdown/structured JSON via the backend parser.
"""
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
"""Bind the tool to the invoker and its conversation/run scope."""
self.config: Dict[str, Any] = tool_config or {}
self.user_id: Optional[str] = user_id
self.tool_id: Optional[str] = self.config.get("tool_id")
self.conversation_id: Optional[str] = self.config.get("conversation_id")
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
self._last_artifact_id: Optional[str] = None
# ------------------------------------------------------------------
# Tool ABC
# ------------------------------------------------------------------
def get_actions_metadata(self) -> List[Dict[str, Any]]:
"""Return JSON metadata describing the ``read_document`` action for tool schemas."""
return [
{
"name": "read_document",
"description": (
"Read a document artifact (pdf/docx/pptx/...) and return its parsed content as "
"markdown, plain text, structured JSON (with tables), or chunks. Optionally "
"validate the structured result against a json_schema and persist it as a "
"downloadable data artifact."
),
"active": True,
"parameters": {
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Document (from this conversation/run) to read; accepts the short "
"ref like `A1` returned by a previous artifact action, or the full artifact id.",
},
"output": {
"type": "string",
"enum": ["markdown", "text", "structured", "chunks"],
"description": "Shape of the parsed result (default: markdown). Note: "
"`structured` always uses the Docling engine regardless of `engine` "
"(the `fast` engine is markdown/text only).",
},
"ocr": {
"type": "string",
"enum": ["auto", "on", "off"],
"description": "OCR mode for scanned pages/images (default: auto, follows server config).",
},
"pages": {
"type": "string",
"description": "Optional page range to read, e.g. `1-3` or `2` (best-effort).",
},
"engine": {
"type": "string",
"enum": ["auto", "docling", "fast"],
"description": "Parser engine (default: auto). `fast` is a lighter "
"markdown/text-only engine; it is ignored when `output='structured'`, "
"which always uses Docling.",
},
"max_chars": {
"type": "integer",
"description": "Optional cap on returned characters.",
},
"include_tables": {
"type": "boolean",
"description": "Include extracted tables in the result (default: true).",
},
"json_schema": {
"type": "object",
"description": "Optional JSON schema the structured payload must satisfy.",
},
"persist": {
"type": "boolean",
"description": "Persist the parsed result as a downloadable data artifact (default true).",
},
},
"required": ["input"],
},
}
]
def get_config_requirements(self) -> Dict[str, Any]:
"""Return configuration requirements (none beyond a running parsing worker)."""
return {}
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
"""Return the persisted parse artifact id so the UI artifact rail lights up."""
return self._last_artifact_id
# ------------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------------
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
"""Dispatch a tool action; only ``read_document`` is supported."""
self._last_artifact_id = None
if action_name != "read_document":
return {"status": "error", "error": f"unknown action: {action_name}"}
if not self.user_id:
return {"status": "error", "error": "read_document requires a valid user_id."}
if self.conversation_id is None and self.workflow_run_id is None:
return {"status": "error", "error": "read_document requires a conversation_id or workflow_run_id."}
return self._read(**kwargs)
# ------------------------------------------------------------------
# Read
# ------------------------------------------------------------------
def _read(self, **kwargs: Any) -> Dict[str, Any]:
"""Resolve the input run-scoped (reject cross-tenant before enqueue), enqueue+await, validate."""
input_id = kwargs.get("input")
json_schema = kwargs.get("json_schema")
if not isinstance(input_id, str) or not input_id.strip():
return {"status": "error", "error": "input artifact id is required."}
if json_schema is not None:
schema_err = self._check_schema(json_schema)
if schema_err is not None:
return schema_err
artifact_id = self._resolve_input(input_id.strip())
if isinstance(artifact_id, dict):
return artifact_id # error payload
options = {
"output": kwargs.get("output", "markdown"),
"ocr": kwargs.get("ocr", "auto"),
"pages": kwargs.get("pages"),
"engine": kwargs.get("engine", "auto"),
"max_chars": kwargs.get("max_chars"),
"include_tables": kwargs.get("include_tables", True),
"persist": kwargs.get("persist", True),
"tool_id": self.tool_id,
}
result = self._dispatch(artifact_id, options)
if result.get("status") == "error":
return result
if json_schema is not None:
valid = self._validate(json_schema, result.get("structured"))
if valid is not None:
return valid
artifact = result.get("artifact")
if isinstance(artifact, dict) and artifact.get("artifact_id"):
self._last_artifact_id = artifact["artifact_id"]
return result
def _dispatch(self, artifact_id: str, options: Dict[str, Any]) -> Dict[str, Any]:
"""Enqueue ``parse_document`` on the dedicated parsing queue and await with a timeout.
Deadlock note: this tool may run in the WEB process (/stream) OR inside a
Celery worker (headless/scheduled agents). Routing to the dedicated
``parsing`` queue + separate parsing workers avoids worker-self-deadlock;
the await degrades (returns an error) on timeout/failure rather than hanging.
"""
from celery.exceptions import TimeoutError as CeleryTimeoutError
from application.api.user.tasks import parse_document
parent = self._parent()
timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
try:
async_result = parse_document.apply_async(
args=[artifact_id, parent, self.user_id, options], queue=queue
)
result = async_result.get(timeout=timeout)
except (CeleryTimeoutError, TimeoutError):
return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."}
except Exception as exc:
logger.exception("read_document: parse task failed")
return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
if not isinstance(result, dict):
return {"status": "error", "error": "document parsing produced an unexpected result."}
return result
# ------------------------------------------------------------------
# Input resolution (run-scoped gate, before enqueue)
# ------------------------------------------------------------------
def _resolve_input(self, raw_id: str) -> Any:
"""Resolve a short ref/uuid to a parent-scoped artifact id; an error dict on miss/cross-tenant."""
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
artifact_id = resolve_artifact_id(
repo,
raw_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
artifact = (
repo.get_artifact_in_parent(
artifact_id,
conversation_id=self.conversation_id,
workflow_run_id=self.workflow_run_id,
)
if artifact_id is not None
else None
)
except Exception:
logger.exception("read_document: failed to resolve input artifact")
return {"status": "error", "error": f"failed to resolve input artifact {raw_id}."}
if artifact is None:
return {"status": "error", "error": f"input artifact {raw_id} not found in this conversation/run."}
return str(artifact_id)
def _parent(self) -> Dict[str, Any]:
"""Build the run-scoped parent dict passed to the worker for its independent re-resolve."""
if self.conversation_id is not None:
return {"conversation_id": self.conversation_id}
return {"workflow_run_id": self.workflow_run_id}
# ------------------------------------------------------------------
# Schema validation
# ------------------------------------------------------------------
@staticmethod
def _check_schema(json_schema: Any) -> Optional[Dict[str, Any]]:
"""Return an error payload when ``json_schema`` itself is malformed, else None."""
try:
normalize_json_schema_payload(json_schema)
except JsonSchemaValidationError as exc:
return {"status": "error", "error": f"invalid json_schema: {exc}"}
return None
@staticmethod
def _validate(json_schema: Any, instance: Any) -> Optional[Dict[str, Any]]:
"""Validate ``instance`` against the (already-normalized) json_schema; error payload on mismatch."""
if jsonschema is None:
return {"status": "error", "error": "jsonschema is required for json_schema validation."}
if instance is None:
return {"status": "error", "error": "json_schema validation requires output='structured'."}
schema = normalize_json_schema_payload(json_schema)
try:
jsonschema.validate(instance=instance, schema=schema)
except jsonschema.exceptions.ValidationError as exc:
return {"status": "error", "error": f"parsed structure did not match json_schema: {exc.message}"}
return None
application/agents/tools/tool_manager.py
+2 -2
View File
@@ -39,7 +39,7 @@ class ToolManager:
"remote_device",
"code_executor",
"artifact_generator",
"document_extractor",
"read_document",
}
and user_id
):
@@ -61,7 +61,7 @@ class ToolManager:
"remote_device",
"code_executor",
"artifact_generator",
"document_extractor",
"read_document",
}
and user_id
):
application/agents/workflows/workflow_engine.py
+53 -22
View File
@@ -566,7 +566,9 @@ class WorkflowEngine:
attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
native_count += 1
else:
content = self._extract_attachment_text(storage_path, mime_type, filename, max_bytes)
content = self._extract_attachment_text(
artifact_id, storage_path, mime_type, filename, max_bytes
)
if content is None:
logger.warning(
"Workflow node %s: could not extract text from %s; skipping",
@@ -598,37 +600,66 @@ class WorkflowEngine:
return mime_type == "application/pdf" and supports_images
def _extract_attachment_text(
self, storage_path: str, mime_type: str, filename: str, max_bytes: int
self, artifact_id: str, storage_path: str, mime_type: str, filename: str, max_bytes: int
) -> Optional[str]:
"""Get an attachment's text: inline already-text formats, else extract via Docling; None on failure."""
from application.agents.tools.document_extractor import (
extract_markdown_from_bytes,
truncate_text_head_tail,
)
"""Get an attachment's text: inline already-text formats, else parse via the parsing worker; None on failure."""
from application.parser.document_reader import truncate_text_head_tail
from application.storage.storage_creator import StorageCreator
try:
data = StorageCreator.get_storage().get_file(storage_path).read()
except Exception:
logger.exception("Workflow node: failed to read document bytes for extraction")
return None
# Defensive size gate: a NULL/missing version size skips the pre-read cap,
# so re-check the actual bytes before inlining or sandbox-extracting.
if len(data) > max_bytes:
logger.warning(
"Workflow node: document at %s (%d bytes) exceeds the %d-byte cap; skipping",
storage_path, len(data), max_bytes,
)
return None
if self._is_inline_text_mime(mime_type):
try:
data = StorageCreator.get_storage().get_file(storage_path).read()
except Exception:
logger.exception("Workflow node: failed to read document bytes for extraction")
return None
# Defensive size gate: a NULL/missing version size skips the pre-read cap,
# so re-check the actual bytes before inlining.
if len(data) > max_bytes:
logger.warning(
"Workflow node: document at %s (%d bytes) exceeds the %d-byte cap; skipping",
storage_path, len(data), max_bytes,
)
return None
try:
text = data.decode("utf-8", errors="replace")
except Exception:
return None
# Bound the inlined text to a head+tail window so a large-but-under-cap
# text file can't blow the context (the Docling branch is already bounded).
# text file can't blow the context (the parse branch is already bounded).
return truncate_text_head_tail(text)
return extract_markdown_from_bytes(data, filename, self._session_id())
# Non-text mimes parse via the dedicated parsing queue (works on any backend,
# no sandbox): the worker re-resolves the artifact run-scoped and reads its bytes.
return self._parse_document_text(artifact_id)
def _parse_document_text(self, artifact_id: str) -> Optional[str]:
"""Enqueue ``parse_document`` for this run and await the bounded markdown; None on failure."""
from celery.exceptions import TimeoutError as CeleryTimeoutError
from application.api.user.tasks import parse_document
from application.core.settings import settings
user_id = self._resolve_user_id()
if not user_id:
return None
options = {"output": "markdown", "include_tables": False, "persist": False}
queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
try:
async_result = parse_document.apply_async(
args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options],
queue=queue,
)
result = async_result.get(timeout=timeout)
except (CeleryTimeoutError, TimeoutError):
logger.warning("Workflow node: document parse timed out for %s", artifact_id)
return None
except Exception:
logger.exception("Workflow node: document parse failed")
return None
if isinstance(result, dict) and result.get("status") == "ok":
content = result.get("content")
return content if isinstance(content, str) else None
return None
@staticmethod
def _is_inline_text_mime(mime_type: str) -> bool:
application/api/user/tasks.py
+11
View File
@@ -8,6 +8,7 @@ from application.worker import (
attachment_worker,
ingest_worker,
mcp_oauth,
parse_document_worker,
reembed_wiki_page_worker,
remote_worker,
sync,
@@ -246,6 +247,16 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None):
return resp
# Not DURABLE: the read_document tool awaits this synchronously with a timeout, so a
# blind autoretry would double-parse and the caller would already have degraded. The
# task is routed to the dedicated ``parsing`` queue (celeryconfig task_routes) so a
# parse enqueued from inside a Celery worker (headless/scheduled agent) is served by a
# separate parsing worker and never self-deadlocks the awaiting worker.
@celery.task(bind=True, acks_late=False, autoretry_for=())
def parse_document(self, artifact_id, parent, user_id, options=None):
return parse_document_worker(self, artifact_id, parent, user_id, options or {})
@celery.task(**DURABLE_TASK)
@with_idempotency(
task_name="ingest_connector_task", on_poison=_emit_ingest_poison_event,
application/celeryconfig.py
+8
View File
@@ -19,6 +19,14 @@ task_default_queue = "docsgpt"
task_default_exchange = "docsgpt"
task_default_routing_key = "docsgpt"
# Route document parsing to a dedicated queue so a parse enqueued from inside a
# Celery worker (headless/scheduled agent) is served by a separate parsing worker
# and never self-deadlocks the awaiting worker. The tool also passes the queue at
# apply_async time, so this routing is the default for any other enqueuer.
task_routes = {
"application.api.user.tasks.parse_document": {"queue": settings.DOCUMENT_PARSE_QUEUE},
}
beat_scheduler = "redbeat.RedBeatScheduler"
redbeat_redis_url = broker_url
redbeat_key_prefix = "redbeat:docsgpt:"
application/core/settings.py
+4
View File
@@ -343,6 +343,10 @@ class Settings(BaseSettings):
SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024 # cap on buffered stdout+stderr per exec
SANDBOX_MAX_FILE_BYTES: int = 10 * 1024 * 1024 # cap on get_file size routed through stdout
SANDBOX_MAX_INPUT_BYTES: int = 25 * 1024 * 1024 # cap on an input document staged into a sandbox session
# ``read_document`` parsing on a dedicated Celery ``parsing`` queue (backend parser).
DOCUMENT_PARSE_QUEUE: str = "parsing" # queue the parse_document task is routed to
DOCUMENT_PARSE_TIMEOUT: int = 120 # seconds the tool awaits the enqueued parse before degrading
DOCUMENT_PARSE_MAX_BYTES: int = 0 # cap on a parsed document's bytes (0 = reuse SANDBOX_MAX_INPUT_BYTES)
# Per-agent-node cap on files passed natively to the node's LLM (vision/doc
# inputs). Files past the cap are extracted to text or dropped, not attached
# natively, to bound context/cost. Re-uses SANDBOX_MAX_INPUT_BYTES per file.
application/parser/document_reader.py
+357
View File
@@ -0,0 +1,357 @@
"""In-process document parsing for the ``read_document`` tool, run on the Celery parsing worker.
``parse_document_bytes`` turns untrusted document bytes into a bounded, shaped
result (markdown/text/structured/chunks) using the BACKEND parsers (Docling by
default). It applies the same untrusted-content safeguards as uploads — an
extension whitelist, a byte cap, ``safe_filename`` staging into a temp file, and
temp cleanup — so a hostile filename or document is treated as inert data.
"""
from __future__ import annotations
import logging
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional
from application.core.settings import settings
from application.parser.file.bulk import get_default_file_extractor
from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
from application.utils import safe_filename
logger = logging.getLogger(__name__)
# Cap the text returned to the LLM so a huge document can't flood context; the
# full result is still persisted as a ``data`` artifact. When the text exceeds
# the cap a head+tail window keeps both the document's beginning AND end (e.g.
# totals/signatures) within the byte budget.
_TEXT_MAX_BYTES = 8000
_MAX_TABLES_RETURNED = 20
_MAX_TABLE_ROWS = 50
_MAX_CELL_CHARS = 200
# Caps applied to the bounded view that rides back through the Redis result
# backend (the full result still lives in the persisted artifact).
_MAX_CHUNKS_RETURNED = 50
_VALID_OUTPUTS = ("markdown", "text", "structured", "chunks")
_VALID_OCR = ("auto", "on", "off")
_VALID_ENGINES = ("auto", "docling", "fast")
def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
"""Bound text to a head+tail byte window so a large file can't flood context."""
cap = int(max_bytes or _TEXT_MAX_BYTES)
if cap <= 0:
return text
encoded = text.encode("utf-8")
if len(encoded) <= cap:
return text
head = cap // 2
tail = cap - head
dropped = len(encoded) - head - tail
head_text = encoded[:head].decode("utf-8", errors="ignore")
tail_text = encoded[-tail:].decode("utf-8", errors="ignore")
return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
def bound_parse_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
"""Bound every shape of a parse payload so the Redis-backed result stays small.
``content`` is re-windowed and ``chunks`` is capped in count and per-chunk
length. ``structured`` is left as-is: it rides back so json_schema validation
in the tool can run against it, and it is already bounded by the input byte
cap plus the table caps (``_compact_table`` / ``summary``); the full result is
also persisted as a ``data`` artifact. The dict is mutated in place.
"""
content = payload.get("content")
if isinstance(content, str):
payload["content"] = truncate_text_head_tail(content)
chunks = payload.get("chunks")
if isinstance(chunks, list):
bounded = [
truncate_text_head_tail(chunk) if isinstance(chunk, str) else chunk
for chunk in chunks[:_MAX_CHUNKS_RETURNED]
]
if len(chunks) > _MAX_CHUNKS_RETURNED:
payload["chunks_truncated"] = True
payload["total_chunks"] = len(chunks)
payload["chunks"] = bounded
return payload
def _max_input_bytes() -> int:
"""Return the size cap for a parsed document (its own setting, else the sandbox cap)."""
explicit = int(getattr(settings, "DOCUMENT_PARSE_MAX_BYTES", 0) or 0)
if explicit > 0:
return explicit
return int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
def _resolve_ocr_enabled(ocr: str) -> bool:
"""Resolve the OCR flag from the ``ocr`` arg and the deployment setting."""
if ocr == "on":
return True
if ocr == "off":
return False
return bool(getattr(settings, "DOCLING_OCR_ENABLED", False))
def _pick_parser(suffix: str, *, ocr_enabled: bool, engine: str):
"""Select the parser for ``suffix`` honoring the requested engine; None when unsupported."""
if engine == "fast":
legacy = _legacy_parser_for(suffix)
if legacy is not None:
return legacy
extractor = get_default_file_extractor(ocr_enabled=ocr_enabled)
return extractor.get(suffix)
def _legacy_parser_for(suffix: str):
"""Return a non-Docling parser for ``suffix`` (the ``fast`` engine), or None."""
from application.parser.file.docs_parser import DocxParser, PDFParser
from application.parser.file.html_parser import HTMLParser
from application.parser.file.markdown_parser import MarkdownParser
from application.parser.file.tabular_parser import ExcelParser, PandasCSVParser
legacy = {
".pdf": PDFParser,
".docx": DocxParser,
".csv": PandasCSVParser,
".xlsx": ExcelParser,
".html": HTMLParser,
".md": MarkdownParser,
".mdx": MarkdownParser,
}
cls = legacy.get(suffix)
return cls() if cls is not None else None
def _parse_to_text(parser: Any, path: Path) -> str:
"""Run a parser and coerce its ``str | List[str]`` result to a single text blob."""
if not parser.parser_config_set:
parser.init_parser()
parsed = parser.parse_file(path, errors="ignore")
if isinstance(parsed, list):
return "\n\n".join(str(part) for part in parsed)
return str(parsed)
def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]:
"""Bound a single table's rows and cell sizes so one giant table can't bloat context."""
def _cell(value: Any) -> Any:
if isinstance(value, str) and len(value) > _MAX_CELL_CHARS:
return value[:_MAX_CELL_CHARS] + "...[truncated]"
return value
rows = table.get("rows")
if not isinstance(rows, list):
return table
capped = [[_cell(c) for c in row] if isinstance(row, list) else _cell(row) for row in rows[:_MAX_TABLE_ROWS]]
compact = dict(table)
compact["rows"] = capped
if len(rows) > _MAX_TABLE_ROWS:
compact["rows_truncated"] = True
compact["total_rows"] = len(rows)
return compact
def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool) -> Dict[str, Any]:
"""Convert a document with Docling and return markdown + structured dict + bounded tables."""
from docling.document_converter import DocumentConverter
converter = DocumentConverter()
doc = converter.convert(str(path)).document
markdown = doc.export_to_markdown()
structured = doc.export_to_dict()
tables: List[Dict[str, Any]] = []
if include_tables:
for tbl in getattr(doc, "tables", []) or []:
try:
df = tbl.export_to_dataframe()
tables.append({"columns": [str(c) for c in df.columns], "rows": df.astype(str).values.tolist()})
except Exception:
try:
tables.append({"markdown": tbl.export_to_markdown()})
except Exception:
continue
if len(tables) >= _MAX_TABLES_RETURNED:
break
page_count = len(getattr(doc, "pages", {}) or {})
return {"markdown": markdown, "structured": structured, "tables": tables, "page_count": page_count}
def _structure_summary(structured: Any) -> Dict[str, int]:
"""Summarize the Docling structured dict by top-level element counts (keeps context compact)."""
if not isinstance(structured, dict):
return {}
counts: Dict[str, int] = {}
for key in ("texts", "tables", "pictures", "groups", "pages"):
value = structured.get(key)
if isinstance(value, (list, dict)):
counts[key] = len(value)
return counts
def _apply_pages(text: str, pages: Any) -> str:
"""Best-effort page-range slice on a page-delimited markdown blob (``\\f`` separated)."""
if not pages:
return text
parts = text.split("\f")
if len(parts) <= 1:
return text
selected = _selected_page_indices(pages, len(parts))
if not selected:
return text
return "\f".join(parts[i] for i in selected if 0 <= i < len(parts))
def _selected_page_indices(pages: Any, total: int) -> List[int]:
"""Parse ``pages`` ("1-3", "2", [1,2]) into 0-based indices bounded by ``total``."""
indices: List[int] = []
tokens = pages if isinstance(pages, list) else str(pages).split(",")
for token in tokens:
token = str(token).strip()
if "-" in token:
try:
lo, hi = (int(p) for p in token.split("-", 1))
except ValueError:
continue
indices.extend(range(lo - 1, hi))
else:
try:
indices.append(int(token) - 1)
except ValueError:
continue
return [i for i in indices if 0 <= i < total]
def _to_chunks(text: str, max_chars: Optional[int]) -> List[str]:
"""Chunk parsed text via the ingestion chunker; bounded and JSON-safe for the result."""
from application.parser.chunking_creator import ChunkerCreator
from application.parser.schema.base import Document
chunker = ChunkerCreator.create_chunker("classic_chunk")
chunks = chunker.chunk([Document(text=text)])
cap = int(max_chars or 0)
out: List[str] = []
for chunk in chunks:
body = getattr(chunk, "text", str(chunk))
out.append(body[:cap] if cap > 0 else body)
if len(out) >= 200:
break
return out
def parse_document_bytes(
data: bytes,
filename: str,
*,
output: str = "markdown",
ocr: str = "auto",
pages: Any = None,
engine: str = "auto",
max_chars: Optional[int] = None,
include_tables: bool = True,
) -> Dict[str, Any]:
"""Parse untrusted document bytes into a bounded shaped result; whitelist + size + cleanup guarded."""
if output not in _VALID_OUTPUTS:
return {"error": f"unsupported output '{output}'; expected one of {_VALID_OUTPUTS}."}
if ocr not in _VALID_OCR:
return {"error": f"unsupported ocr '{ocr}'; expected one of {_VALID_OCR}."}
if engine not in _VALID_ENGINES:
return {"error": f"unsupported engine '{engine}'; expected one of {_VALID_ENGINES}."}
safe_name = safe_filename(filename) or "document"
suffix = os.path.splitext(safe_name)[1].lower()
if suffix not in SUPPORTED_SOURCE_EXTENSIONS:
return {"error": f"unsupported file type '{suffix or filename}'."}
cap = _max_input_bytes()
if len(data) > cap:
return {"error": f"input document is too large: {len(data)} bytes exceeds the {cap}-byte cap."}
ocr_enabled = _resolve_ocr_enabled(ocr)
tmp_dir = tempfile.mkdtemp(prefix="docparse-")
tmp_path = Path(tmp_dir) / safe_name
try:
tmp_path.write_bytes(data)
return _shape(tmp_path, suffix, output, ocr_enabled, engine, pages, max_chars, include_tables)
except Exception as exc:
logger.exception("parse_document_bytes: parsing failed")
return {"error": f"parsing failed: {type(exc).__name__}: {exc}"}
finally:
try:
tmp_path.unlink(missing_ok=True)
os.rmdir(tmp_dir)
except OSError:
logger.warning("parse_document_bytes: temp cleanup failed for %s", tmp_dir, exc_info=True)
def _shape(
path: Path,
suffix: str,
output: str,
ocr_enabled: bool,
engine: str,
pages: Any,
max_chars: Optional[int],
include_tables: bool,
) -> Dict[str, Any]:
"""Run the selected parser/engine and shape the result per ``output``; bounded throughout.
``output='structured'`` always uses Docling regardless of ``engine`` — the ``fast``
engine is markdown/text only and cannot produce the structured dict.
"""
if output == "structured":
try:
extracted = _docling_structured(path, ocr_enabled=ocr_enabled, include_tables=include_tables)
except Exception as exc:
return {"error": f"structured parsing requires Docling: {type(exc).__name__}: {exc}"}
bounded, truncated = _bounded(extracted["markdown"], max_chars)
return {
"output": "structured",
"content": bounded,
"truncated": truncated,
"tables": [_compact_table(t) for t in extracted["tables"]],
"structured": extracted["structured"],
"summary": _structure_summary(extracted["structured"]),
"page_count": extracted["page_count"],
}
parser = _pick_parser(suffix, ocr_enabled=ocr_enabled, engine=engine)
if parser is None:
# A whitelisted extension with no dedicated parser (e.g. .txt) reads as plain
# text, matching SimpleDirectoryReader's standard-read fallback.
text = path.read_text(errors="ignore")
else:
text = _parse_to_text(parser, path)
text = _apply_pages(text, pages)
if output == "chunks":
return {"output": "chunks", "chunks": _to_chunks(text, max_chars), "truncated": False}
tables: List[Dict[str, Any]] = []
if include_tables and engine != "fast":
try:
tables = [_compact_table(t) for t in _docling_structured(
path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]]
except Exception:
tables = []
bounded, truncated = _bounded(text, max_chars)
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
if tables:
payload["tables"] = tables
return payload
def _bounded(text: str, max_chars: Optional[int]) -> tuple[str, bool]:
"""Bound text to ``max_chars`` (chars) or the default byte window; flag truncation."""
if max_chars and int(max_chars) > 0:
capped = text[: int(max_chars)]
return capped, len(capped) < len(text)
bounded = truncate_text_head_tail(text)
return bounded, bounded != text
application/worker.py
+107
View File
@@ -34,6 +34,7 @@ from application.parser.schema.base import Document
from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.repositories.attachments import AttachmentsRepository
from application.storage.db.repositories.ingest_chunk_progress import (
IngestChunkProgressRepository,
@@ -1609,6 +1610,112 @@ def attachment_worker(self, file_info, user):
raise
def parse_document_worker(self, artifact_id, parent, user_id, options):
"""Parse an input artifact's bytes to a shaped result on the dedicated parsing queue.
Security: the artifact is re-resolved through the run-scoped gate IN THE WORKER
(never trusting a raw storage path) so authz is enforced independently here, in
addition to the pre-enqueue check in the tool.
"""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.parser.document_reader import bound_parse_payload, parse_document_bytes
options = options or {}
parent = parent or {}
conversation_id = parent.get("conversation_id")
workflow_run_id = parent.get("workflow_run_id")
if conversation_id is None and workflow_run_id is None:
return {"status": "error", "error": "parse_document requires a conversation_id or workflow_run_id."}
# Re-resolve through the parent-scoped gate so a forged/cross-run id is rejected
# in the worker too; resolve a short ref to an id within this parent only.
try:
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
resolved_id = resolve_artifact_id(
repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
artifact = (
repo.get_artifact_in_parent(
resolved_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
)
if resolved_id is not None
else None
)
if artifact is None:
return {"status": "error", "error": f"input artifact {artifact_id} not found in this conversation/run."}
version = repo.get_version(resolved_id, artifact["current_version"])
except Exception:
logging.error("parse_document_worker: failed to resolve input artifact", exc_info=True)
return {"status": "error", "error": f"failed to load input artifact {artifact_id}."}
if not version or not version.get("storage_path"):
return {"status": "error", "error": f"input artifact {artifact_id} has no stored content."}
display_name = version.get("filename") or artifact.get("title") or str(resolved_id)
filename = safe_filename(display_name)
try:
data = StorageCreator.get_storage().get_file(version["storage_path"]).read()
except Exception:
logging.error("parse_document_worker: failed to read input artifact bytes", exc_info=True)
return {"status": "error", "error": f"failed to read input artifact {artifact_id}."}
result = parse_document_bytes(
data,
filename,
output=options.get("output", "markdown"),
ocr=options.get("ocr", "auto"),
pages=options.get("pages"),
engine=options.get("engine", "auto"),
max_chars=options.get("max_chars"),
include_tables=bool(options.get("include_tables", True)),
)
if result.get("error"):
return {"status": "error", "error": result["error"]}
payload = {"status": "ok", **result}
if options.get("persist"):
# The full shaped result is persisted by reference; only a bounded view rides
# back through the Redis result backend (the bytes live in the artifact).
artifact_ref = _persist_parse_result(result, display_name, user_id, parent, options)
if isinstance(artifact_ref, dict) and artifact_ref.get("error"):
payload["artifact_error"] = artifact_ref["error"]
elif artifact_ref is not None:
payload["artifact"] = artifact_ref
# Bound the Redis-backed view across all shapes: content is re-windowed, chunks are
# capped, and structured (needed for json_schema validation) is bounded by the input
# cap + table caps. The FULL result already lives in the persisted artifact above.
payload = bound_parse_payload(payload)
return payload
def _persist_parse_result(result, title, user_id, parent, options):
"""Persist the full shaped parse result as an owner/parent-scoped ``data`` artifact; return its ref."""
from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact
try:
data = json.dumps(result).encode("utf-8")
except (TypeError, ValueError):
logging.error("parse_document_worker: parse result is not JSON-serializable", exc_info=True)
return {"error": "parse result is not JSON-serializable."}
base = safe_filename(title) or "document"
filename = f"{base}.parsed.json"
try:
return persist_new_artifact(
user_id=user_id,
kind="data",
data=data,
filename=filename,
mime_type="application/json",
title=f"{title} (parsed)",
conversation_id=parent.get("conversation_id"),
workflow_run_id=parent.get("workflow_run_id"),
produced_by={"tool": "read_document", "action": "read_document", "tool_id": options.get("tool_id")},
)
except QuotaExceeded as exc:
return {"error": str(exc)}
def agent_webhook_worker(self, agent_id, payload):
"""Process the webhook payload for an agent.
deployment/docker-compose-azure.yaml
+2 -1
View File
@@ -32,7 +32,8 @@ services:
worker:
build: ../application
command: celery -A application.app.celery worker -l INFO
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
command: celery -A application.app.celery worker -l INFO -Q docsgpt,parsing
env_file:
- ../.env
environment:
deployment/docker-compose-hub.yaml
+2 -1
View File
@@ -39,7 +39,8 @@ services:
worker:
user: root
image: arc53/docsgpt:develop
command: celery -A application.app.celery worker -l INFO -B
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing
env_file:
- ../.env
environment:
deployment/docker-compose.yaml
+5 -1
View File
@@ -45,7 +45,11 @@ services:
worker:
user: root
build: ../application
command: celery -A application.app.celery worker -l INFO -B
# Consumes the default queue AND the dedicated `parsing` queue (read_document /
# parse_document). Without `parsing` here the read_document await never resolves.
# For heavy/OCR parsing run a separate worker with `-Q parsing` (see
# deployment/sandbox/README.md).
command: celery -A application.app.celery worker -l INFO -B -Q docsgpt,parsing
env_file:
- ../.env
environment:
deployment/k8s/deployments/docsgpt-deploy.yaml
+3 -1
View File
@@ -93,7 +93,9 @@ spec:
containers:
- name: docsgpt-worker
image: arc53/docsgpt
command: ["celery", "-A", "application.app.celery", "worker", "-l", "INFO", "-n", "worker.%h"]
# `parsing` queue carries read_document/parse_document; required for its await to resolve.
# For heavy/OCR parsing, run a separate deployment with `-Q parsing` (and GPU env).
command: ["celery", "-A", "application.app.celery", "worker", "-l", "INFO", "-n", "worker.%h", "-Q", "docsgpt,parsing"]
resources:
limits:
memory: "4Gi"
deployment/sandbox/Dockerfile
+7 -5
View File
@@ -32,11 +32,13 @@ RUN pip install --no-cache-dir \
pandas==2.2.3 \
matplotlib==3.9.2
# Docling (MIT) for the document_extractor tool — OFF by default because it pulls
# torch + models and makes the image multi-GB. Build the "extract" variant with
# `--build-arg INSTALL_DOCLING=true` (see README) when document extraction is
# needed. Docling is MIT and uses its own PDF backend; PyMuPDF (AGPL) is NOT
# installed here. The base image and the app's requirements stay docling-free.
# Docling (MIT) — OFF by default because it pulls torch + models and makes the
# image multi-GB. NOTE: document parsing now runs on the Celery `parsing` worker
# (see read_document / parse_document_worker), NOT in this sandbox, so this
# INSTALL_DOCLING build-arg path is effectively unused for the read_document flow;
# it remains only for sandbox code that opts into Docling explicitly. The build
# mechanics are kept for that case. Docling is MIT and uses its own PDF backend;
# PyMuPDF (AGPL) is NOT installed here. The base image stays docling-free.
ARG INSTALL_DOCLING=false
RUN if [ "$INSTALL_DOCLING" = "true" ]; then \
pip install --no-cache-dir docling==2.8.3; \
deployment/sandbox/README.md
+24 -13
View File
@@ -109,24 +109,35 @@ app chooses it). The same applies to k8s: `SANDBOX_KERNEL_NAME=docsgpt-python`
is set on the `docsgpt-api` and `docsgpt-worker` deployments in
`deployment/k8s/deployments/docsgpt-deploy.yaml`.
## Document extraction variant (Docling)
## Document reading (parsing worker — not the sandbox)
The `document_extractor` tool runs Docling (MIT) inside the runner to convert
documents (pdf/docx/pptx/...) to schema-validated JSON. Docling pulls `torch` and
ML models, which makes the image multi-GB, so it is **off by default**: it is not
in the base image, not in `application/requirements.txt`, and not required in the
dev `.venv`. Build the extract variant only where extraction is needed:
Document reading no longer runs in this sandbox. The `read_document` tool and the
workflow native-file extract branch enqueue a `parse_document` Celery task that
parses the document **in the backend** (Docling, already in
`application/requirements.txt`) and awaits the result. The task is routed to a
dedicated **`parsing` queue** (`settings.DOCUMENT_PARSE_QUEUE`, default
`"parsing"`) so a parse enqueued from inside a Celery worker (headless/scheduled
agent) is served by a separate worker and never self-deadlocks the awaiting one.
Run a dedicated parsing worker that consumes the `parsing` queue:
```bash
docker build -t docsgpt-sandbox-extract \
--build-arg INSTALL_DOCLING=true deployment/sandbox
docker run --rm -p 8888:8888 docsgpt-sandbox-extract
celery -A application.app.celery worker -Q parsing -l INFO
```
Point the app at this runner via `SANDBOX_GATEWAY_URL` exactly as the base image.
Docling uses its own permissive PDF backend; **PyMuPDF (AGPL) is intentionally not
installed.** If the base (non-extract) image is used, `extract_document` returns a
clean "docling is not available in the sandbox runner" error rather than crashing.
It can be GPU-enabled with its own env (`DOCLING_OCR_ENABLED=true` plus GPU
libraries) so OCR-heavy parsing runs on a separate, optionally larger pool.
**Dev / single-worker setups:** without a dedicated parsing worker the default
worker must also consume `parsing`, or the tool's await never resolves:
```bash
celery -A application.app.celery worker -Q docsgpt,parsing -l INFO
```
Tuning settings: `DOCUMENT_PARSE_TIMEOUT` (seconds the tool awaits before
degrading to an error), `DOCUMENT_PARSE_MAX_BYTES` (per-document byte cap; 0
reuses `SANDBOX_MAX_INPUT_BYTES`).
## Network egress / SSRF
scripts/e2e/up.sh
+1 -1
View File
@@ -276,7 +276,7 @@ log "starting Celery worker (solo pool)"
PYTHONPATH="$REPO_ROOT${PYTHONPATH:+:$PYTHONPATH}" \
PYTHONUNBUFFERED=1 \
nohup "$PY_BIN" -m celery -A application.app.celery worker \
-l INFO --pool=solo \
-l INFO --pool=solo -Q docsgpt,parsing \
--without-gossip --without-mingle --without-heartbeat \
>"$CELERY_LOG" 2>&1 &
echo $! > "$CELERY_PID"
tests/agents/test_workflow_node_attachments.py
+7 -18
View File
@@ -225,18 +225,13 @@ def test_native_on_unsupported_mime_raises(monkeypatch):
@pytest.mark.unit
def test_extract_non_text_uses_docling(monkeypatch):
"""A non-text mime under extract routes through the Docling helper."""
def test_extract_non_text_uses_parse_worker(monkeypatch):
"""A non-text mime under extract routes through the parsing-worker path (no sandbox)."""
aid, rec = _artifact(RUN_ID, "application/vnd.openxmlformats", filename="r.docx")
_patch_repo(monkeypatch, {aid: rec})
eng = _engine(monkeypatch)
monkeypatch.setattr(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"\x00binary")),
)
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "EXTRACTED MD",
WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED MD"
)
cfg = _node_config(input_documents=[aid], file_passing="extract")
@@ -298,10 +293,9 @@ def test_native_file_cap_bounds_native_then_extracts(monkeypatch):
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"img-bytes")),
)
# extract of a non-text image falls back to Docling; stub it so it returns text.
# extract of a non-text image routes through the parsing worker; stub it so it returns text.
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "fallback text",
WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "fallback text"
)
cfg = _node_config(input_documents=ids, file_passing="auto")
@@ -354,7 +348,7 @@ def test_oversize_text_skipped_by_post_read_guard_when_size_missing(monkeypatch)
@pytest.mark.unit
def test_large_under_cap_text_is_windowed_not_inlined_whole(monkeypatch):
"""A large-but-under-cap text file is bounded to a head+tail window, not inlined whole."""
from application.agents.tools.document_extractor import _MARKDOWN_MAX_BYTES
from application.parser.document_reader import _TEXT_MAX_BYTES as _MARKDOWN_MAX_BYTES
big_text = ("A" * (_MARKDOWN_MAX_BYTES * 3)).encode("utf-8")
aid, rec = _artifact(RUN_ID, "text/plain", filename="notes.txt", size=len(big_text))
@@ -490,12 +484,7 @@ def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch):
_patch_repo(monkeypatch, {aid: rec})
_patch_capabilities(monkeypatch)
monkeypatch.setattr(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"image-bytes")),
)
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "EXTRACTED",
WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED"
)
eng = _engine(monkeypatch)
tests/agents/tools/test_document_extractor_live.py
-56
View File
@@ -1,56 +0,0 @@
"""Optional live Docling extraction test, SKIPPED unless RUN_DOCLING_LIVE=1.
Default CI/dev runs never import docling (which pulls torch + models), so this is
skipped unless explicitly opted in. It runs the FIXED extraction program in-process
against a tiny generated PDF and asserts the compact result shape, exercising the
exact program string the tool ships to the sandbox.
"""
from __future__ import annotations
import json
import os
import tempfile
import pytest
from application.agents.tools.document_extractor import _EXTRACT_PROGRAM
RUN_LIVE = os.environ.get("RUN_DOCLING_LIVE") == "1"
pytestmark = pytest.mark.skipif(not RUN_LIVE, reason="set RUN_DOCLING_LIVE=1 to run live Docling extraction")
def _make_pdf(path: str) -> None:
"""Write a one-line PDF using reportlab (kept out of the import path unless live)."""
from reportlab.pdfgen import canvas
c = canvas.Canvas(path)
c.drawString(72, 720, "Compliance balance: 1000")
c.showPage()
c.save()
def test_live_docling_extracts_pdf():
pytest.importorskip("docling")
pytest.importorskip("reportlab")
workdir = tempfile.mkdtemp()
input_path = os.path.join(workdir, "inputs", "doc.pdf")
os.makedirs(os.path.dirname(input_path), exist_ok=True)
_make_pdf(input_path)
params_path = os.path.join(workdir, "params.json")
result_path = os.path.join(workdir, "result.json")
with open(params_path, "w") as fh:
json.dump({"input_path": input_path, "markdown_max_bytes": 8000, "max_tables": 20}, fh)
program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
namespace: dict = {}
exec(compile(program, "<extractor>", "exec"), namespace, namespace) # noqa: S102
with open(result_path) as fh:
result = json.load(fh)
assert result.get("ok") is True, result
assert "balance" in result["markdown"].lower()
assert isinstance(result["structured"], dict)
tests/agents/tools/test_document_extractor_unit.py
-624
View File
@@ -1,624 +0,0 @@
"""Unit tests for DocumentExtractorTool: staging, fixed-program shaping, schema validation, and injection safety.
Docling and the sandbox exec are MOCKED so no heavy import (torch/models) or real
kernel is touched. These cover: a parent-scoped input is staged and extracted into
a compact payload; json_schema validation (pass + fail); a cross-tenant input id is
denied; and the extraction program is FIXED with params travelling as data so a
malicious filename/parameter is never executed.
"""
from __future__ import annotations
import json
from typing import Any, Dict, List, Optional
import os
import sys
import tempfile
import types
import uuid
import application.agents.tools.document_extractor as de
from application.agents.tools.document_extractor import (
_EXTRACT_PROGRAM,
_MARKDOWN_MAX_BYTES,
_MAX_CELL_CHARS,
_MAX_TABLE_ROWS,
DocumentExtractorTool,
)
from application.sandbox.base import ExecResult
# Production artifact ids are uuids; the resolver gates non-uuid/non-ref inputs to
# "not found", so the stubbed happy-path input uses a real uuid.
_ART_ID = str(uuid.uuid4())
# A canned Docling-style extraction result the mocked program "writes" to result.json.
_CANNED_EXTRACT = {
"ok": True,
"markdown": "# Statement\n\nBalance: 1000",
"markdown_truncated": False,
"tables": [{"columns": ["item", "amount"], "rows": [["fee", "10"]]}],
"page_count": 2,
"structured": {
"schema_name": "DoclingDocument",
"texts": [{"text": "Balance: 1000"}],
"tables": [{"data": "..."}],
"pages": {"1": {}, "2": {}},
},
}
class _FakeManager:
"""In-memory sandbox stand-in: records put_file calls and the exact exec program/code."""
def __init__(self, result: ExecResult, result_file: bytes) -> None:
self._result = result
self._result_file = result_file
self.put_files: Dict[str, bytes] = {}
self.exec_programs: List[str] = []
self.opened: List[Any] = []
self.closed: List[str] = []
self.removed: List[str] = []
def open(self, session_id, ttl=None):
self.opened.append((session_id, ttl))
return session_id
def remove_path(self, session_id, path):
self.removed.append(path)
def put_file(self, session_id, dest_path, data):
self.put_files[dest_path] = data
def exec(self, session_id, code, timeout=None):
self.exec_programs.append(code)
return self._result
def get_file(self, session_id, path):
return self._result_file
def close(self, session_id):
self.closed.append(session_id)
class _FakeVersion(dict):
pass
def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str]):
"""Patch db_readonly + ArtifactsRepository so input scoping is exercised without a DB."""
class _Repo:
def __init__(self, conn):
pass
def artifact_id_at_position(self, n, *, conversation_id=None, workflow_run_id=None):
# The lone artifact lives at position 1, scoped to the right parent only.
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):
# Mirror the real scope gate: the artifact resolves only for the right parent.
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"}
def get_version(self, artifact_id, version):
return _FakeVersion(
{"filename": "statement.pdf", "storage_path": f"inputs/u/artifacts/{artifact_id}/v1/statement.pdf"}
)
class _Conn:
def __enter__(self):
return object()
def __exit__(self, *exc):
return False
monkeypatch.setattr(de, "db_readonly", lambda: _Conn())
monkeypatch.setattr(de, "ArtifactsRepository", _Repo)
class _FakeStorage:
"""Returns canned input bytes for any storage path."""
def __init__(self, data: bytes = b"%PDF-1.4 fake") -> None:
self._data = data
def get_file(self, path):
import io
return io.BytesIO(self._data)
def _patch_storage(monkeypatch):
monkeypatch.setattr(de.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage()))
def _patch_manager(monkeypatch, manager):
monkeypatch.setattr(de.SandboxCreator, "get_manager", lambda: manager)
def _patch_no_persist(monkeypatch):
# Default the persist path to a no-op so most tests don't touch the artifact store.
monkeypatch.setattr(de, "persist_new_artifact", lambda **kwargs: None)
def _tool(**config) -> DocumentExtractorTool:
base = {"conversation_id": "conv-1", "tool_id": "t-1"}
base.update(config)
return DocumentExtractorTool(tool_config=base, user_id="u-1")
def _manager_with_extract(extract: Dict[str, Any]) -> _FakeManager:
return _FakeManager(ExecResult(status="ok"), json.dumps(extract).encode("utf-8"))
# ---------------------------------------------------------------------------
# Guards
# ---------------------------------------------------------------------------
def test_unknown_action_rejected():
out = _tool().execute_action("nope", input="a")
assert out["status"] == "error" and "unknown action" in out["error"]
def test_requires_user_and_parent():
no_user = DocumentExtractorTool({"conversation_id": "c"}, user_id=None)
out = no_user.execute_action("extract_document", input="a")
assert out["status"] == "error" and "user_id" in out["error"]
no_parent = DocumentExtractorTool({}, user_id="u")
out2 = no_parent.execute_action("extract_document", input="a")
assert out2["status"] == "error" and "conversation_id" in out2["error"]
def test_input_id_required():
out = _tool().execute_action("extract_document", input=" ")
assert out["status"] == "error" and "input artifact id is required" in out["error"]
def test_action_metadata_shape():
meta = _tool().get_actions_metadata()[0]
assert meta["name"] == "extract_document"
assert meta["parameters"]["required"] == ["input"]
assert "json_schema" in meta["parameters"]["properties"]
# ---------------------------------------------------------------------------
# Happy path: staging + compact payload shaping
# ---------------------------------------------------------------------------
def test_extract_stages_input_and_shapes_compact_payload(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "ok"
structured = out["structured"]
# Compact payload: markdown + tables + page_count + a structure summary, no raw bytes.
assert structured["markdown"].startswith("# Statement")
assert structured["tables"] == _CANNED_EXTRACT["tables"]
assert structured["page_count"] == 2
assert structured["summary"] == {"texts": 1, "tables": 1, "pages": 2}
assert "bytes" not in out
# The input document was staged into the workspace as a DATA file (not code).
staged = [p for p in manager.put_files if p.endswith("statement.pdf")]
assert staged, "input document should be staged into the workspace"
assert manager.put_files[staged[0]] == b"%PDF-1.4 fake"
# The session is closed after a one-shot extraction.
assert manager.closed == ["conv-1"]
def test_extract_removes_its_scratch_token_dir(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
_tool().execute_action("extract_document", input=_ART_ID, persist=False)
# The per-render token dir (extract/<token>) is removed before/at teardown so
# a reused session doesn't accumulate staged inputs + results on disk.
assert manager.removed, "extractor should remove its scratch token dir"
removed_dir = manager.removed[0]
assert removed_dir.startswith("extract/")
# The dir prefix must cover the files that were staged under it.
assert all(p.startswith(removed_dir + "/") for p in manager.put_files)
def test_extract_persists_data_artifact_by_reference(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
captured: Dict[str, Any] = {}
def _fake_persist(**kwargs):
captured.update(kwargs)
return {"artifact_id": "new-art", "version": 1, "filename": "x.json",
"mime_type": "application/json", "size": 10}
monkeypatch.setattr(de, "persist_new_artifact", _fake_persist)
tool = _tool()
out = tool.execute_action("extract_document", input=_ART_ID) # persist defaults to true
assert out["status"] == "ok"
assert out["artifact"]["artifact_id"] == "new-art"
assert tool.get_artifact_id("extract_document") == "new-art"
# Persisted as a JSON data artifact carrying the full extraction (by reference).
assert captured["kind"] == "data"
assert captured["mime_type"] == "application/json"
assert json.loads(captured["data"].decode("utf-8")) == _CANNED_EXTRACT
def test_quota_exceeded_on_persist_surfaces_cleanly(monkeypatch):
from application.sandbox.artifacts_capture import QuotaExceeded
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
def _quota_blocked(**kwargs):
raise QuotaExceeded("artifact storage quota reached")
monkeypatch.setattr(de, "persist_new_artifact", _quota_blocked)
out = _tool().execute_action("extract_document", input=_ART_ID)
# Extraction still succeeds; the quota failure surfaces as a non-fatal note.
assert out["status"] == "ok"
assert "artifact" not in out
assert "quota" in out["artifact_error"].lower()
# ---------------------------------------------------------------------------
# Cross-tenant scope
# ---------------------------------------------------------------------------
def test_cross_tenant_input_is_denied(monkeypatch):
# Tool is bound to conv-1; the artifact only exists under conv-OTHER, so the
# parent-scoped fetch returns None and the extraction is refused.
_stub_repo(monkeypatch, found=True, conv="conv-OTHER", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "error"
assert "not found in this conversation/run" in out["error"]
# Nothing was staged or executed for a denied input.
assert manager.exec_programs == []
def test_extract_accepts_short_ref_input(monkeypatch):
# A short ref (A1) resolves to the parent-scoped artifact at position 1 and extracts.
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input="A1", persist=False)
assert out["status"] == "ok"
assert out["structured"]["markdown"].startswith("# Statement")
# The resolved input was staged and the program executed.
assert manager.exec_programs
def test_extract_short_ref_out_of_range_is_clean_error(monkeypatch):
# A2 has no artifact behind it (stub only fills position 1) -> clean not-found.
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input="A2", persist=False)
assert out["status"] == "error"
assert "not found in this conversation/run" in out["error"]
assert manager.exec_programs == []
def test_missing_input_artifact_is_denied(monkeypatch):
_stub_repo(monkeypatch, found=False, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input="ghost", persist=False)
assert out["status"] == "error" and "not found" in out["error"]
assert manager.exec_programs == []
# ---------------------------------------------------------------------------
# json_schema validation (pass + fail)
# ---------------------------------------------------------------------------
def _schema_requiring_texts() -> Dict[str, Any]:
return {
"type": "object",
"properties": {"texts": {"type": "array"}, "schema_name": {"type": "string"}},
"required": ["texts", "schema_name"],
}
def test_json_schema_validation_passes(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
_patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT))
out = _tool().execute_action(
"extract_document", input=_ART_ID, json_schema=_schema_requiring_texts(), persist=False
)
assert out["status"] == "ok"
def test_json_schema_validation_fails_with_clean_error(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
# structured payload lacks the required "amount" field the schema demands.
schema = {"type": "object", "required": ["amount"], "properties": {"amount": {"type": "number"}}}
_patch_manager(monkeypatch, _manager_with_extract(_CANNED_EXTRACT))
out = _tool().execute_action("extract_document", input=_ART_ID, json_schema=schema, persist=False)
assert out["status"] == "error"
assert "did not match json_schema" in out["error"]
def test_malformed_json_schema_rejected_before_run(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
# No "type"/"schema" key -> normalize_json_schema_payload rejects it up front.
out = _tool().execute_action("extract_document", input=_ART_ID, json_schema={"properties": {}}, persist=False)
assert out["status"] == "error" and "invalid json_schema" in out["error"]
assert manager.exec_programs == []
# ---------------------------------------------------------------------------
# Docling-unavailable / extractor error surfacing
# ---------------------------------------------------------------------------
def test_docling_unavailable_surfaces_clean_error(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
# The program's own "docling not available" branch writes an {"error": ...} result.
err_result = {
"error": "docling is not available in the sandbox runner: ModuleNotFoundError: No module named 'docling'"
}
_patch_manager(monkeypatch, _manager_with_extract(err_result))
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "error"
assert "docling is not available" in out["error"]
def test_exec_error_surfaces_clean_error(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _FakeManager(
ExecResult(status="error", error_name="TimeoutError", error_value="exceeded 60s"), b""
)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "error"
assert "TimeoutError: exceeded 60s" in out["error"]
# ---------------------------------------------------------------------------
# Injection safety: the extraction program is FIXED; params travel as data
# ---------------------------------------------------------------------------
def test_extraction_program_is_fixed_and_params_are_data(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
_tool().execute_action("extract_document", input=_ART_ID, persist=False)
program = manager.exec_programs[0]
# The executed program is the fixed template with only server-controlled path
# literals substituted — no input id, filename, or document content appears.
assert "json.load(open(" in program
assert "DocumentConverter" in program
assert _ART_ID not in program
assert "statement.pdf" not in program
# Params reached the program as a JSON DATA file, not via interpolation.
params_files = [p for p in manager.put_files if p.endswith("params.json")]
assert params_files
params = json.loads(manager.put_files[params_files[0]].decode("utf-8"))
assert params["input_path"].endswith("statement.pdf")
def test_malicious_param_value_is_not_executed(monkeypatch):
"""A hostile filename ends up only in a DATA file the program json.loads, never in the program text."""
# The artifact's stored filename is attacker-controlled in the worst case;
# prove a code-shaped filename never lands in the executed program.
payload = "__import__('os').system('echo PWNED'); x = 'a.pdf"
class _Repo:
def __init__(self, conn):
pass
def get_artifact_in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
return {"id": artifact_id, "current_version": 1, "title": payload}
def get_version(self, artifact_id, version):
return {"filename": payload, "storage_path": "inputs/u/artifacts/x/v1/f"}
class _Conn:
def __enter__(self):
return object()
def __exit__(self, *exc):
return False
monkeypatch.setattr(de, "db_readonly", lambda: _Conn())
monkeypatch.setattr(de, "ArtifactsRepository", _Repo)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "ok"
program = manager.exec_programs[0]
# The injection payload is absent from the program; it (sanitized) only rides
# in the params DATA file, so it is parsed as a string, never executed.
assert "PWNED" not in program
assert "__import__" not in program
params_files = [p for p in manager.put_files if p.endswith("params.json")]
params = json.loads(manager.put_files[params_files[0]].decode("utf-8"))
# The path is sanitized but even unsanitized it is pure data inside params.json.
assert "input_path" in params
# ---------------------------------------------------------------------------
# The fixed program template itself
# ---------------------------------------------------------------------------
def test_program_template_has_no_unbound_interpolation_points():
# Only the two server-controlled path placeholders may be formatted in; the
# template must format cleanly with just those, proving nothing else is a hole.
rendered = _EXTRACT_PROGRAM.format(params_path="extract/x/params.json", result_path="extract/x/result.json")
assert "params.json" in rendered
assert "result.json" in rendered
# Compiles as valid Python (no broken escapes from the doubled braces).
compile(rendered, "<extractor>", "exec")
# ---------------------------------------------------------------------------
# Markdown head+tail window (finding: the END of compliance docs must survive)
# ---------------------------------------------------------------------------
def _run_program_with_fake_docling(monkeypatch, markdown: str, md_cap: int) -> Dict[str, Any]:
"""Exec the REAL extraction program with docling stubbed so we test the markdown windowing."""
class _FakeDoc:
tables: List[Any] = []
pages: Dict[str, Any] = {}
def export_to_markdown(self):
return markdown
def export_to_dict(self):
return {"texts": []}
class _FakeConverter:
def convert(self, src):
return types.SimpleNamespace(document=_FakeDoc())
fake_mod = types.ModuleType("docling")
fake_sub = types.ModuleType("docling.document_converter")
fake_sub.DocumentConverter = _FakeConverter
fake_mod.document_converter = fake_sub
monkeypatch.setitem(sys.modules, "docling", fake_mod)
monkeypatch.setitem(sys.modules, "docling.document_converter", fake_sub)
workdir = tempfile.mkdtemp()
params_path = os.path.join(workdir, "params.json")
result_path = os.path.join(workdir, "result.json")
with open(params_path, "w") as fh:
json.dump({"input_path": "x", "markdown_max_bytes": md_cap, "max_tables": 20}, fh)
program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
ns: Dict[str, Any] = {}
exec(compile(program, "<extractor>", "exec"), ns, ns) # noqa: S102
with open(result_path) as fh:
return json.load(fh)
def test_markdown_window_keeps_head_and_tail_when_truncated(monkeypatch):
head_marker = "HEAD_START_OF_DOC"
tail_marker = "TAIL_END_OF_DOC_TOTALS"
body = "x" * (_MARKDOWN_MAX_BYTES * 2)
markdown = head_marker + body + tail_marker
result = _run_program_with_fake_docling(monkeypatch, markdown, _MARKDOWN_MAX_BYTES)
assert result["markdown_truncated"] is True
out = result["markdown"]
# Both the beginning AND the end of the document survive the byte budget.
assert head_marker in out
assert tail_marker in out
assert "...[truncated" in out
# The window respects the budget (plus the small truncation marker).
assert len(out) <= _MARKDOWN_MAX_BYTES + 64
def test_markdown_full_content_when_under_cap(monkeypatch):
markdown = "# Short doc\n\nBalance: 1000\n\nSignature: Jane"
result = _run_program_with_fake_docling(monkeypatch, markdown, _MARKDOWN_MAX_BYTES)
assert result["markdown_truncated"] is False
assert result["markdown"] == markdown
# ---------------------------------------------------------------------------
# Per-table content cap (a single giant table must not bloat context)
# ---------------------------------------------------------------------------
def test_compact_payload_bounds_table_rows_and_cell_bytes(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_storage(monkeypatch)
_patch_no_persist(monkeypatch)
big_cell = "y" * (_MAX_CELL_CHARS * 3)
huge_table = {
"columns": ["a", "b"],
"rows": [[str(i), big_cell] for i in range(_MAX_TABLE_ROWS * 4)],
}
extract = dict(_CANNED_EXTRACT)
extract["tables"] = [huge_table]
_patch_manager(monkeypatch, _manager_with_extract(extract))
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "ok"
table = out["structured"]["tables"][0]
# Rows are capped and flagged with the original total.
assert len(table["rows"]) == _MAX_TABLE_ROWS
assert table["rows_truncated"] is True
assert table["total_rows"] == _MAX_TABLE_ROWS * 4
# Long cell strings are truncated.
assert len(table["rows"][0][1]) <= _MAX_CELL_CHARS + len("...[truncated]")
assert table["rows"][0][1].endswith("...[truncated]")
# ---------------------------------------------------------------------------
# Input size cap (don't load an unbounded artifact into backend + kernel memory)
# ---------------------------------------------------------------------------
def test_oversized_input_is_rejected_before_exec(monkeypatch):
_stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
_patch_no_persist(monkeypatch)
monkeypatch.setattr(de.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage(b"P" * 4096)))
monkeypatch.setattr(de.settings, "SANDBOX_MAX_INPUT_BYTES", 1024, raising=False)
manager = _manager_with_extract(_CANNED_EXTRACT)
_patch_manager(monkeypatch, manager)
out = _tool().execute_action("extract_document", input=_ART_ID, persist=False)
assert out["status"] == "error"
assert "too large" in out["error"]
# Nothing was staged or executed for an oversized input.
assert manager.exec_programs == []
assert manager.opened == []
tests/agents/tools/test_read_document_tool.py
+265
View File
@@ -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"]
tests/api/user/test_tasks.py
+4 -1
View File
@@ -276,7 +276,7 @@ class TestSetupPeriodicTasks:
setup_periodic_tasks(sender)
assert sender.add_periodic_task.call_count == 11
assert sender.add_periodic_task.call_count == 12
calls = sender.add_periodic_task.call_args_list
@@ -308,6 +308,9 @@ class TestSetupPeriodicTasks:
# schedule runs cleanup (24h)
assert calls[10][0][0] == timedelta(hours=24)
assert calls[10][1].get("name") == "cleanup-schedule-runs"
# sandbox session reaper (60s)
assert calls[11][0][0] == timedelta(seconds=60)
assert calls[11][1].get("name") == "reap-sandbox-sessions"
class TestMcpOauthTask:
tests/parser/test_document_reader.py
+262
View File
@@ -0,0 +1,262 @@
"""Unit tests for parse_document_bytes: output shapes, whitelist/size guards, params, and cleanup.
Docling-heavy paths are stubbed or skipped; these cover the shaping and the
untrusted-content safeguards (extension whitelist, byte cap, temp cleanup).
"""
from __future__ import annotations
from pathlib import Path
from typing import Any, Dict, List
import pytest
import application.parser.document_reader as dr
from application.parser.document_reader import (
bound_parse_payload,
parse_document_bytes,
truncate_text_head_tail,
)
# ---------------------------------------------------------------------------
# Guards: whitelist + size cap
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_unknown_extension_is_rejected():
out = parse_document_bytes(b"data", "evil.exe")
assert "error" in out and "unsupported file type" in out["error"]
@pytest.mark.unit
def test_size_cap_rejects_oversize(monkeypatch):
monkeypatch.setattr(dr.settings, "DOCUMENT_PARSE_MAX_BYTES", 8, raising=False)
out = parse_document_bytes(b"P" * 64, "note.txt", output="text")
assert "error" in out and "too large" in out["error"]
@pytest.mark.unit
def test_bad_output_ocr_engine_rejected():
assert "unsupported output" in parse_document_bytes(b"x", "a.txt", output="nope")["error"]
assert "unsupported ocr" in parse_document_bytes(b"x", "a.txt", ocr="maybe")["error"]
assert "unsupported engine" in parse_document_bytes(b"x", "a.txt", engine="ghost")["error"]
# ---------------------------------------------------------------------------
# Plain-text path: .txt has no dedicated parser -> standard read
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_text_output_for_plain_text():
out = parse_document_bytes(b"hello world\n", "note.txt", output="text", include_tables=False)
assert out["output"] == "text"
assert out["content"] == "hello world\n"
assert out["truncated"] is False
@pytest.mark.unit
def test_markdown_output_default():
out = parse_document_bytes(b"# Title\n", "note.txt", include_tables=False)
assert out["output"] == "markdown"
assert "# Title" in out["content"]
@pytest.mark.unit
def test_max_chars_truncates_and_flags():
out = parse_document_bytes(("A" * 100).encode(), "note.txt", output="text", max_chars=10, include_tables=False)
assert out["truncated"] is True
assert len(out["content"]) == 10
@pytest.mark.unit
def test_default_window_truncates_large_text():
big = ("A" * (dr._TEXT_MAX_BYTES * 3)).encode()
out = parse_document_bytes(big, "note.txt", output="text", include_tables=False)
assert out["truncated"] is True
assert "...[truncated" in out["content"]
assert len(out["content"].encode("utf-8")) <= dr._TEXT_MAX_BYTES + 64
# ---------------------------------------------------------------------------
# chunks output
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_chunks_output_returns_list():
out = parse_document_bytes(b"para one.\n\npara two.\n", "note.txt", output="chunks", include_tables=False)
assert out["output"] == "chunks"
assert isinstance(out["chunks"], list)
assert all(isinstance(c, str) for c in out["chunks"])
# ---------------------------------------------------------------------------
# engine selection: a mapped parser is chosen and run with the right text shape
# ---------------------------------------------------------------------------
class _FakeParser:
"""Records that it was used and returns a fixed string or list of strings."""
def __init__(self, result):
self._result = result
self.parser_config_set = True
self.inited = False
def init_parser(self):
self.inited = True
def parse_file(self, file: Path, errors: str = "ignore"):
return self._result
@pytest.mark.unit
def test_engine_picks_parser_and_coerces_list(monkeypatch):
fake = _FakeParser(["chunk A", "chunk B"])
monkeypatch.setattr(dr, "get_default_file_extractor", lambda ocr_enabled=None: {".pdf": fake})
out = parse_document_bytes(b"%PDF-1.4", "doc.pdf", output="text", engine="docling", include_tables=False)
assert out["content"] == "chunk A\n\nchunk B"
@pytest.mark.unit
def test_fast_engine_uses_legacy_parser(monkeypatch):
fake = _FakeParser("legacy text")
monkeypatch.setattr(dr, "_legacy_parser_for", lambda suffix: fake)
out = parse_document_bytes(b"%PDF-1.4", "doc.pdf", output="text", engine="fast", include_tables=False)
assert out["content"] == "legacy text"
# ---------------------------------------------------------------------------
# pages: page-range slice on a form-feed delimited blob
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_pages_slices_form_feed_blob(monkeypatch):
fake = _FakeParser("page1\fpage2\fpage3")
monkeypatch.setattr(dr, "get_default_file_extractor", lambda ocr_enabled=None: {".pdf": fake})
out = parse_document_bytes(b"%PDF", "doc.pdf", output="text", pages="2", engine="docling", include_tables=False)
assert out["content"] == "page2"
# ---------------------------------------------------------------------------
# structured output (Docling stubbed)
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_structured_output_shapes_via_docling(monkeypatch):
def _fake_structured(path, *, ocr_enabled, include_tables):
return {
"markdown": "# Statement",
"structured": {"texts": [{}], "tables": [{}], "pages": {"1": {}}},
"tables": [{"columns": ["a"], "rows": [["1"]]}],
"page_count": 1,
}
monkeypatch.setattr(dr, "_docling_structured", _fake_structured)
out = parse_document_bytes(b"%PDF", "doc.pdf", output="structured")
assert out["output"] == "structured"
assert out["content"].startswith("# Statement")
assert out["structured"]["texts"]
assert out["summary"] == {"texts": 1, "tables": 1, "pages": 1}
assert out["page_count"] == 1
assert out["tables"] == [{"columns": ["a"], "rows": [["1"]]}]
@pytest.mark.unit
def test_structured_output_missing_docling_is_clean_error(monkeypatch):
def _boom(path, *, ocr_enabled, include_tables):
raise ImportError("No module named 'docling'")
monkeypatch.setattr(dr, "_docling_structured", _boom)
out = parse_document_bytes(b"%PDF", "doc.pdf", output="structured")
assert "error" in out and "structured parsing requires Docling" in out["error"]
# ---------------------------------------------------------------------------
# table bounding
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_table_rows_and_cells_are_bounded():
big_cell = "y" * (dr._MAX_CELL_CHARS * 3)
table: Dict[str, Any] = {"columns": ["a", "b"], "rows": [[str(i), big_cell] for i in range(dr._MAX_TABLE_ROWS * 4)]}
compact = dr._compact_table(table)
assert len(compact["rows"]) == dr._MAX_TABLE_ROWS
assert compact["rows_truncated"] is True
assert compact["total_rows"] == dr._MAX_TABLE_ROWS * 4
assert compact["rows"][0][1].endswith("...[truncated]")
# ---------------------------------------------------------------------------
# temp cleanup: the staged temp file is removed even on parser failure
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_temp_file_cleaned_up_on_success(monkeypatch):
seen: List[Path] = []
real_mkdtemp = dr.tempfile.mkdtemp
def _tracking_mkdtemp(*a, **k):
d = real_mkdtemp(*a, **k)
seen.append(Path(d))
return d
monkeypatch.setattr(dr.tempfile, "mkdtemp", _tracking_mkdtemp)
parse_document_bytes(b"hi", "note.txt", output="text", include_tables=False)
assert seen and not seen[0].exists()
@pytest.mark.unit
def test_temp_file_cleaned_up_on_parser_error(monkeypatch):
seen: List[Path] = []
real_mkdtemp = dr.tempfile.mkdtemp
def _tracking_mkdtemp(*a, **k):
d = real_mkdtemp(*a, **k)
seen.append(Path(d))
return d
monkeypatch.setattr(dr.tempfile, "mkdtemp", _tracking_mkdtemp)
fake = _FakeParser("x")
fake.parse_file = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))
monkeypatch.setattr(dr, "get_default_file_extractor", lambda ocr_enabled=None: {".pdf": fake})
out = parse_document_bytes(b"%PDF", "doc.pdf", output="text", engine="docling", include_tables=False)
assert "error" in out and "parsing failed" in out["error"]
assert seen and not seen[0].exists()
# ---------------------------------------------------------------------------
# ocr resolution
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_ocr_resolution(monkeypatch):
monkeypatch.setattr(dr.settings, "DOCLING_OCR_ENABLED", True, raising=False)
assert dr._resolve_ocr_enabled("off") is False
assert dr._resolve_ocr_enabled("on") is True
assert dr._resolve_ocr_enabled("auto") is True
monkeypatch.setattr(dr.settings, "DOCLING_OCR_ENABLED", False, raising=False)
assert dr._resolve_ocr_enabled("auto") is False
@pytest.mark.unit
def test_truncate_head_tail_keeps_both_ends():
text = "HEAD" + ("x" * 200) + "TAIL"
out = truncate_text_head_tail(text, 40)
assert "HEAD" in out and "TAIL" in out and "...[truncated" in out
# ---------------------------------------------------------------------------
# bound_parse_payload: every shape stays bounded for the Redis result backend
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_bound_parse_payload_bounds_content_and_chunks():
huge = "Z" * (dr._TEXT_MAX_BYTES * 3)
chunks = [huge for _ in range(dr._MAX_CHUNKS_RETURNED * 2)]
out = bound_parse_payload({"output": "chunks", "content": huge, "chunks": chunks})
assert len(out["content"].encode("utf-8")) <= dr._TEXT_MAX_BYTES + 64
assert len(out["chunks"]) == dr._MAX_CHUNKS_RETURNED
assert out["chunks_truncated"] is True
assert out["total_chunks"] == dr._MAX_CHUNKS_RETURNED * 2
assert all("...[truncated" in c for c in out["chunks"])
@pytest.mark.unit
def test_bound_parse_payload_keeps_structured_for_validation():
structured = {"texts": [{}], "tables": [{}]}
out = bound_parse_payload({"output": "structured", "content": "# ok", "structured": structured})
# structured must survive so the tool's json_schema validation can run on it.
assert out["structured"] == structured
tests/worker/test_parse_document_worker.py
+208
View File
@@ -0,0 +1,208 @@
"""Unit tests for ``application.worker.parse_document_worker``.
The worker re-resolves the artifact through the run-scoped gate (independent of
the tool), reads its bytes, shapes the result, and persists when asked. The DB,
storage, parser, and persistence boundaries are mocked so no live worker / DB is
needed (call the underlying function directly).
"""
from __future__ import annotations
import uuid
from typing import Any, Dict, Optional
import pytest
import application.worker as worker
_ART_ID = str(uuid.uuid4())
class _FakeFile:
def __init__(self, data: bytes) -> None:
self._data = data
def read(self) -> bytes:
return self._data
class _FakeStorage:
def __init__(self, data: bytes = b"%PDF-1.4 fake") -> None:
self._data = data
def get_file(self, path):
return _FakeFile(self._data)
def _patch_repo(monkeypatch, *, found: bool, 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 or (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 or (run is not None and workflow_run_id != run):
return None
return {"id": artifact_id, "current_version": 1, "title": "statement.pdf"}
def get_version(self, artifact_id, version):
return {"filename": "statement.pdf", "storage_path": f"inputs/u/artifacts/{artifact_id}/v1/x.pdf"}
class _Conn:
def __enter__(self):
return object()
def __exit__(self, *exc):
return False
monkeypatch.setattr(worker, "db_readonly", lambda: _Conn())
monkeypatch.setattr(worker, "ArtifactsRepository", _Repo)
monkeypatch.setattr(worker.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage()))
def _patch_parse(monkeypatch, result: Dict[str, Any]):
import application.parser.document_reader as dr
monkeypatch.setattr(dr, "parse_document_bytes", lambda data, filename, **opts: result)
@pytest.mark.unit
def test_requires_parent():
out = worker.parse_document_worker(None, _ART_ID, {}, "u-1", {})
assert out["status"] == "error" and "conversation_id or workflow_run_id" in out["error"]
@pytest.mark.unit
def test_happy_path_shapes_result(monkeypatch):
_patch_repo(monkeypatch, found=True, run="run-1")
_patch_parse(monkeypatch, {"output": "markdown", "content": "# Hi", "truncated": False})
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"output": "markdown", "persist": False}
)
assert out["status"] == "ok"
assert out["content"] == "# Hi"
assert out["output"] == "markdown"
@pytest.mark.unit
def test_cross_run_artifact_is_rejected(monkeypatch):
# The artifact only resolves for run-OTHER; the worker is asked for run-1 -> denied.
_patch_repo(monkeypatch, found=True, run="run-OTHER")
_patch_parse(monkeypatch, {"output": "markdown", "content": "x", "truncated": False})
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"persist": False}
)
assert out["status"] == "error" and "not found in this conversation/run" in out["error"]
@pytest.mark.unit
def test_missing_artifact_is_rejected(monkeypatch):
_patch_repo(monkeypatch, found=False, run="run-1")
_patch_parse(monkeypatch, {"output": "markdown", "content": "x", "truncated": False})
out = worker.parse_document_worker(
None, "ghost", {"workflow_run_id": "run-1"}, "u-1", {"persist": False}
)
assert out["status"] == "error" and "not found in this conversation/run" in out["error"]
@pytest.mark.unit
def test_parse_error_is_surfaced(monkeypatch):
_patch_repo(monkeypatch, found=True, run="run-1")
_patch_parse(monkeypatch, {"error": "unsupported file type '.exe'."})
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"persist": False}
)
assert out["status"] == "error" and "unsupported file type" in out["error"]
@pytest.mark.unit
def test_persist_stores_full_result_and_returns_ref(monkeypatch):
_patch_repo(monkeypatch, found=True, run="run-1")
full = {"output": "structured", "content": "# Big", "structured": {"texts": [{}]}, "truncated": False}
_patch_parse(monkeypatch, full)
captured: Dict[str, Any] = {}
def _fake_persist(**kwargs):
captured.update(kwargs)
return {"artifact_id": "new-art", "version": 1, "filename": "x.json",
"mime_type": "application/json", "size": 10}
import application.sandbox.artifacts_capture as ac
monkeypatch.setattr(ac, "persist_new_artifact", _fake_persist)
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"output": "structured", "persist": True}
)
assert out["status"] == "ok"
assert out["artifact"]["artifact_id"] == "new-art"
# The FULL shaped result is persisted by reference (not just the bounded view).
import json
assert captured["kind"] == "data"
assert json.loads(captured["data"].decode("utf-8")) == full
assert captured["workflow_run_id"] == "run-1"
@pytest.mark.unit
def test_persist_quota_surfaces_as_artifact_error(monkeypatch):
_patch_repo(monkeypatch, found=True, run="run-1")
_patch_parse(monkeypatch, {"output": "markdown", "content": "x", "truncated": False})
import application.sandbox.artifacts_capture as ac
def _quota(**kwargs):
raise ac.QuotaExceeded("artifact storage quota reached")
monkeypatch.setattr(ac, "persist_new_artifact", _quota)
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"persist": True}
)
# Parse still succeeded; quota failure is a non-fatal note.
assert out["status"] == "ok"
assert "artifact" not in out
assert "quota" in out["artifact_error"].lower()
@pytest.mark.unit
def test_result_payload_content_is_bounded(monkeypatch):
_patch_repo(monkeypatch, found=True, run="run-1")
huge = "Z" * 50000
_patch_parse(monkeypatch, {"output": "markdown", "content": huge, "truncated": False})
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"persist": False}
)
# The bounded view ridden back to the (Redis) result backend is far smaller.
assert len(out["content"]) < len(huge)
assert "...[truncated" in out["content"]
@pytest.mark.unit
def test_result_payload_chunks_are_bounded(monkeypatch):
import application.parser.document_reader as dr
_patch_repo(monkeypatch, found=True, run="run-1")
# Many oversized chunks: count is capped AND each chunk is windowed.
huge_chunk = "Z" * 50000
chunks = [huge_chunk for _ in range(dr._MAX_CHUNKS_RETURNED * 3)]
_patch_parse(monkeypatch, {"output": "chunks", "chunks": chunks, "truncated": False})
out = worker.parse_document_worker(
None, _ART_ID, {"workflow_run_id": "run-1"}, "u-1", {"persist": False}
)
assert len(out["chunks"]) == dr._MAX_CHUNKS_RETURNED
assert out["chunks_truncated"] is True
assert out["total_chunks"] == dr._MAX_CHUNKS_RETURNED * 3
assert all("...[truncated" in c for c in out["chunks"])
assert all(len(c) < len(huge_chunk) for c in out["chunks"])