LLM Toolkit: Validation is all you need

by Jeff Schomay
May 20, 2024

Forget chains—structured output and validation are all you need.

At Mechanical Orchard, we’re building complex, bespoke, explainable AI agents to interrogate legacy mainframe systems autonomously. We’ve used a wide range of popular AI libraries and frameworks to bring LLM best practices into our codebase. The one I’ve been absolutely loving is Instructor, because modeling data is so much more powerful than modeling prompts. In this technical article, I’ll show you why.

Use case

We have a RAG tool that takes a plain English question, rephrases it for better results based on a known schema, converts it to a graph database query, runs the query, performs multiple retries on errors using self-healing from error messages, interprets the results, and returns a conversational style answer with citations for transparency and groundedness.

Earlier, I wrote about our first approach that used a hard-coded 9-step chain to do this reliably. With Instructor, we can accomplish equivalent (and sometimes better) results in only 20 lines of code!

Implementation

Instructor works by making LLMs format their output in a way that can be run through Pydantic, which is Python’s premier data validation library. This simple trick unlocks many power features:

  • Structured output with well-defined types
  • Validation at the field or model level
  • Easy to apply, common LLM enhancement patterns

Let me show you how:

from pydantic import BaseModel, Field, model_validator
import instructor

class GraphDatabaseResult(BaseModel):
    rephrased_question: str = Field(..., description="Rephrase the question in terms of the graph schema.")
    reasoning: str = Field(..., description="Think step by step on how to make a graph database query that can answer this question.")
    query: str = Field(..., description="The query to run")
    result: str = Field(description="leave this blank", default="Pending...")

@model_validator(mode='after')
    def try_to_run_query(self) -> 'Self':
        self.result = call_graph_db(self.query)
        return self

def ask_the_graph(user_question: str) -> GraphDatabaseResult:
    client = instructor.from_openai(OpenAI())
    return client.chat.completions.create(
        model="gpt-4-turbo",
        messages=[{"role": "system", "content": SCHEMA}, {"role": "user", "content": user_question}],
        response_model=GraphDatabaseResult,
        max_retries=3
    )

This short chunk of code is doing a lot! Let’s walk through it:

  • First, notice the lack of an obvious prompt. Instead, we provide a declarative annotated class that inherits from Pydantic. Instructor tells the LLM to make this GraphDatabaseResult class. Its shape, types, and field descriptions provide the context for the LLM.

  • Now look at its fields. Both rephrased_question and reasoning are a form of “Chain of Thought” (CoT) to prime the LLM output with sound logic before it gets to generating actual query tokens. Adding CoT couldn’t be easier!

  • The model_validator runs automatically after the LLM returns the raw response that Instructor feeds into GraphDatabaseResult. This is where the magic happens. All it does is actually run the query. call_graph_db will raise a database exception if there is a problem with the query, which will make the validator fail, which will make Instructor send the error to the LLM to try again. If all goes well, it stores the result in the result field and finishes. That’s powerful!

  • Finally, the LLM call is very much like a standard OpenAI chat completion, with just a few tweaks from Instructor such as the response_model and max_retries fields. Also notice how minimal the messages are, just the graph schema so the LLM knows what is in the database, and the user’s question.

You might feel a little odd about hitting the database as part of validation, but let me ask, how else do you know if the query is valid without actually running it?

And that’s it. Here’s what output looks like:

ask_the_graph("How many tables get used by jobs starting with IND?")

Output:

>GraphDatabaseResult(
    rephrased_question="Count the number of unique database tables that are read or written by jobs whose names start with 'IND'.",
    reasoning="To find the number of unique database tables that are accessed by jobs with names starting with 'IND', the query will match jobs with names starting with 'IND' and find connections to database tables either through direct reads or writes. After identifying such connections, the query will count the distinct database tables connected to these jobs.",
    query="MATCH (j:Job)-[r:READS|WRITES]->(db:DBTable)\nWHERE j.name STARTS WITH 'IND'\nRETURN COUNT(DISTINCT db) AS distinct_tables",
    result='+-------------------+\n| distinct_tables    |\n|-------------------+|\n| 89                |\n+-------------------+'\n)

Do you still need chains?

What if you want an LLM to summarize the results for us? How would you chain that together?

That’s a trick question, you don’t need chains, you just need normal functions. Notice how ask_the_graph returns a GraphDatabaseResult which you can just pass into a function that takes that type, like you normally do when working with Python code. This pattern lets you mix AI and traditional code with ease.

from fructose import Fructose

ai = Fructose()

@ai
def interpret_results(results: GraphDatabaseResult) -> str:
    """Summarize the result to answer the rephrased question in 1-2 sentences."""

interpret_results(ask_the_graph("How many tables get used by jobs starting with IND?"))

Output:

>89 different database tables are accessed by jobs whose names begin with 'IND'.

This uses Fructose, a slick, tiny library that hides away the LLM call for you if you include @ai on a function, based on the function name, arguments, types, and docstring. It even coerces the output into a structured type like Instructor.

And there you have it. With just a little more work, you get live self-evaluation to make sure answers are grounded, reliable and relevant, with automatic retries and self-healing if the score falls below a threshold.

Conclusion

Rather than letting the LLM “ramble on” in prose, you can now get high-quality, well-typed responses thanks to Instructor and Pydantic. And rather than building complex chains, you can get solid LLM performance improvement patterns with some simple validation hacks. Not only will your code be more compact and powerful, your outputs will be of higher quality too.