From f8bdf998ac1c4a2fed661c1d31e91620e4ab87f6 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 12:50:45 +0100 Subject: [PATCH] Add document extraction tool using Docling Add a document extractor tool with an extract_document action that converts an input artifact (PDF, docx, pptx, ...) into structured JSON by running a fixed Docling program in the sandbox (the document and parameters travel as data, so nothing in them can execute). Output is a compact payload bounded by an input size cap, a head-and-tail markdown window, and per-table caps, can be validated against a JSON schema, and is persisted as a data artifact by reference. Docling is heavy, so it stays out of the base image and ships in an opt-in sandbox image variant. --- .../agents/tools/document_extractor.py | 422 ++++++++++++++ application/agents/tools/tool_manager.py | 2 + application/core/settings.py | 1 + deployment/sandbox/Dockerfile | 10 + deployment/sandbox/README.md | 19 + .../tools/test_document_extractor_live.py | 56 ++ .../tools/test_document_extractor_unit.py | 538 ++++++++++++++++++ 7 files changed, 1048 insertions(+) create mode 100644 application/agents/tools/document_extractor.py create mode 100644 tests/agents/tools/test_document_extractor_live.py create mode 100644 tests/agents/tools/test_document_extractor_unit.py diff --git a/application/agents/tools/document_extractor.py b/application/agents/tools/document_extractor.py new file mode 100644 index 00000000..94f76959 --- /dev/null +++ b/application/agents/tools/document_extractor.py @@ -0,0 +1,422 @@ +"""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.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 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" +) + + +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": "Artifact id (from this conversation/run) of the document to extract.", + }, + "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 + input_path = f"extract/{token}/inputs/{loaded['filename']}" + params_path = f"extract/{token}/params.json" + result_path = f"extract/{token}/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: + 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: + ref = self._persist(extracted, loaded["title"]) + if ref is not None: + self._last_artifact_id = ref["artifact_id"] + payload["artifact"] = ref + return payload + + # ------------------------------------------------------------------ + # Input / payload helpers + # ------------------------------------------------------------------ + def _load_input(self, artifact_id: str) -> Dict[str, Any]: + """Fetch the parent-scoped input artifact's current-version bytes; never cross-tenant.""" + try: + with db_readonly() as conn: + repo = ArtifactsRepository(conn) + artifact = repo.get_artifact_in_parent( + artifact_id, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + ) + if artifact is None: + return {"error": f"input artifact {artifact_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 {artifact_id}."} + if not version or not version.get("storage_path"): + return {"error": f"input artifact {artifact_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 {artifact_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 diff --git a/application/agents/tools/tool_manager.py b/application/agents/tools/tool_manager.py index 579d5353..4b7bcf42 100644 --- a/application/agents/tools/tool_manager.py +++ b/application/agents/tools/tool_manager.py @@ -39,6 +39,7 @@ class ToolManager: "remote_device", "code_executor", "artifact_generator", + "document_extractor", } and user_id ): @@ -60,6 +61,7 @@ class ToolManager: "remote_device", "code_executor", "artifact_generator", + "document_extractor", } and user_id ): diff --git a/application/core/settings.py b/application/core/settings.py index 4b5e9f07..5ebb0cad 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -337,6 +337,7 @@ class Settings(BaseSettings): SANDBOX_HTTP_TIMEOUT: int = 10 # fixed cap (s) for REST control calls (create/delete/alive/interrupt) 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 # Runner container resource caps — consumed by the docsgpt-sandbox compose # service (deployment/sandbox), not by the app client. cgroup CPU/mem caps # are part of the untrusted-code security boundary. diff --git a/deployment/sandbox/Dockerfile b/deployment/sandbox/Dockerfile index 8e36addf..e12fcd2b 100644 --- a/deployment/sandbox/Dockerfile +++ b/deployment/sandbox/Dockerfile @@ -32,6 +32,16 @@ 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. +ARG INSTALL_DOCLING=false +RUN if [ "$INSTALL_DOCLING" = "true" ]; then \ + pip install --no-cache-dir docling==2.8.3; \ + fi + USER sandbox WORKDIR /home/sandbox diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md index 0a46116c..0b60b479 100644 --- a/deployment/sandbox/README.md +++ b/deployment/sandbox/README.md @@ -43,6 +43,25 @@ The `docsgpt-sandbox` service is defined in `deployment/docker-compose.yaml` on an internal-only network. The backend and worker reach it at `http://docsgpt-sandbox:8888`. +## Document extraction variant (Docling) + +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: + +```bash +docker build -t docsgpt-sandbox-extract \ + --build-arg INSTALL_DOCLING=true deployment/sandbox +docker run --rm -p 8888:8888 docsgpt-sandbox-extract +``` + +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. + ## Hardening (separate slice — not in this image) Egress/SSRF blocks (drop RFC1918, link-local, `169.254.169.254`), the gVisor diff --git a/tests/agents/tools/test_document_extractor_live.py b/tests/agents/tools/test_document_extractor_live.py new file mode 100644 index 00000000..0687de7b --- /dev/null +++ b/tests/agents/tools/test_document_extractor_live.py @@ -0,0 +1,56 @@ +"""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, "", "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) diff --git a/tests/agents/tools/test_document_extractor_unit.py b/tests/agents/tools/test_document_extractor_unit.py new file mode 100644 index 00000000..60f84dba --- /dev/null +++ b/tests/agents/tools/test_document_extractor_unit.py @@ -0,0 +1,538 @@ +"""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 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 + +# 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] = [] + + def open(self, session_id, ttl=None): + self.opened.append((session_id, ttl)) + return session_id + + 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 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-1", 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_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-1") # 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 + + +# --------------------------------------------------------------------------- +# 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-1", 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_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-1", 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-1", 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-1", 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-1", 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-1", 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-1", 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-1" 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-1", 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, "", "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, "", "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-1", 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-1", 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 == []