Design Validation and Retry for a Multi-Step LLM Extraction Pipeline
Keep a chain of language-model extraction steps correct when one step's JSON output does not match what the next step expects, while wasting few tokens and finishing inside a strict deadline.
Problem statement
For this exercise, assume a document extraction service turns a document into a structured record by running a chain of language-model steps. Each step returns JSON, and the next step reads that JSON as its input.
Now and then a step's output is slightly off from the schema the next step expects, and the next step then fails. Assume the usual problems are a missing or renamed field, a value of the wrong type, extra text around the JSON, or output that stops before the end.
Design how the pipeline catches output that is off and recovers from it, so that few tokens are wasted, latency stays low, and every request finishes inside its strict timeout.
Assume the pipeline has three model steps: step 1 reads the document text and extracts raw fields, step 2 turns them into the target record, and step 3 checks the record across fields and fills in derived fields.
Between the steps, ordinary code transforms the JSON, for example parsing dates and amounts or looking up reference codes in another service; this code always gives the same output for the same input. Assume every step's input and output has a versioned JSON schema.
Assume callers wait for the answer, and each request has a hard deadline of 10 seconds end to end: after that the caller gives up and any result is wasted. A wrong record is worse than a clear failure: the caller must get either a record that matches the final schema or a failure that says what went wrong.
Model training, text recognition from images, and the document upload flow are out of scope.
What to build
- Accept an extraction request for a document and return, before its deadline, either a record that matches the final schema or a clear failure that says which step failed and why.
- Make sure the next step and the caller only ever receive output that matches the schema they expect, and keep the reason whenever output is rejected.
- Recover from malformed or schema-breaking step output without repeating work that already succeeded, staying within the request's time and token budget.
- Record each step attempt with its errors, tokens and latency, so operators can see failure and retry rates and what retries cost.
At scale
- For this exercise, assume about 3 million documents a day, with peaks of about 100 requests a second.
- For this exercise, assume a model call takes about 1.5 seconds at the median and up to 3 seconds at the 99th percentile, and the code between steps adds a few hundred milliseconds in total.
- For this exercise, assume step 1 reads about 8,000 tokens of document text, steps 2 and 3 each read about 1,500 tokens, and each step writes about 1,000 tokens of JSON.
- For this exercise, assume 2 to 5 percent of model step outputs fail the schema check, and most failures are small: one field missing, renamed or of the wrong type, text around the JSON, or a cut-off ending.