From d6d1930cd8c73be9d127f4e31976c52dc16322c0 Mon Sep 17 00:00:00 2001 From: Simon Coombes Date: Tue, 28 Jul 2026 05:41:15 -0400 Subject: [PATCH 1/3] Update Unstructured Transform MCP tool names Signed-off-by: Simon Coombes --- examples/unstructured_transform_mcp/README.md | 26 +++++----- .../configs/config.yml | 6 +-- .../register.py | 22 ++++---- .../tests/test_transform_document.py | 52 +++++++++---------- 4 files changed, 53 insertions(+), 53 deletions(-) diff --git a/examples/unstructured_transform_mcp/README.md b/examples/unstructured_transform_mcp/README.md index 0cdf5dc..c7c396f 100644 --- a/examples/unstructured_transform_mcp/README.md +++ b/examples/unstructured_transform_mcp/README.md @@ -37,9 +37,9 @@ This example is hosted in the examples repository because it requires an Unstruc The Unstructured Transform MCP server exposes an asynchronous job protocol as four tools: 1. `request_file_upload_url`: Returns a pre-signed upload URL and a file reference for a local file. -2. `transform_files`: Starts a transform job for one or more file references (or public HTTP(S) URLs) and returns a job ID. -3. `check_transform_status`: Reports whether the job is `SCHEDULED`, `IN_PROGRESS`, or `COMPLETED` (any other state means the job failed, and the function reports it as an error). -4. `get_transform_results`: Returns a pre-signed download URL for the Markdown output of each transformed file. +2. `start_transform_job`: Starts a transform job for one or more file references (or public HTTP(S) URLs) and returns a job ID. +3. `check_job_status`: Reports whether the job is `SCHEDULED`, `IN_PROGRESS`, or `COMPLETED` (any other state means the job failed, and the function reports it as an error). +4. `get_job_results`: Returns a pre-signed download URL for the Markdown output of each transformed file. Two steps of the protocol are plain HTTP transfers rather than MCP calls: uploading the raw document bytes to the pre-signed upload URL and downloading the Markdown from the pre-signed download URL. An agent cannot perform those byte transfers with MCP tools alone, and letting an LLM drive the polling loop is slow and unreliable. The `transform_document` function in `src/nat_unstructured_transform_mcp/register.py` therefore performs the whole sequence deterministically: @@ -47,13 +47,13 @@ Two steps of the protocol are plain HTTP transfers rather than MCP calls: upload agent -> transform_document(source) |-- request_file_upload_url (MCP) # local files only |-- PUT raw bytes to upload URL # plain HTTP, no bearer token - |-- transform_files (MCP) - |-- check_transform_status (MCP) # polled until COMPLETED - |-- get_transform_results (MCP) + |-- start_transform_job (MCP) + |-- check_job_status (MCP) # polled until COMPLETED + |-- get_job_results (MCP) `-- GET Markdown from download URL # plain HTTP, no bearer token ``` -The function accepts either a local file path or a public HTTP(S) URL (public URLs are passed directly to `transform_files`, skipping the upload). Transforms take from a few seconds up to several minutes depending on page count, and the maximum file size is 50 MB. Each `transform_document` call processes a single document, and the example always requests the default Markdown output; the Transform service also supports element JSON, HTML, and plain-text output, which would require extending `transform_document`. +The function accepts either a local file path or a public HTTP(S) URL (public URLs are passed directly to `start_transform_job`, skipping the upload). Transforms take from a few seconds up to several minutes depending on page count, and the maximum file size is 50 MB. Each `transform_document` call processes a single document, and the example always requests the default Markdown output; the Transform service also supports element JSON, HTML, and plain-text output, which would require extending `transform_document`. > [!IMPORTANT] > Trust boundary: for a local path, `transform_document` reads that file and uploads its contents to the hosted Transform service. Because the path comes from the agent, a crafted prompt could point it at a sensitive file (for example a private key or a credentials file) and cause that file to leave the host. Run this example with documents and prompts you trust, and if you adapt it for untrusted input, restrict the accepted paths to a designated directory. @@ -144,9 +144,9 @@ function_groups: auth_provider: unstructured_auth include: - request_file_upload_url - - transform_files - - check_transform_status - - get_transform_results + - start_transform_job + - check_job_status + - get_job_results authentication: unstructured_auth: @@ -175,9 +175,9 @@ function_groups: Authorization: "Bearer ${UNSTRUCTURED_API_KEY}" include: - request_file_upload_url - - transform_files - - check_transform_status - - get_transform_results + - start_transform_job + - check_job_status + - get_job_results ``` Keep the `include` list in this variant as well; it provides the same fail-fast contract check described above. diff --git a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml index 5c54a71..2987cfb 100644 --- a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml +++ b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml @@ -33,9 +33,9 @@ function_groups: # The four tools that make up the asynchronous transform protocol. include: - request_file_upload_url - - transform_files - - check_transform_status - - get_transform_results + - start_transform_job + - check_job_status + - get_job_results authentication: unstructured_auth: diff --git a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py index 9bf91f3..128bb14 100644 --- a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py +++ b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py @@ -44,7 +44,7 @@ # Maximum file size accepted by the Unstructured Transform service (50 MB). MAX_FILE_SIZE_BYTES = 52_428_800 -# Job states reported by the check_transform_status MCP tool which mean "keep waiting". +# Job states reported by the check_job_status MCP tool which mean "keep waiting". _PENDING_JOB_STATES = frozenset({"SCHEDULED", "IN_PROGRESS"}) _COMPLETED_JOB_STATE = "COMPLETED" @@ -90,9 +90,9 @@ class TransformTools(typing.NamedTuple): """The four Unstructured Transform MCP tools, resolved from the function group.""" request_file_upload_url: Function - transform_files: Function - check_transform_status: Function - get_transform_results: Function + start_transform_job: Function + check_job_status: Function + get_job_results: Function def resolve_tools(group_functions: dict[str, Function]) -> TransformTools: @@ -169,7 +169,7 @@ async def _upload_source(tools: TransformTools, http_client: httpx.AsyncClient, max_file_size_bytes: int) -> str: """Return a file reference for the document, uploading local files first. - Public HTTP or HTTPS URLs are passed through unchanged because the transform_files + Public HTTP or HTTPS URLs are passed through unchanged because the start_transform_job MCP tool accepts them directly. Local files are uploaded to a pre-signed URL minted by the request_file_upload_url MCP tool. """ @@ -226,7 +226,7 @@ async def _wait_for_job(tools: TransformTools, while True: try: - status_payload = await _invoke_tool(tools.check_transform_status, job_id=job_id) + status_payload = await _invoke_tool(tools.check_job_status, job_id=job_id) except MCPToolError: raise except RuntimeError: @@ -262,7 +262,7 @@ async def _fetch_results(tools: TransformTools, """Fetch the transform results, tolerating the brief window after the job reports completion. The status endpoint can report ``COMPLETED`` slightly before the results are - materialized, in which case get_transform_results returns a ``job_not_complete`` + materialized, in which case get_job_results returns a ``job_not_complete`` error. Retry within the overall transform deadline. Transport-level blips are retried the same way as in the status loop. """ @@ -271,7 +271,7 @@ async def _fetch_results(tools: TransformTools, while True: try: - return await _invoke_tool(tools.get_transform_results, job_id=job_id) + return await _invoke_tool(tools.get_job_results, job_id=job_id) except MCPToolError as e: if e.code != "job_not_complete": raise @@ -309,8 +309,8 @@ async def transform_source(tools: TransformTools, config: TransformDocumentConfi async with httpx.AsyncClient(timeout=config.http_timeout_seconds) as http_client: file_ref = await _upload_source(tools, http_client, source, config.max_file_size_bytes) - job = await _invoke_tool(tools.transform_files, file_refs=[file_ref]) - job_id = _require(job, "job_id", "transform_files") + job = await _invoke_tool(tools.start_transform_job, file_refs=[file_ref]) + job_id = _require(job, "job_id", "start_transform_job") logger.info("Started transform job %s for '%s'", job_id, source) # The status wait and the results fetch share one overall deadline. @@ -331,7 +331,7 @@ async def transform_source(tools: TransformTools, config: TransformDocumentConfi f"{str(first_file)[:300]!r}") # The download URL is pre-signed as well, so no bearer token is sent here either. - download = await http_client.get(_require(first_file, "download_url", "get_transform_results")) + download = await http_client.get(_require(first_file, "download_url", "get_job_results")) _check_http_response(download, "result download") return TransformResult(markdown=download.text, diff --git a/examples/unstructured_transform_mcp/tests/test_transform_document.py b/examples/unstructured_transform_mcp/tests/test_transform_document.py index e54da7a..48f5579 100644 --- a/examples/unstructured_transform_mcp/tests/test_transform_document.py +++ b/examples/unstructured_transform_mcp/tests/test_transform_document.py @@ -79,14 +79,14 @@ async def _request_file_upload_url(**_kwargs) -> str: "file_ref": self.file_ref, }) - async def _transform_files(**_kwargs) -> str: + async def _start_transform_job(**_kwargs) -> str: return json.dumps({"job_id": self.job_id, "status": "SCHEDULED"}) - async def _check_transform_status(**_kwargs) -> str: + async def _check_job_status(**_kwargs) -> str: status = self.statuses.pop(0) if len(self.statuses) > 1 else self.statuses[0] return json.dumps({"job_id": self.job_id, "status": status}) - async def _get_transform_results(**_kwargs) -> str: + async def _get_job_results(**_kwargs) -> str: return json.dumps({ "job_id": self.job_id, @@ -100,9 +100,9 @@ async def _get_transform_results(**_kwargs) -> str: self.group_functions = { f"{_GROUP}__request_file_upload_url": _FakeTool("request_file_upload_url", _request_file_upload_url), - f"{_GROUP}__transform_files": _FakeTool("transform_files", _transform_files), - f"{_GROUP}__check_transform_status": _FakeTool("check_transform_status", _check_transform_status), - f"{_GROUP}__get_transform_results": _FakeTool("get_transform_results", _get_transform_results), + f"{_GROUP}__start_transform_job": _FakeTool("start_transform_job", _start_transform_job), + f"{_GROUP}__check_job_status": _FakeTool("check_job_status", _check_job_status), + f"{_GROUP}__get_job_results": _FakeTool("get_job_results", _get_job_results), } @property @@ -146,9 +146,9 @@ async def test_local_file_transform(httpserver: HTTPServer, fast_config: Transfo assert upload_request_args["content_type"] == "application/pdf" assert upload_request_args["size_bytes"] == sample_document.stat().st_size - assert service.tool("transform_files").calls == [{"file_refs": [service.file_ref]}] - assert all(call == {"job_id": service.job_id} for call in service.tool("check_transform_status").calls) - assert service.tool("get_transform_results").calls == [{"job_id": service.job_id}] + assert service.tool("start_transform_job").calls == [{"file_refs": [service.file_ref]}] + assert all(call == {"job_id": service.job_id} for call in service.tool("check_job_status").calls) + assert service.tool("get_job_results").calls == [{"job_id": service.job_id}] put_requests = [request for request, _ in httpserver.log if request.method == "PUT"] assert len(put_requests) == 1 @@ -159,7 +159,7 @@ async def test_local_file_transform(httpserver: HTTPServer, fast_config: Transfo async def test_public_url_skips_upload(httpserver: HTTPServer, fast_config: TransformDocumentConfig): - """A public URL is passed straight to transform_files, skipping the upload step.""" + """A public URL is passed straight to start_transform_job, skipping the upload step.""" service = _FakeTransformService(httpserver, statuses=["COMPLETED"]) url = "https://example.com/whitepaper.pdf" @@ -167,7 +167,7 @@ async def test_public_url_skips_upload(httpserver: HTTPServer, fast_config: Tran assert result.markdown == service.markdown assert service.tool("request_file_upload_url").calls == [] - assert service.tool("transform_files").calls == [{"file_refs": [url]}] + assert service.tool("start_transform_job").calls == [{"file_refs": [url]}] async def test_failed_job_raises(httpserver: HTTPServer, fast_config: TransformDocumentConfig, sample_document: Path): @@ -216,7 +216,7 @@ async def test_results_not_ready_retries_until_available(httpserver: HTTPServer, sample_document: Path): """The status endpoint can report COMPLETED slightly before the results exist.""" service = _FakeTransformService(httpserver, statuses=["COMPLETED"]) - real_results_tool = service.tool("get_transform_results") + real_results_tool = service.tool("get_job_results") not_ready_responses = 2 attempts = [] @@ -230,7 +230,7 @@ async def _flaky_results(**kwargs) -> str: }) return await real_results_tool._handler(**kwargs) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _flaky_results) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _flaky_results) result = await transform_source(service.tools, fast_config, str(sample_document)) @@ -247,7 +247,7 @@ async def test_tool_error_payload_raises(httpserver: HTTPServer, async def _unauthorized(**_kwargs) -> str: return json.dumps({"error": {"code": "unauthorized", "message": "Invalid API key", "status": 401}}) - service.group_functions[f"{_GROUP}__transform_files"] = _FakeTool("transform_files", _unauthorized) + service.group_functions[f"{_GROUP}__start_transform_job"] = _FakeTool("start_transform_job", _unauthorized) with pytest.raises(RuntimeError, match="unauthorized"): await transform_source(service.tools, fast_config, str(sample_document)) @@ -271,9 +271,9 @@ async def _error_text(**_kwargs) -> str: def test_resolve_tools_reports_missing(httpserver: HTTPServer): """resolve_tools reports which required tool is absent from the group.""" service = _FakeTransformService(httpserver, statuses=["COMPLETED"]) - del service.group_functions[f"{_GROUP}__get_transform_results"] + del service.group_functions[f"{_GROUP}__get_job_results"] - with pytest.raises(ValueError, match="get_transform_results"): + with pytest.raises(ValueError, match="get_job_results"): resolve_tools(service.group_functions) # type: ignore[arg-type] @@ -310,7 +310,7 @@ async def test_results_error_not_retried(httpserver: HTTPServer, async def _unauthorized(**_kwargs) -> str: return json.dumps({"error": {"code": "unauthorized", "message": "Invalid API key", "status": 401}}) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _unauthorized) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _unauthorized) with pytest.raises(RuntimeError, match="unauthorized"): await transform_source(service.tools, fast_config, str(sample_document)) @@ -323,7 +323,7 @@ async def test_results_fetch_deadline(httpserver: HTTPServer, sample_document: P async def _never_ready(**_kwargs) -> str: return json.dumps({"error": {"code": "job_not_complete", "message": "not yet"}}) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _never_ready) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _never_ready) config = TransformDocumentConfig(poll_interval_seconds=0.01, transform_timeout_seconds=0.05, http_timeout_seconds=5.0) @@ -341,7 +341,7 @@ async def test_empty_files_list_raises(httpserver: HTTPServer, async def _no_files(**_kwargs) -> str: return json.dumps({"job_id": service.job_id, "files": []}) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _no_files) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _no_files) with pytest.raises(RuntimeError, match="returned no files"): await transform_source(service.tools, fast_config, str(sample_document)) @@ -382,7 +382,7 @@ async def test_transient_status_blip_is_retried(httpserver: HTTPServer, sample_document: Path): """A momentary transport failure during status polling must not abandon the job.""" service = _FakeTransformService(httpserver, statuses=["IN_PROGRESS", "COMPLETED"]) - real_status_tool = service.tool("check_transform_status") + real_status_tool = service.tool("check_job_status") blips = [] async def _flaky_status(**kwargs) -> str: @@ -391,7 +391,7 @@ async def _flaky_status(**kwargs) -> str: return "MCPToolClient tool call failed: connection reset" return await real_status_tool._handler(**kwargs) - service.group_functions[f"{_GROUP}__check_transform_status"] = _FakeTool("check_transform_status", _flaky_status) + service.group_functions[f"{_GROUP}__check_job_status"] = _FakeTool("check_job_status", _flaky_status) result = await transform_source(service.tools, fast_config, str(sample_document)) @@ -410,7 +410,7 @@ async def _always_broken(**_kwargs) -> str: attempts.append(True) return "MCPToolClient tool call failed: connection reset" - service.group_functions[f"{_GROUP}__check_transform_status"] = _FakeTool("check_transform_status", _always_broken) + service.group_functions[f"{_GROUP}__check_job_status"] = _FakeTool("check_job_status", _always_broken) with pytest.raises(RuntimeError, match="Expected a JSON payload"): await transform_source(service.tools, fast_config, str(sample_document)) @@ -450,7 +450,7 @@ async def test_status_error_envelope_raises(httpserver: HTTPServer, async def _status_error(**_kwargs) -> str: return json.dumps({"error": {"code": "job_not_found", "message": "Job not found", "status": 404}}) - service.group_functions[f"{_GROUP}__check_transform_status"] = _FakeTool("check_transform_status", _status_error) + service.group_functions[f"{_GROUP}__check_job_status"] = _FakeTool("check_job_status", _status_error) with pytest.raises(RuntimeError, match="job_not_found"): await transform_source(service.tools, fast_config, str(sample_document)) @@ -461,7 +461,7 @@ async def test_transient_results_blip_is_retried(httpserver: HTTPServer, sample_document: Path): """A momentary transport failure during the results fetch is retried.""" service = _FakeTransformService(httpserver, statuses=["COMPLETED"]) - real_results_tool = service.tool("get_transform_results") + real_results_tool = service.tool("get_job_results") blips = [] async def _flaky_results(**kwargs) -> str: @@ -470,7 +470,7 @@ async def _flaky_results(**kwargs) -> str: return "MCPToolClient tool call failed: connection reset" return await real_results_tool._handler(**kwargs) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _flaky_results) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _flaky_results) result = await transform_source(service.tools, fast_config, str(sample_document)) @@ -487,7 +487,7 @@ async def test_non_dict_files_entry_raises(httpserver: HTTPServer, async def _string_files(**_kwargs) -> str: return json.dumps({"job_id": service.job_id, "files": ["not-a-mapping"]}) - service.group_functions[f"{_GROUP}__get_transform_results"] = _FakeTool("get_transform_results", _string_files) + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _string_files) with pytest.raises(RuntimeError, match="unexpected files entry"): await transform_source(service.tools, fast_config, str(sample_document)) From 165fad9f8b7015924e5c55e5a0c845103938aa4a Mon Sep 17 00:00:00 2001 From: Simon Coombes Date: Tue, 28 Jul 2026 07:08:53 -0400 Subject: [PATCH 2/3] Add structured data extraction to the Unstructured Transform example Signed-off-by: Simon Coombes --- examples/unstructured_transform_mcp/README.md | 88 ++- .../configs/config.yml | 16 +- .../register.py | 317 ++++++++++- .../tests/test_extract_structured_data.py | 512 ++++++++++++++++++ .../tests/test_workflow_integration.py | 51 ++ 5 files changed, 956 insertions(+), 28 deletions(-) create mode 100644 examples/unstructured_transform_mcp/tests/test_extract_structured_data.py diff --git a/examples/unstructured_transform_mcp/README.md b/examples/unstructured_transform_mcp/README.md index c7c396f..b780f68 100644 --- a/examples/unstructured_transform_mcp/README.md +++ b/examples/unstructured_transform_mcp/README.md @@ -21,7 +21,7 @@ limitations under the License. This example demonstrates how the NVIDIA NeMo Agent Toolkit connects to a third-party remote MCP server that is protected by static bearer-token authentication: the hosted [Unstructured Transform](https://transform.unstructured.io/get-started) service, which converts documents (PDF, DOCX, PPTX, XLSX, HTML, images, and 40+ other formats) into clean Markdown that agents can reason over. -It also demonstrates a useful composition pattern: the remote MCP server exposes an asynchronous, multi-step protocol, and this example wraps that protocol in a single deterministic custom function so the agent only needs one reliable tool call. +It also demonstrates a useful composition pattern: the remote MCP server exposes an asynchronous, multi-step protocol, and this example wraps that protocol in deterministic custom functions so the agent only needs one reliable tool call. Two functions are registered, one that converts a document to Markdown and one that extracts named fields from it as JSON. This example is hosted in the examples repository because it requires an Unstructured API key and depends on an external MCP server whose data, schema, availability, and responses are not controlled by the toolkit. It is a reference integration, and it targets NeMo Agent Toolkit 1.8. @@ -30,17 +30,23 @@ This example is hosted in the examples repository because it requires an Unstruc - **Bearer-token MCP authentication:** Uses the `api_key` authentication provider with `auth_scheme: Bearer` to authenticate against a remote MCP server with a static API key supplied through an environment variable. The other MCP examples in the NeMo Agent Toolkit cover unauthenticated servers and OAuth2 flows; this example covers the common "API key in a header" case. - **Remote MCP client over streamable HTTP:** Declares the Transform server as an `mcp_client` function group using the `streamable-http` transport. - **Deterministic composition of MCP tools:** A custom function (`transform_document`) resolves the four Transform MCP tools from the function group and orchestrates the upload, transform, poll, and download flow in plain Python, exposing one dependable tool to the ReAct agent. +- **Chaining two async jobs:** A second function (`extract_structured_data`) shows the same pattern over a dependent pair of jobs, parsing a document and then extracting named fields from the parse output as JSON. - **Document parsing for agents:** Turns binary documents into Markdown the LLM can summarize, query, and extract from. ## How It Works -The Unstructured Transform MCP server exposes an asynchronous job protocol as four tools: +The Unstructured Transform MCP server exposes an asynchronous job protocol. Four tools cover parsing: 1. `request_file_upload_url`: Returns a pre-signed upload URL and a file reference for a local file. 2. `start_transform_job`: Starts a transform job for one or more file references (or public HTTP(S) URLs) and returns a job ID. 3. `check_job_status`: Reports whether the job is `SCHEDULED`, `IN_PROGRESS`, or `COMPLETED` (any other state means the job failed, and the function reports it as an error). 4. `get_job_results`: Returns a pre-signed download URL for the Markdown output of each transformed file. +Two more cover structured data extraction, described in [Structured Data Extraction](#structured-data-extraction) below: + +5. `suggest_extraction_schema_for_file`: Drafts a JSON Schema from one parsed document. +6. `start_extraction_job`: Runs an extraction against a JSON Schema and returns a job ID. Status and results come from the same two tools as a parse. + Two steps of the protocol are plain HTTP transfers rather than MCP calls: uploading the raw document bytes to the pre-signed upload URL and downloading the Markdown from the pre-signed download URL. An agent cannot perform those byte transfers with MCP tools alone, and letting an LLM drive the polling loop is slow and unreliable. The `transform_document` function in `src/nat_unstructured_transform_mcp/register.py` therefore performs the whole sequence deterministically: ```text @@ -56,7 +62,30 @@ agent -> transform_document(source) The function accepts either a local file path or a public HTTP(S) URL (public URLs are passed directly to `start_transform_job`, skipping the upload). Transforms take from a few seconds up to several minutes depending on page count, and the maximum file size is 50 MB. Each `transform_document` call processes a single document, and the example always requests the default Markdown output; the Transform service also supports element JSON, HTML, and plain-text output, which would require extending `transform_document`. > [!IMPORTANT] -> Trust boundary: for a local path, `transform_document` reads that file and uploads its contents to the hosted Transform service. Because the path comes from the agent, a crafted prompt could point it at a sensitive file (for example a private key or a credentials file) and cause that file to leave the host. Run this example with documents and prompts you trust, and if you adapt it for untrusted input, restrict the accepted paths to a designated directory. +> Trust boundary: for a local path, `transform_document` reads that file and uploads its contents to the hosted Transform service. Because the path comes from the agent, a crafted prompt could point it at a sensitive file (for example a private key or a credentials file) and cause that file to leave the host. Run this example with documents and prompts you trust, and if you adapt it for untrusted input, restrict the accepted paths to a designated directory. The same applies to `extract_structured_data`, which takes a path the same way. + +### Structured Data Extraction + +`extract_structured_data` returns named fields as JSON instead of the whole document as text, which is what you want for invoices, forms, and contracts. It shows the composition pattern over two *dependent* jobs, because the extractor reads the Element JSON a parse produces rather than a raw file: + +```text +agent -> extract_structured_data(source, extraction_schema, guidance) + |-- request_file_upload_url (MCP) # local files only + |-- PUT raw bytes to upload URL # plain HTTP, no bearer token + |-- start_transform_job (MCP) # with a partition strategy + |-- check_job_status / get_job_results (MCP) # polled; yields output_ref + |-- suggest_extraction_schema_for_file (MCP) # only when no schema is given + |-- start_extraction_job (MCP) # consumes the output_ref + `-- check_job_status / get_job_results (MCP) # polled; returns JSON inline +``` + +Three details of that flow are worth knowing if you adapt it: + +- **The `output_ref` is the handle, not the rendered output.** Each file entry in the results of a completed parse carries a durable `output_ref` alongside its `download_url`, and it is present whichever output format you render. The function passes that reference straight to the extractor and never downloads the Markdown, which saves a transfer of output nobody reads. +- **Parse fidelity sets the quality ceiling.** The extractor can only surface what the parse captured, so the function picks the partition strategy from the input type: `vlm` for PDFs, images, and PowerPoint files, and `fast` for everything else, where the higher-fidelity strategies silently fall back anyway. If a result comes back sparse, re-parsing with `hi_res` plus the `image_description`, `generative_ocr`, and `table_to_html` enrichment steps is the documented next step. +- **Results keep their provenance.** Each record wraps `extracted_data` with `filename`, `filetype`, `processed_date_utc`, and `source_file_uri`. The function returns the whole wrapper rather than the bare data, since that is what ties a record back to the document it came from. + +The schema is optional. Pass `extraction_schema` as a JSON Schema encoded in a JSON string to fix the output shape, or leave it out and the server drafts one from the document itself. `guidance` is free text that steers both the drafted schema and how fields are filled. The agent-facing argument is deliberately named `extraction_schema` rather than `schema`, because these parameter names become fields of a generated Pydantic model and `schema` shadows a `BaseModel` attribute. ## Prerequisites @@ -130,6 +159,39 @@ The section headings of the transformed document are: - # References ``` +### Extract Structured Data + +Ask for named fields instead, and the agent picks `extract_structured_data`: + +```bash +nat run --config_file examples/unstructured_transform_mcp/configs/config.yml \ + --input "Extract each plant's name, light needs, watering instructions, and humidity level from https://docs.unstructured.io/img/pipelines/data-extractor/house-plant-care.png as JSON." +``` + +The agent parses the image, has the server draft a schema (no schema was supplied), runs the extraction, and returns one record per document: + +```text +Workflow Result: +[ + { + "filename": "house-plant-care.png", + "filetype": "image/png", + "processed_date_utc": "2026-07-28T10:45:04.158141Z", + "source_file_uri": "u10d://output/_house-plant-care.json", + "extracted_data": { + "plants": [ + { + "plant_name": "MONSTERA DELICIOSA", + "light_requirements": "Bright Indirect - Some direct", + "watering_instructions": "Water when 80% dry", + "humidity_level": "Low - Medium" + } + ] + } + } +] +``` + ## Configuration Details The complete configuration is in `configs/config.yml`. The MCP client and authentication sections are the interesting parts: @@ -145,6 +207,8 @@ function_groups: include: - request_file_upload_url - start_transform_job + - suggest_extraction_schema_for_file + - start_extraction_job - check_job_status - get_job_results @@ -157,8 +221,8 @@ authentication: - The `api_key` authentication provider attaches `Authorization: Bearer ` to every request the MCP client makes, including the initial handshake. - `${UNSTRUCTURED_API_KEY}` is interpolated from the environment when the configuration is loaded. See [workflow configuration](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/main/docs/source/build-workflows/workflow-configuration.md) for the interpolation syntax. -- The `include` list documents the four tools the example depends on and fails fast if the server stops exposing any of them. -- The `transform_document` function references the function group through its `mcp_group` setting, so the group does not need to appear in the workflow `tool_names` and the agent never sees the low-level tools. +- The `include` list documents the six tools the example depends on and fails fast if the server stops exposing any of them. +- Both functions reference the function group through their `mcp_group` setting, so the group does not need to appear in the workflow `tool_names` and the agent never sees the low-level tools. ### Alternative: Custom Headers @@ -176,6 +240,8 @@ function_groups: include: - request_file_upload_url - start_transform_job + - suggest_extraction_schema_for_file + - start_extraction_job - check_job_status - get_job_results ``` @@ -194,6 +260,17 @@ The `transform_document` function accepts a few settings in the `functions` sect | `max_file_size_bytes` | `52428800` | Maximum document size accepted by the Transform service (50 MB at the time of writing). | | `max_output_characters` | `50000` | Truncates the returned Markdown to protect the context window of the agent; a short truncation notice is appended. | +`extract_structured_data` takes the same settings, except that it splits the job timeout in two because it runs a parse and then an extraction: + +| Setting | Default | Purpose | +|---|---|---| +| `poll_interval_seconds` | `5.0` | Delay between job status checks. | +| `parse_timeout_seconds` | `900.0` | Maximum time to wait for the parse that precedes the extraction. | +| `extraction_timeout_seconds` | `600.0` | Maximum time to wait for the extraction job itself. | +| `http_timeout_seconds` | `120.0` | Timeout for the raw upload request. | +| `max_file_size_bytes` | `52428800` | Maximum document size accepted by the Transform service (50 MB at the time of writing). | +| `max_output_characters` | `50000` | Truncates the returned JSON to protect the context window of the agent. Truncated output is no longer valid JSON, so the appended notice says so. | + ## Testing Unit tests mock the MCP tools and the HTTP transfers, so they run without network access or credentials: @@ -221,3 +298,4 @@ pytest --run_integration --run_slow examples/unstructured_transform_mcp/tests - [`kaggle_mcp`](../kaggle_mcp/README.md): Another remote MCP server reached over `streamable-http` with bearer-token authentication. - [MCP client documentation](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/main/docs/source/build-workflows/mcp-client.md): All `mcp_client` configuration options. - [API authentication documentation](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/main/docs/source/components/auth/api-authentication.md): Details of the `api_key` authentication provider. +- [Unstructured structured data extraction](https://docs.unstructured.io/transform/sde): How the extractor works, and prompt patterns for it. diff --git a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml index 2987cfb..f510ec5 100644 --- a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml +++ b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/configs/config.yml @@ -22,6 +22,16 @@ functions: poll_interval_seconds: 5.0 transform_timeout_seconds: 900.0 + # Structured data extraction. Extraction runs on the Element JSON a parse produces, so + # this function parses the document first and then chains a second asynchronous job onto + # it, which is why it carries its own two timeouts. + extract_structured_data: + _type: extract_structured_data + mcp_group: unstructured_transform + poll_interval_seconds: 5.0 + parse_timeout_seconds: 900.0 + extraction_timeout_seconds: 600.0 + function_groups: unstructured_transform: _type: mcp_client @@ -30,10 +40,13 @@ function_groups: # The Unstructured Transform MCP server is served at the root path. url: https://mcp.transform.unstructured.io auth_provider: unstructured_auth - # The four tools that make up the asynchronous transform protocol. + # The four tools that make up the asynchronous transform protocol, plus the two + # structured data extraction tools. include: - request_file_upload_url - start_transform_job + - suggest_extraction_schema_for_file + - start_extraction_job - check_job_status - get_job_results @@ -58,6 +71,7 @@ workflow: _type: react_agent tool_names: - transform_document + - extract_structured_data llm_name: nim_llm verbose: true retry_parsing_errors: true diff --git a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py index 128bb14..5eccc5b 100644 --- a/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py +++ b/examples/unstructured_transform_mcp/src/nat_unstructured_transform_mcp/register.py @@ -12,12 +12,18 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Register the ``transform_document`` function. +"""Register the ``transform_document`` and ``extract_structured_data`` functions. The Unstructured Transform MCP server exposes an asynchronous, multi-step protocol (request an upload URL, upload the raw bytes over plain HTTP, start a transform job, -poll for completion, download the result). This module composes those MCP tools into a -single deterministic function so the agent only needs one reliable tool call. +poll for completion, download the result). ``transform_document`` composes those MCP tools +into a single deterministic function so the agent only needs one reliable tool call. + +``extract_structured_data`` adds the server's structured data extraction tools on top of +that flow. Extraction consumes the Element JSON a parse produces rather than a raw file, +so it is a second asynchronous job chained onto the first: parse the document, take the +``output_ref`` of the completed parse, optionally have the server draft a JSON Schema for +it, then run the extraction and poll it with the same status and results tools. """ import asyncio @@ -26,6 +32,7 @@ import mimetypes import typing from pathlib import Path +from urllib.parse import urlsplit import httpx from pydantic import BaseModel @@ -95,7 +102,29 @@ class TransformTools(typing.NamedTuple): get_job_results: Function -def resolve_tools(group_functions: dict[str, Function]) -> TransformTools: +class ExtractionTools(typing.NamedTuple): + """The Transform MCP tools needed to parse a document and then extract from it. + + A superset of ``TransformTools``: extraction runs on the Element JSON of a completed + parse, so the parse tools are required as well as the two extraction tools. + """ + + request_file_upload_url: Function + start_transform_job: Function + suggest_extraction_schema_for_file: Function + start_extraction_job: Function + check_job_status: Function + get_job_results: Function + + +# The helpers below drive whichever job kind they are given, so they accept either bundle. +AnyTools = TransformTools | ExtractionTools + +_ToolsT = typing.TypeVar("_ToolsT", TransformTools, ExtractionTools) + + +def resolve_tools(group_functions: dict[str, Function], + tools_type: type[_ToolsT] = TransformTools) -> _ToolsT: # type: ignore[assignment] """Look up the Transform MCP tools in a function group by their MCP tool names. Function groups expose members as ``__``; the tool name is @@ -103,6 +132,7 @@ def resolve_tools(group_functions: dict[str, Function]) -> TransformTools: Args: group_functions: Accessible functions of the mcp_client function group. + tools_type: The bundle to resolve, naming the tools that are required. Returns: The resolved Transform MCP tools. @@ -112,12 +142,12 @@ def resolve_tools(group_functions: dict[str, Function]) -> TransformTools: """ by_tool_name = {FunctionGroup.decompose(full_name)[1]: fn for full_name, fn in group_functions.items()} - missing = [name for name in TransformTools._fields if name not in by_tool_name] + missing = [name for name in tools_type._fields if name not in by_tool_name] if missing: raise ValueError(f"Required MCP tools {missing} were not found in the function group. " f"Available tools: {sorted(by_tool_name)}") - return TransformTools(**{name: by_tool_name[name] for name in TransformTools._fields}) + return tools_type(**{name: by_tool_name[name] for name in tools_type._fields}) class MCPToolError(RuntimeError): @@ -165,8 +195,7 @@ def _check_http_response(response: httpx.Response, description: str) -> None: raise RuntimeError(f"The {description} failed with HTTP {response.status_code} from {safe_url}") -async def _upload_source(tools: TransformTools, http_client: httpx.AsyncClient, source: str, - max_file_size_bytes: int) -> str: +async def _upload_source(tools: AnyTools, http_client: httpx.AsyncClient, source: str, max_file_size_bytes: int) -> str: """Return a file reference for the document, uploading local files first. Public HTTP or HTTPS URLs are passed through unchanged because the start_transform_job @@ -205,16 +234,17 @@ async def _upload_source(tools: TransformTools, http_client: httpx.AsyncClient, return _require(upload, "file_ref", "request_file_upload_url") -def _timeout_error(job_id: str, timeout_seconds: float) -> TimeoutError: - return TimeoutError(f"Transform job {job_id} did not complete within {timeout_seconds} seconds. " - "Large documents can take several minutes; consider raising transform_timeout_seconds.") +def _timeout_error(job_id: str, timeout_seconds: float, job_kind: str = "Transform") -> TimeoutError: + return TimeoutError(f"{job_kind} job {job_id} did not complete within {timeout_seconds} seconds. " + "Large documents can take several minutes; consider raising the job timeout.") -async def _wait_for_job(tools: TransformTools, +async def _wait_for_job(tools: AnyTools, job_id: str, poll_interval_seconds: float, deadline: float, - timeout_seconds: float) -> None: + timeout_seconds: float, + job_kind: str = "Transform") -> None: """Poll the transform job status until it completes, fails, or times out. Transport-level blips (which the MCP client surfaces as plain-text tool output rather @@ -233,7 +263,8 @@ async def _wait_for_job(tools: TransformTools, consecutive_failures += 1 if consecutive_failures >= _MAX_CONSECUTIVE_POLL_FAILURES or loop.time() >= deadline: raise - logger.warning("Status check for transform job %s failed (%d consecutive); retrying", + logger.warning("Status check for %s job %s failed (%d consecutive); retrying", + job_kind.lower(), job_id, consecutive_failures) await asyncio.sleep(poll_interval_seconds) @@ -245,20 +276,21 @@ async def _wait_for_job(tools: TransformTools, if status == _COMPLETED_JOB_STATE: return if status not in _PENDING_JOB_STATES: - raise RuntimeError(f"Transform job {job_id} ended in unexpected state '{status}': " + raise RuntimeError(f"{job_kind} job {job_id} ended in unexpected state '{status}': " f"{str(status_payload)[:300]!r}") if loop.time() >= deadline: - raise _timeout_error(job_id, timeout_seconds) + raise _timeout_error(job_id, timeout_seconds, job_kind) - logger.debug("Transform job %s is %s; polling again in %.1f seconds", job_id, status, poll_interval_seconds) + logger.debug("%s job %s is %s; polling again in %.1f seconds", job_kind, job_id, status, poll_interval_seconds) await asyncio.sleep(poll_interval_seconds) -async def _fetch_results(tools: TransformTools, +async def _fetch_results(tools: AnyTools, job_id: str, poll_interval_seconds: float, deadline: float, - timeout_seconds: float) -> dict[str, typing.Any]: + timeout_seconds: float, + job_kind: str = "Transform") -> dict[str, typing.Any]: """Fetch the transform results, tolerating the brief window after the job reports completion. The status endpoint can report ``COMPLETED`` slightly before the results are @@ -279,17 +311,19 @@ async def _fetch_results(tools: TransformTools, # successful (isError=false) tool result; this is the observed behavior of the # brief status-vs-results consistency window and is safe to retry. if loop.time() >= deadline: - raise _timeout_error(job_id, timeout_seconds) from e + raise _timeout_error(job_id, timeout_seconds, job_kind) from e consecutive_failures = 0 except RuntimeError: consecutive_failures += 1 if consecutive_failures >= _MAX_CONSECUTIVE_POLL_FAILURES or loop.time() >= deadline: raise - logger.warning("Results fetch for transform job %s failed (%d consecutive); retrying", + logger.warning("Results fetch for %s job %s failed (%d consecutive); retrying", + job_kind.lower(), job_id, consecutive_failures) - logger.debug("Transform job %s results are not ready yet; polling again in %.1f seconds", + logger.debug("%s job %s results are not ready yet; polling again in %.1f seconds", + job_kind, job_id, poll_interval_seconds) await asyncio.sleep(poll_interval_seconds) @@ -384,3 +418,242 @@ async def _transform(source: str) -> str: "a local file or a public HTTP or HTTPS URL of the document. Returns the extracted content as " "Markdown. Transformation is asynchronous on the server and can take from a few seconds up to " "several minutes depending on document size.")) + + +class ExtractStructuredDataConfig(FunctionBaseConfig, name="extract_structured_data"): + """Configuration for the ``extract_structured_data`` function.""" + + mcp_group: FunctionGroupRef = Field( + default=FunctionGroupRef("unstructured_transform"), + description="Reference to the mcp_client function group connected to the Unstructured Transform MCP server.") + poll_interval_seconds: float = Field(default=5.0, gt=0, description="Delay between job status checks.") + parse_timeout_seconds: float = Field( + default=900.0, + gt=0, + description="Maximum time to wait for the parse that precedes an extraction. Large documents take longer.") + extraction_timeout_seconds: float = Field(default=600.0, + gt=0, + description="Maximum time to wait for the extraction job itself.") + http_timeout_seconds: float = Field(default=120.0, + gt=0, + description="Timeout for the plain HTTP file upload request.") + max_file_size_bytes: int = Field( + default=MAX_FILE_SIZE_BYTES, + gt=0, + description="Maximum document size accepted by the Transform service (50 MB at the time of writing).") + max_output_characters: int = Field( + default=50_000, + gt=0, + description="Truncate the returned JSON beyond this length to protect the context window of the agent. " + "Truncated output is no longer valid JSON, so a notice saying so is appended.") + + +class ExtractionResult(BaseModel): + """Outcome of a structured data extraction over one document.""" + + records_json: str + filename: str + schema_was_suggested: bool + + +# Input types whose parse quality benefits from a vision model reading the page. The +# server's guidance is to start with the vlm strategy for these, and to use fast for +# everything else, where vlm and hi_res silently fall back to fast anyway. Extraction can +# only surface what the parse captured, so this choice sets the quality ceiling. +_VLM_PARTITION_SUFFIXES = frozenset({ + ".bmp", + ".heic", + ".jpeg", + ".jpg", + ".pdf", + ".png", + ".ppt", + ".pptx", + ".tif", + ".tiff", + ".webp", +}) + + +def _partition_strategy_for(source: str) -> str: + """Return the parse strategy to use for a document that will be extracted from.""" + path = urlsplit(source).path if source.startswith(("http://", "https://")) else source + return "vlm" if Path(path).suffix.lower() in _VLM_PARTITION_SUFFIXES else "fast" + + +async def _parse_to_output_ref(tools: ExtractionTools, + config: ExtractStructuredDataConfig, + http_client: httpx.AsyncClient, + source: str) -> tuple[str, str]: + """Parse the document and return the Element JSON reference the extractor consumes. + + Returns: + The completed parse job's ``output_ref`` and the filename the server reported. + """ + file_ref = await _upload_source(tools, http_client, source, config.max_file_size_bytes) + + stages = {"partition": {"strategy": _partition_strategy_for(source)}} + job = await _invoke_tool(tools.start_transform_job, file_refs=[file_ref], stages=stages) + job_id = _require(job, "job_id", "start_transform_job") + logger.info("Started parse job %s for '%s' ahead of extraction", job_id, source) + + deadline = asyncio.get_running_loop().time() + config.parse_timeout_seconds + await _wait_for_job(tools, job_id, config.poll_interval_seconds, deadline, config.parse_timeout_seconds, "Parse") + results = await _fetch_results(tools, + job_id, + config.poll_interval_seconds, + deadline, + config.parse_timeout_seconds, + "Parse") + + files = results.get("files") + if not isinstance(files, list) or not files: + raise RuntimeError(f"Parse job {job_id} completed but returned no files.") + first_file = files[0] + if not isinstance(first_file, dict): + raise RuntimeError(f"Parse job {job_id} returned an unexpected files entry: {str(first_file)[:300]!r}") + + # The rendered output is deliberately not downloaded. The extractor consumes the + # durable Element JSON handle, which is present whichever output format was rendered, + # so fetching the Markdown here would transfer a result nobody reads. + output_ref = _require(first_file, "output_ref", "get_job_results") + return str(output_ref), str(first_file.get("filename") or Path(source).name) + + +async def _suggest_schema(tools: ExtractionTools, output_ref: str, guidance: str | None) -> str: + """Ask the server to draft a JSON Schema for a parsed document.""" + args: dict[str, typing.Any] = {"element_json_ref": output_ref} + if guidance: + args["guidance"] = guidance + + payload = await _invoke_tool(tools.suggest_extraction_schema_for_file, **args) + schema = _require(payload, "schema", "suggest_extraction_schema_for_file") + # The tool returns the schema as a JSON string; tolerate an object for robustness. + return schema if isinstance(schema, str) else json.dumps(schema) + + +async def extract_source(tools: ExtractionTools, + config: ExtractStructuredDataConfig, + source: str, + extraction_schema: str | None = None, + guidance: str | None = None) -> ExtractionResult: + """Parse one document, then extract structured data from it against a JSON Schema. + + Args: + tools: The resolved Transform MCP tools. + config: The function configuration. + source: Local file path or public HTTP or HTTPS URL of the document. + extraction_schema: JSON Schema as a JSON string. When omitted, the server drafts one. + guidance: Free-text guidance on which fields matter and how to fill them. + + Returns: + The extracted records as JSON, with provenance, plus which schema was used. + + Raises: + ValueError: If a supplied schema is not valid JSON. + """ + if extraction_schema is not None: + try: + json.loads(extraction_schema) + except json.JSONDecodeError as e: + raise ValueError(f"The supplied extraction schema is not valid JSON: {e}") from e + + async with httpx.AsyncClient(timeout=config.http_timeout_seconds) as http_client: + output_ref, filename = await _parse_to_output_ref(tools, config, http_client, source) + + schema_was_suggested = extraction_schema is None + if extraction_schema is None: + extraction_schema = await _suggest_schema(tools, output_ref, guidance) + logger.info("The server drafted an extraction schema for '%s'", filename) + + extraction_args: dict[str, typing.Any] = { + "element_json_refs": [output_ref], "schema_to_extract": extraction_schema + } + if guidance: + extraction_args["extraction_guidance"] = guidance + + job = await _invoke_tool(tools.start_extraction_job, **extraction_args) + job_id = _require(job, "job_id", "start_extraction_job") + logger.info("Started extraction job %s for '%s'", job_id, filename) + + deadline = asyncio.get_running_loop().time() + config.extraction_timeout_seconds + await _wait_for_job(tools, + job_id, + config.poll_interval_seconds, + deadline, + config.extraction_timeout_seconds, + "Extraction") + results = await _fetch_results(tools, + job_id, + config.poll_interval_seconds, + deadline, + config.extraction_timeout_seconds, + "Extraction") + + files = results.get("files") + if not isinstance(files, list) or not files: + raise RuntimeError(f"Extraction job {job_id} completed but returned no records.") + + # Each entry wraps extracted_data with the provenance that ties it back to its source + # document: filename, filetype, processed_date_utc, and source_file_uri. The whole + # wrapper is returned rather than the bare data so that audit trail is not lost. + return ExtractionResult(records_json=json.dumps(files, indent=2), + filename=filename, + schema_was_suggested=schema_was_suggested) + + +@register_function(config_type=ExtractStructuredDataConfig) +async def extract_structured_data(config: ExtractStructuredDataConfig, builder: Builder): + """Register the ``extract_structured_data`` function. + + Resolves the parse and extraction MCP tools from the configured mcp_client function + group and exposes the chained parse-then-extract flow to the agent as one tool. + """ + group = await builder.get_function_group(config.mcp_group) + tools = resolve_tools(await group.get_accessible_functions(), ExtractionTools) + + # The agent-facing parameter is deliberately not called "schema": these names become + # fields of a generated Pydantic input model, and "schema" shadows a BaseModel attribute. + async def _extract(source: str, extraction_schema: str = "", guidance: str = "") -> str: + """Extract structured data from a document as JSON. + + Args: + source: Local file path or public HTTP or HTTPS URL of the document. + extraction_schema: Optional JSON Schema, as a JSON string, naming the fields to + extract. When empty, the server drafts a schema from the document itself. + guidance: Optional free text describing which fields matter. + + Returns: + The extracted records as JSON, or an error description. + """ + try: + result = await extract_source(tools, config, source, extraction_schema or None, guidance or None) + except (OSError, ValueError, RuntimeError, TimeoutError, httpx.HTTPError, httpx.InvalidURL) as e: + logger.exception("Failed to extract structured data from '%s'", source) + return f"ERROR: failed to extract structured data from '{source}': {e}" + + logger.info("Extracted structured data from '%s' using a schema %s", + result.filename, + "drafted by the server" if result.schema_was_suggested else "supplied by the caller") + + records_json = result.records_json + if len(records_json) > config.max_output_characters: + return (f"{records_json[:config.max_output_characters]}\n\n" + f"[Output truncated at {config.max_output_characters} characters, so the JSON above is " + f"incomplete. The full result has {len(records_json)} characters. Extract fewer fields, or " + f"raise max_output_characters.]") + return records_json + + yield FunctionInfo.from_fn( + _extract, + description=("Extract specific named fields from a document (PDF, DOCX, PPTX, XLSX, HTML, images, and 40+ " + "other formats) as structured JSON, using the Unstructured Transform service. Prefer this over " + "transform_document when the user wants named values such as an invoice's line items or a " + "form's answers, rather than the whole document as text. 'source' must be the path to a local " + "file or a public HTTP or HTTPS URL. 'extraction_schema' is optional: pass a JSON Schema as a " + "JSON string to fix the output shape, or leave it out to have the server draft one from the " + "document. " + "'guidance' is optional free text describing which fields matter. Returns one JSON record per " + "document, wrapping the extracted data with the source filename. The document is parsed first " + "and the extraction runs as a second asynchronous job, so this can take from a few seconds up " + "to several minutes.")) diff --git a/examples/unstructured_transform_mcp/tests/test_extract_structured_data.py b/examples/unstructured_transform_mcp/tests/test_extract_structured_data.py new file mode 100644 index 0000000..633abc8 --- /dev/null +++ b/examples/unstructured_transform_mcp/tests/test_extract_structured_data.py @@ -0,0 +1,512 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the extract_structured_data orchestration. No network access is required.""" + +import json +import typing +from contextlib import asynccontextmanager +from pathlib import Path + +import pytest +from pydantic import BaseModel +from pydantic import ConfigDict +from pytest_httpserver import HTTPServer + +from nat_unstructured_transform_mcp.register import ExtractionTools +from nat_unstructured_transform_mcp.register import ExtractStructuredDataConfig +from nat_unstructured_transform_mcp.register import _partition_strategy_for +from nat_unstructured_transform_mcp.register import extract_source +from nat_unstructured_transform_mcp.register import extract_structured_data +from nat_unstructured_transform_mcp.register import resolve_tools + +_GROUP = "unstructured_transform" + + +class _ToolArgs(BaseModel): + """Permissive input-schema stand-in that accepts any tool arguments.""" + model_config = ConfigDict(extra="allow") + + +class _FakeTool: + """Mimics an MCP tool exposed through a function group: JSON text in, JSON text out.""" + + def __init__(self, name: str, handler: typing.Callable[..., typing.Awaitable[str]]): + self.instance_name = f"{_GROUP}__{name}" + self.input_schema = _ToolArgs + self.calls: list[dict[str, typing.Any]] = [] + self._handler = handler + + async def ainvoke(self, value: _ToolArgs) -> str: + """Record the call arguments and return the handler's canned response.""" + args = value.model_dump() + self.calls.append(args) + return await self._handler(**args) + + +class _FakeExtractionService: + """Canned Transform MCP responses for the chained parse-then-extract flow. + + The parse job and the extraction job have distinct ids, and the shared status and + results tools dispatch on the ``job_id`` they are given, exactly as the server does. + """ + + def __init__(self, + httpserver: HTTPServer, + parse_statuses: list[str] | None = None, + extraction_statuses: list[str] | None = None): + self.file_ref = "u10d://file/test-file-ref" + self.parse_job_id = "parse-job-1" + self.extraction_job_id = "extract-job-1" + self.output_ref = "u10d://output/parse-job-1_sample.json" + self.suggested_schema = json.dumps({ + "type": "object", "properties": { + "total": { + "type": "string" + } + }, "required": ["total"] + }) + self.records = [{ + "filename": "sample.pdf", + "filetype": "application/pdf", + "processed_date_utc": "2026-07-28T10:45:04.158141Z", + "source_file_uri": self.output_ref, + "extracted_data": { + "total": "42.00" + }, + }] + + self.upload_url = httpserver.url_for("/upload/test-file-ref") + # Registered so that an unwanted fetch of the rendered parse output would succeed + # and show up in the request log rather than erroring out. + self.download_url = httpserver.url_for("/download/test-file-ref") + httpserver.expect_request("/upload/test-file-ref", method="PUT").respond_with_data("") + httpserver.expect_request("/download/test-file-ref", method="GET").respond_with_data("# Sample") + + self._parse_statuses = list(parse_statuses or ["COMPLETED"]) + self._extraction_statuses = list(extraction_statuses or ["COMPLETED"]) + + async def _request_file_upload_url(**kwargs) -> str: + return json.dumps({ + "upload_url": self.upload_url, + "method": "PUT", + "headers": { + "Content-Type": kwargs.get("content_type", "application/octet-stream") + }, + "file_ref": self.file_ref, + }) + + async def _start_transform_job(**_kwargs) -> str: + return json.dumps({"job_id": self.parse_job_id, "status": "SCHEDULED"}) + + async def _start_extraction_job(**_kwargs) -> str: + return json.dumps({"job_id": self.extraction_job_id, "status": "SCHEDULED"}) + + async def _suggest_extraction_schema_for_file(**_kwargs) -> str: + return json.dumps({ + "schema": self.suggested_schema, + "reconciliation_guidance": "Reconcile across representative documents." + }) + + def _next(statuses: list[str]) -> str: + return statuses.pop(0) if len(statuses) > 1 else statuses[0] + + async def _check_job_status(**kwargs) -> str: + job_id = kwargs.get("job_id") + statuses = self._parse_statuses if job_id == self.parse_job_id else self._extraction_statuses + return json.dumps({"job_id": job_id, "status": _next(statuses)}) + + async def _get_job_results(**kwargs) -> str: + if kwargs.get("job_id") == self.parse_job_id: + return json.dumps({ + "job_type": + "transform", + "job_id": + self.parse_job_id, + "files": [{ + "filename": "sample.pdf", + "download_url": self.download_url, + "output_ref": self.output_ref, + "element_count": 12, + }], + }) + return json.dumps({"job_type": "extraction", "job_id": self.extraction_job_id, "files": self.records}) + + self.group_functions = { + f"{_GROUP}__request_file_upload_url": + _FakeTool("request_file_upload_url", _request_file_upload_url), + f"{_GROUP}__start_transform_job": + _FakeTool("start_transform_job", _start_transform_job), + f"{_GROUP}__suggest_extraction_schema_for_file": + _FakeTool("suggest_extraction_schema_for_file", _suggest_extraction_schema_for_file), + f"{_GROUP}__start_extraction_job": + _FakeTool("start_extraction_job", _start_extraction_job), + f"{_GROUP}__check_job_status": + _FakeTool("check_job_status", _check_job_status), + f"{_GROUP}__get_job_results": + _FakeTool("get_job_results", _get_job_results), + } + + @property + def tools(self) -> ExtractionTools: + """Resolve the fake group into the ExtractionTools the code under test expects.""" + return resolve_tools(self.group_functions, ExtractionTools) # type: ignore[arg-type] + + def tool(self, name: str) -> _FakeTool: + """Return the fake tool registered under the given MCP tool name.""" + return typing.cast(_FakeTool, self.group_functions[f"{_GROUP}__{name}"]) + + +@pytest.fixture(name="fast_config") +def fast_config_fixture() -> ExtractStructuredDataConfig: + """Config with tiny timeouts so the polling paths run quickly in tests.""" + return ExtractStructuredDataConfig(poll_interval_seconds=0.01, + parse_timeout_seconds=5.0, + extraction_timeout_seconds=5.0, + http_timeout_seconds=5.0) + + +@pytest.fixture(name="sample_document") +def sample_document_fixture(tmp_path: Path) -> Path: + """Write a small local file for the upload path to read.""" + document = tmp_path / "sample.pdf" + document.write_bytes(b"%PDF-1.4 fake test document bytes") + return document + + +_SUPPLIED_SCHEMA = json.dumps({"type": "object", "properties": {"total": {"type": "string"}}}) + + +async def test_supplied_schema_drives_extraction(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """A caller-supplied schema is passed through and no schema suggestion is requested.""" + service = _FakeExtractionService(httpserver, parse_statuses=["IN_PROGRESS", "COMPLETED"]) + + result = await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + assert result.schema_was_suggested is False + assert result.filename == "sample.pdf" + assert json.loads(result.records_json) == service.records + + assert service.tool("suggest_extraction_schema_for_file").calls == [] + assert service.tool("start_extraction_job").calls == [{ + "element_json_refs": [service.output_ref], "schema_to_extract": _SUPPLIED_SCHEMA + }] + + +async def test_parse_uses_vlm_strategy_and_output_ref(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """The parse ahead of an extraction requests a high-fidelity strategy for a PDF.""" + service = _FakeExtractionService(httpserver) + + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + assert service.tool("start_transform_job").calls == [{ + "file_refs": [service.file_ref], "stages": { + "partition": { + "strategy": "vlm" + } + } + }] + + +async def test_rendered_parse_output_is_not_downloaded(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """Extraction consumes the output_ref, so the rendered parse output is never fetched.""" + service = _FakeExtractionService(httpserver) + + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + get_requests = [request for request, _ in httpserver.log if request.method == "GET"] + assert get_requests == [] + + +async def test_schema_is_suggested_when_absent(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """Without a schema, the server drafts one from the parsed document and it is used.""" + service = _FakeExtractionService(httpserver) + + result = await extract_source(service.tools, fast_config, str(sample_document)) + + assert result.schema_was_suggested is True + assert service.tool("suggest_extraction_schema_for_file").calls == [{"element_json_ref": service.output_ref}] + assert service.tool("start_extraction_job").calls == [{ + "element_json_refs": [service.output_ref], "schema_to_extract": service.suggested_schema + }] + + +async def test_guidance_is_forwarded_to_both_tools(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """Guidance steers the suggested schema and the extraction itself.""" + service = _FakeExtractionService(httpserver) + + await extract_source(service.tools, fast_config, str(sample_document), guidance="Only monetary totals.") + + assert service.tool("suggest_extraction_schema_for_file").calls == [{ + "element_json_ref": service.output_ref, "guidance": "Only monetary totals." + }] + assert service.tool("start_extraction_job").calls[0]["extraction_guidance"] == "Only monetary totals." + + +async def test_provenance_wrapper_is_preserved(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """The result keeps the provenance fields that tie each record to its document.""" + service = _FakeExtractionService(httpserver) + + result = await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + record = json.loads(result.records_json)[0] + assert set(record) == {"filename", "filetype", "processed_date_utc", "source_file_uri", "extracted_data"} + assert record["source_file_uri"] == service.output_ref + + +async def test_invalid_supplied_schema_rejected(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """A schema that is not valid JSON is rejected before any job is started.""" + service = _FakeExtractionService(httpserver) + + with pytest.raises(ValueError, match="not valid JSON"): + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema="{not json") + + assert service.tool("start_transform_job").calls == [] + assert service.tool("start_extraction_job").calls == [] + + +async def test_missing_output_ref_raises(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """A parse result without an output_ref names the missing field.""" + service = _FakeExtractionService(httpserver) + + async def _no_output_ref(**_kwargs) -> str: + return json.dumps({"job_id": service.parse_job_id, "files": [{"filename": "sample.pdf"}]}) + + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _no_output_ref) + + with pytest.raises(RuntimeError, match="'output_ref'"): + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + +async def test_failed_extraction_job_names_the_job_kind(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """A failed extraction job reports the extraction state, not the parse state.""" + service = _FakeExtractionService(httpserver, extraction_statuses=["FAILED"]) + + with pytest.raises(RuntimeError, match="Extraction job .* unexpected state 'FAILED'"): + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + +async def test_extraction_timeout_reports_extraction(httpserver: HTTPServer, sample_document: Path): + """An extraction that never completes times out against its own deadline.""" + service = _FakeExtractionService(httpserver, extraction_statuses=["IN_PROGRESS"]) + config = ExtractStructuredDataConfig(poll_interval_seconds=0.01, + parse_timeout_seconds=5.0, + extraction_timeout_seconds=0.05, + http_timeout_seconds=5.0) + + with pytest.raises(TimeoutError, match="Extraction job .* did not complete"): + await extract_source(service.tools, config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + +async def test_empty_records_raises(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """A completed extraction that returns nothing raises rather than reporting success.""" + service = _FakeExtractionService(httpserver) + real_results = service.tool("get_job_results") + + async def _empty_extraction(**kwargs) -> str: + if kwargs.get("job_id") == service.extraction_job_id: + return json.dumps({"job_type": "extraction", "job_id": service.extraction_job_id, "files": []}) + return await real_results._handler(**kwargs) + + service.group_functions[f"{_GROUP}__get_job_results"] = _FakeTool("get_job_results", _empty_extraction) + + with pytest.raises(RuntimeError, match="returned no records"): + await extract_source(service.tools, fast_config, str(sample_document), extraction_schema=_SUPPLIED_SCHEMA) + + +async def test_public_url_skips_upload(httpserver: HTTPServer, fast_config: ExtractStructuredDataConfig): + """A public URL is parsed directly, with no upload step.""" + service = _FakeExtractionService(httpserver) + url = "https://example.com/invoice.pdf" + + await extract_source(service.tools, fast_config, url, extraction_schema=_SUPPLIED_SCHEMA) + + assert service.tool("request_file_upload_url").calls == [] + assert service.tool("start_transform_job").calls[0]["file_refs"] == [url] + + +def test_resolve_tools_reports_missing_extraction_tool(httpserver: HTTPServer): + """resolve_tools names an absent extraction tool, which is the include-list contract check.""" + service = _FakeExtractionService(httpserver) + del service.group_functions[f"{_GROUP}__start_extraction_job"] + + with pytest.raises(ValueError, match="start_extraction_job"): + resolve_tools(service.group_functions, ExtractionTools) # type: ignore[arg-type] + + +def test_shipped_config_exposes_the_extraction_tools(monkeypatch: pytest.MonkeyPatch): + """The shipped config registers the function and includes the tools it resolves. + + ``resolve_tools`` fails fast on a missing tool, so the include list and + ``ExtractionTools`` have to stay in step. + """ + from nat.runtime.loader import load_config + from nat.test.utils import locate_example_config + + monkeypatch.setenv("UNSTRUCTURED_API_KEY", "test-key-0123456789") + + config = load_config(locate_example_config(ExtractStructuredDataConfig)) + + function_config = config.functions["extract_structured_data"] + assert isinstance(function_config, ExtractStructuredDataConfig) + assert str(function_config.mcp_group) == "unstructured_transform" + + included = set(config.function_groups["unstructured_transform"].include or []) + assert set(ExtractionTools._fields) <= included + + assert "extract_structured_data" in config.workflow.tool_names + + +@pytest.mark.parametrize( + "source, expected", + [ + ("/tmp/report.pdf", "vlm"), + ("/tmp/scan.PNG", "vlm"), + ("/tmp/deck.pptx", "vlm"), + ("/tmp/notes.docx", "fast"), + ("/tmp/data.xlsx", "fast"), + ("/tmp/no-suffix", "fast"), + ("https://example.com/a/report.pdf?token=abc", "vlm"), + ("https://example.com/a/report.docx?token=abc", "fast"), + ], +) +def test_partition_strategy_by_input_type(source: str, expected: str): + """Images, slides, and PDFs get a vision parse; other formats get the fast parse. + + A query string on a URL must not defeat the suffix check. + """ + assert _partition_strategy_for(source) == expected + + +# --- The registered function: agent-facing wrapper contract --- + + +class _StubGroup: + """Minimal function group that exposes the fake tools to the registered function.""" + + def __init__(self, group_functions: dict[str, _FakeTool]): + self._group_functions = group_functions + + async def get_accessible_functions(self) -> dict[str, _FakeTool]: + """Return the fake group's functions.""" + return self._group_functions + + +class _StubBuilder: + """Minimal builder that hands the registered function its function group.""" + + def __init__(self, group: _StubGroup): + self._group = group + + async def get_function_group(self, _name) -> _StubGroup: + """Return the stub group regardless of the requested name.""" + return self._group + + +@asynccontextmanager +async def _registered_extract_info(service: _FakeExtractionService, config: ExtractStructuredDataConfig): + """Drive the registered async generator exactly as the workflow builder would.""" + builder = _StubBuilder(_StubGroup(service.group_functions)) + async with extract_structured_data(config, builder) as info: # type: ignore[arg-type] + yield info + + +async def test_registered_input_schema_field_names(httpserver: HTTPServer, fast_config: ExtractStructuredDataConfig): + """The generated input model exposes the three documented arguments. + + ``extraction_schema`` rather than ``schema``: the latter shadows a pydantic + ``BaseModel`` attribute and makes the generated model emit a warning. + """ + service = _FakeExtractionService(httpserver) + + async with _registered_extract_info(service, fast_config) as info: + assert list(info.input_schema.model_fields) == ["source", "extraction_schema", "guidance"] + + +async def test_registered_function_happy_path(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """The registered tool returns the extracted records as JSON.""" + service = _FakeExtractionService(httpserver) + + async with _registered_extract_info(service, fast_config) as info: + output = await info.single_fn( + info.input_schema(source=str(sample_document), extraction_extraction_schema=_SUPPLIED_SCHEMA)) + + assert json.loads(output) == service.records + + +async def test_registered_function_without_schema(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + sample_document: Path): + """Called with only a source, the tool has the server draft the schema.""" + service = _FakeExtractionService(httpserver) + + async with _registered_extract_info(service, fast_config) as info: + output = await info.single_fn(info.input_schema(source=str(sample_document))) + + assert json.loads(output) == service.records + assert service.tool("suggest_extraction_schema_for_file").calls == [{"element_json_ref": service.output_ref}] + + +async def test_registered_function_returns_error_string(httpserver: HTTPServer, + fast_config: ExtractStructuredDataConfig, + tmp_path: Path): + """Failures must reach the agent as a readable tool result, never as an exception.""" + service = _FakeExtractionService(httpserver) + missing = tmp_path / "does-not-exist.pdf" + + async with _registered_extract_info(service, fast_config) as info: + output = await info.single_fn( + info.input_schema(source=str(missing), extraction_extraction_schema=_SUPPLIED_SCHEMA)) + + assert output.startswith("ERROR: failed to extract structured data") + assert str(missing) in output + + +async def test_registered_function_truncates_output(httpserver: HTTPServer, sample_document: Path): + """Output longer than the limit is truncated with a notice that the JSON is incomplete.""" + service = _FakeExtractionService(httpserver) + config = ExtractStructuredDataConfig(poll_interval_seconds=0.01, + parse_timeout_seconds=5.0, + extraction_timeout_seconds=5.0, + http_timeout_seconds=5.0, + max_output_characters=20) + + async with _registered_extract_info(service, config) as info: + output = await info.single_fn( + info.input_schema(source=str(sample_document), extraction_extraction_schema=_SUPPLIED_SCHEMA)) + + assert "[Output truncated at 20 characters" in output + assert "incomplete" in output diff --git a/examples/unstructured_transform_mcp/tests/test_workflow_integration.py b/examples/unstructured_transform_mcp/tests/test_workflow_integration.py index b3bb557..ec85f38 100644 --- a/examples/unstructured_transform_mcp/tests/test_workflow_integration.py +++ b/examples/unstructured_transform_mcp/tests/test_workflow_integration.py @@ -111,6 +111,57 @@ async def test_transform_document_function_live(sample_pdf: Path): assert "xylophone" in result.markdown.lower() +@pytest.mark.slow +@pytest.mark.integration +@pytest.mark.timeout(900) +@pytest.mark.usefixtures("unstructured_api_key") +async def test_extract_structured_data_function_live(sample_pdf: Path): + """Exercise the chained parse-then-extract flow against the live server, without an LLM. + + Uses an explicit schema so the assertion does not depend on what the server would + draft; the no-schema path is covered by the unit tests. + """ + import json + + from nat.builder.workflow_builder import WorkflowBuilder + from nat.runtime.loader import load_config + from nat.test.utils import locate_example_config + from nat_unstructured_transform_mcp.register import ExtractionTools + from nat_unstructured_transform_mcp.register import ExtractStructuredDataConfig + from nat_unstructured_transform_mcp.register import extract_source + from nat_unstructured_transform_mcp.register import resolve_tools + + shipped_config = load_config(locate_example_config(ExtractStructuredDataConfig)) + schema = json.dumps({ + "type": "object", + "properties": { + "magic_word": { + "type": "string", "description": "The magic word stated in the document" + } + }, + "required": ["magic_word"], + "additionalProperties": False, + }) + + async with WorkflowBuilder() as builder: + await builder.add_auth_provider("unstructured_auth", shipped_config.authentication["unstructured_auth"]) + group = await builder.add_function_group("unstructured_transform", + shipped_config.function_groups["unstructured_transform"]) + + tools = resolve_tools(await group.get_accessible_functions(), ExtractionTools) + config = ExtractStructuredDataConfig(poll_interval_seconds=3.0, + parse_timeout_seconds=600.0, + extraction_timeout_seconds=300.0) + + result = await extract_source(tools, config, str(sample_pdf), extraction_schema=schema) + + records = json.loads(result.records_json) + assert records, "Expected at least one extracted record" + # The provenance wrapper is what ties each record back to the document it came from. + assert "source_file_uri" in records[0] + assert "xylophone" in json.dumps(records[0]["extracted_data"]).lower() + + @pytest.mark.slow @pytest.mark.integration @pytest.mark.timeout(1000) From 8b1777c00d68453960dd74e8ef76e951a76dbbfa Mon Sep 17 00:00:00 2001 From: Simon Coombes Date: Tue, 28 Jul 2026 09:25:40 -0400 Subject: [PATCH 3/3] Assert the extraction chain rather than the extracted value in the live test Signed-off-by: Simon Coombes --- .../tests/test_workflow_integration.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/examples/unstructured_transform_mcp/tests/test_workflow_integration.py b/examples/unstructured_transform_mcp/tests/test_workflow_integration.py index ec85f38..a654649 100644 --- a/examples/unstructured_transform_mcp/tests/test_workflow_integration.py +++ b/examples/unstructured_transform_mcp/tests/test_workflow_integration.py @@ -120,6 +120,14 @@ async def test_extract_structured_data_function_live(sample_pdf: Path): Uses an explicit schema so the assertion does not depend on what the server would draft; the no-schema path is covered by the unit tests. + + This asserts the shape of the result rather than the extracted value. The fixture is a + synthetic one-sentence PDF, and while the parse reads it correctly, what the extractor + fills into the fields for so degenerate a document is not stable enough to assert on + (observed against production: a plausible-but-wrong value for one schema, and empty + strings for another). The contract this test protects is the chain itself: that the + parse hands its ``output_ref`` to the extractor, that the supplied schema shapes the + output, and that the provenance wrapper survives. """ import json @@ -157,9 +165,15 @@ async def test_extract_structured_data_function_live(sample_pdf: Path): records = json.loads(result.records_json) assert records, "Expected at least one extracted record" + assert result.schema_was_suggested is False + # The provenance wrapper is what ties each record back to the document it came from. - assert "source_file_uri" in records[0] - assert "xylophone" in json.dumps(records[0]["extracted_data"]).lower() + record = records[0] + assert {"filename", "filetype", "processed_date_utc", "source_file_uri"} <= set(record) + assert record["source_file_uri"].startswith("u10d://output/"), "Expected the parse Element JSON ref" + + # The supplied schema shaped the output: its required field is present. + assert set(record["extracted_data"]) == {"magic_word"} @pytest.mark.slow