From spot fixes to measured gains: improving document extraction accuracy with ai_extract and MLflow Prompt Registry

From spot fixes to measured gains: improving document extraction accuracy with ai_extract and MLflow Prompt Registry | Databricks Blog Skip to main content

Summary

You don't have to do it all yourself: Genie Code for document processing automates the core of this workflow, scoring extraction field by field against your labeled data, proposing schema and instruction fixes as reviewable diffs, and re-running the eval to show before-and-after results.

Improve ai_extract accuracy with the right lever: Use schema descriptions for field-level rules, prompts for document-level logic, precision mode for structural complexity, and citations to diagnose errors.

Version prompts and schemas with MLflow Prompt Registry: Track extraction logic as versioned artifacts and deploy or roll back changes through aliases instead of editing pipeline code.

Evaluate changes before promotion: Compare candidates against fixed, human-reviewed ground truth while tracking accuracy, coverage, and run-to-run variation.

Document extraction often starts with a simple loop: spot an incorrect field, adjust the prompt or schema, rerun the document, and check the result. That approach becomes unreliable as fields and document formats multiply. A change may fix one field while introducing regressions elsewhere, work for one layout but fail on another, or appear better simply due to normal variation across LLM runs. We saw this on a recent project that extracted roughly 40 structured fields from partner-submitted order documents across dozens of PDF layouts. Reviewers corrected and approved extracted values before they moved downstream, giving us human-reviewed ground truth for evaluation. The challenge was not getting ai_extract to return the fields. It was improving extraction without degrading the system overall. For example, the revision field sometimes returned status terms such as ORIGINAL instead of a revision number or blank. A small schema change cut errors for that field roughly in half. But across every field, total extraction errors increased by 9%. The lesson was clear: improving one field does not necessarily improve the extraction system. The pipeline already had most of the pieces needed to evaluate changes systematically. Documents landed in object storage and flowed through a Lakeflow SDP pipeline using Auto Loader. Bronze parsed documents with ai_parse_document, Silver extracted structured fields with ai_extract, and Gold published results as a materialized view. Predictions synced to Lakebase, where a Databricks App let reviewers correct and approve fields, and Lakebase CDC returned those approvals to the lakehouse.

Figure 1. A document-extraction pipeline with human review. Prompts and schemas are hardcoded in Silver, while reviewer corrections are not yet used to evaluate extraction changes.

Two pieces were missing from the feedback loop. Reviewer approvals were available, but they were not assembled into a reusable evaluation dataset. At the same time, prompts and schemas were embedded directly in pipeline code, so even a small change to extraction required a code deployment. We addressed both problems by building a workflow around three questions: What should we change? How should we version it? And what evidence is enough to promote it? This post covers three parts of that workflow: Improve ai_extract with the right lever: Use schemas for field rules, global instructions for document logic, precision mode for difficult structures, and citations to diagnose errors. Version prompts and schemas: Store both in MLflow Prompt Registry and use aliases to deploy or roll back configurations without editing pipeline code. Evaluate before promotion : Compare candidates on fixed, human-reviewed ground truth, tracking accuracy, coverage, and run-to-run variation.

Together, these steps create a repeatable path from a proposed extraction change to production, with evidence that the change improves the system as a whole. Start with Genie Code The rest of this post walks through the practices we recommend and why each one matters. You don't have to implement all of them on your own. Genie Code for document processing builds the Lakeflow job (ai_parse_document, ai_classify, and ai_extract with citations and confidence scores) and automates the core of this workflow with an evaluation and hill-climbing loop : Measure : Scores the current extraction field by field against your labeled table, so you can see which fields are weak. Improve : Proposes a fix, most often a sharper field description, as a reviewable diff that you approve before it applies. Re-evaluate : Re-runs the eval set after every accepted change and shows before-and-after results, so regressions are as visible as improvements.

For example, you can ask Genie Code: "Evaluate the extraction against my labeled table and tell me which fields are weakest," followed by "Accuracy on the revision field could be higher. Iteratively improve the schema and instructions, re-run the evaluation, and show me the before and after." Understanding the fundamentals still matters. It helps you judge Genie Code's proposed changes, decide when a result is good enough to promote, and extend the workflow for production, for example by versioning prompts and schemas in MLflow Prompt Registry. Each section below explains a recommended practice and notes how Genie Code can handle it for you. Choose the right ai_extract lever ai_extract takes three inputs: the parsed document, a JSON schema, and optional global instructions. The schema and instructions serve different purposes: Schema: Defines output fields and field-specific extraction rules. Global instructions: Provide guidance that applies across the document.

The result is a structured value with one entry for each schema field. You can access individual fields directly, for example: extracted:response.revision.value . Put field-specific rules in the schema. Start with the schema structure. When a field has a closed set of valid values, use an enum constraint rather than listing the allowed values in the description. Use numeric types for numbers and integers, and use nested arrays or objects for repeated structures. Then use schema descriptions for rules that types alone cannot express, such as: Preserve leading zeros in identifiers such as 004 Exclude status terms such as ORIGINAL , NEW , or REVISED Stop an identifier at the first invalid character Return null when a value isn't present...