diff --git a/.devcontainer/devc-welcome.md b/.devcontainer/devc-welcome.md index 2ee76adf..ad4103ba 100644 --- a/.devcontainer/devc-welcome.md +++ b/.devcontainer/devc-welcome.md @@ -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: diff --git a/AGENTS.md b/AGENTS.md index 033047c0..959e2b8b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -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: diff --git a/application/agents/tools/document_extractor.py b/application/agents/tools/document_extractor.py deleted file mode 100644 index d42f9ee1..00000000 --- a/application/agents/tools/document_extractor.py +++ /dev/null @@ -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 diff --git a/application/agents/tools/read_document.py b/application/agents/tools/read_document.py new file mode 100644 index 00000000..0d19ea77 --- /dev/null +++ b/application/agents/tools/read_document.py @@ -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 diff --git a/application/agents/tools/tool_manager.py b/application/agents/tools/tool_manager.py index 4b7bcf42..d095b472 100644 --- a/application/agents/tools/tool_manager.py +++ b/application/agents/tools/tool_manager.py @@ -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 ): diff --git a/application/agents/workflows/workflow_engine.py b/application/agents/workflows/workflow_engine.py index a2008085..746a0b11 100644 --- a/application/agents/workflows/workflow_engine.py +++ b/application/agents/workflows/workflow_engine.py @@ -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: diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 2415bc54..eb271870 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -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, diff --git a/application/celeryconfig.py b/application/celeryconfig.py index 82d9ca0d..e3e9942c 100644 --- a/application/celeryconfig.py +++ b/application/celeryconfig.py @@ -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:" diff --git a/application/core/settings.py b/application/core/settings.py index e64cf532..9cf6526b 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -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. diff --git a/application/parser/document_reader.py b/application/parser/document_reader.py new file mode 100644 index 00000000..ec884840 --- /dev/null +++ b/application/parser/document_reader.py @@ -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 diff --git a/application/worker.py b/application/worker.py index fea14c75..6cca0728 100755 --- a/application/worker.py +++ b/application/worker.py @@ -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. diff --git a/deployment/docker-compose-azure.yaml b/deployment/docker-compose-azure.yaml index beb667da..71f675a8 100644 --- a/deployment/docker-compose-azure.yaml +++ b/deployment/docker-compose-azure.yaml @@ -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: diff --git a/deployment/docker-compose-hub.yaml b/deployment/docker-compose-hub.yaml index 6ff29131..b20b1e55 100644 --- a/deployment/docker-compose-hub.yaml +++ b/deployment/docker-compose-hub.yaml @@ -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: diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml index 3294b4e5..7e0ce7fe 100644 --- a/deployment/docker-compose.yaml +++ b/deployment/docker-compose.yaml @@ -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: diff --git a/deployment/k8s/deployments/docsgpt-deploy.yaml b/deployment/k8s/deployments/docsgpt-deploy.yaml index 3772ba51..16159f53 100644 --- a/deployment/k8s/deployments/docsgpt-deploy.yaml +++ b/deployment/k8s/deployments/docsgpt-deploy.yaml @@ -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" diff --git a/deployment/sandbox/Dockerfile b/deployment/sandbox/Dockerfile index b36cf107..39f25748 100644 --- a/deployment/sandbox/Dockerfile +++ b/deployment/sandbox/Dockerfile @@ -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; \ diff --git a/deployment/sandbox/README.md b/deployment/sandbox/README.md index fd587ea1..fe6f1040 100644 --- a/deployment/sandbox/README.md +++ b/deployment/sandbox/README.md @@ -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 diff --git a/scripts/e2e/up.sh b/scripts/e2e/up.sh index 51e7e282..6db92c11 100755 --- a/scripts/e2e/up.sh +++ b/scripts/e2e/up.sh @@ -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" diff --git a/tests/agents/test_workflow_node_attachments.py b/tests/agents/test_workflow_node_attachments.py index 3c2e4731..e9412f88 100644 --- a/tests/agents/test_workflow_node_attachments.py +++ b/tests/agents/test_workflow_node_attachments.py @@ -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) diff --git a/tests/agents/tools/test_document_extractor_live.py b/tests/agents/tools/test_document_extractor_live.py deleted file mode 100644 index 0687de7b..00000000 --- a/tests/agents/tools/test_document_extractor_live.py +++ /dev/null @@ -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, "", "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 deleted file mode 100644 index 0a5ef211..00000000 --- a/tests/agents/tools/test_document_extractor_unit.py +++ /dev/null @@ -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/) 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, "", "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_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 == [] diff --git a/tests/agents/tools/test_read_document_tool.py b/tests/agents/tools/test_read_document_tool.py new file mode 100644 index 00000000..8f9b3a5e --- /dev/null +++ b/tests/agents/tools/test_read_document_tool.py @@ -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"] diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index b1170201..b28b101e 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -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: diff --git a/tests/parser/test_document_reader.py b/tests/parser/test_document_reader.py new file mode 100644 index 00000000..06ced7d3 --- /dev/null +++ b/tests/parser/test_document_reader.py @@ -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 diff --git a/tests/worker/test_parse_document_worker.py b/tests/worker/test_parse_document_worker.py new file mode 100644 index 00000000..62e36829 --- /dev/null +++ b/tests/worker/test_parse_document_worker.py @@ -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"])