Control LLM output with LangChain's structured and Pydantic output parsers

In my previous Control LLM output with response type and schema post, I talked about how you can define a JSON response schema and Vertex AI makes sure the output of the Large Language Model (LLM) conforms to that schema.

In this post, I show how you can implement a similar response schema using LangChain’s structured output parser with any model. You can further get the output parsed and populated into Python classes automatically with the Pydantic output parser. This helps you to really narrow down and structure LLM outputs.

The full sample code is in main.py.

Setup

Before we start, make sure your gcloud is set up with your Google Cloud project:

gcloud config set core/project your-google-cloud-project-id

You are logged in:

gcloud auth application-default login

Create and activate a virtual environment:

python -m venv .venv
source .venv/bin/activate

Install dependencies:

pip install -r requirements.txt

Without any schema

First, let’s ask the LLM to give us some cookie recipes without a schema:

model = VertexAI(model_name=MODEL_NAME)
prompt = "List 10 popular cookie recipes"
response = model.invoke(prompt)

Run it:

python main.py without_controlled_generation

You get a response in free text:

Prompt: List 10 popular cookie recipes
Response: 1. **Chocolate Chip Cookies:** The classic, arguably the most popular cookie of all time.  Numerous variations exist, but the basic recipe is universally loved.

2. **Oatmeal Raisin Cookies:** A chewy, comforting classic often featuring oats, raisins, brown sugar, and spices.

3. **Peanut Butter Cookies:** Simple, satisfying, and often made with peanut butter, sugar, and an egg.  Variations include adding chocolate chips or pressing a fork into the top.
...

Not ideal for applications that need to parse the output.

With structured output parser

You can use the StructuredOutputParser of LangChain and define a ResponseSchema for the output:

model = VertexAI(model_name=MODEL_NAME)
prompt = "List a popular cookie recipe"

response_schemas = [
    ResponseSchema(name="recipe_name", description="Name of the recipe", type="str"),
    ResponseSchema(name="calories", description="Calories of the recipe", type="int"),
]

Get format instructions out of the schema:

output_parser = StructuredOutputParser.from_response_schemas(response_schemas)
format_instructions = output_parser.get_format_instructions()
print(f"Format instructions: {format_instructions}")

Pass the format instructions to the model with a prompt template:

prompt_template = PromptTemplate(
    template="Answer the user query.\n{format_instructions}\n{query}",
    input_variables=["query"],
    partial_variables={"format_instructions": format_instructions},
)

chain = prompt_template | model | output_parser
response = chain.invoke({"query": prompt})

Run it:

python main.py with_response_schema

First, you’ll see the format instructions as part of the prompt:

Format instructions: The output should be a markdown code snippet formatted in the following schema, including the leading and trailing "```json" and "```":

```json
{
        "recipe_name": str  // Name of the recipe
        "calories": int  // Calories of the recipe
}

You should also get JSON format back in the response:

Prompt: List a popular cookie recipe
Response: {'recipe_name': 'Chocolate Chip Cookies', 'calories': 78}

This is better, but it gets difficult when you need to define a more complex schema with lists and nested objects. You also need to convert JSON to Python classes yourself. We can do better.

With pydantic parser

Pydantic is the most widely used data validation library for Python. You can use Pydantic to define a model and use LangChain’s Pydantic parser to make sure LLM outputs conform to that model and also get the output parsed into Python classes.

First, define a Recipe and a list of Recipes classes as our model:

class Recipe(BaseModel):
    recipe_name: str = Field(description="Name of the recipe")
    calories: int = Field(description="Calories of the recipe")

class Recipes(BaseModel):
    recipes: List[Recipe]

Define a Pydantic output parser with the Recipes class and get the format instructions:

model = VertexAI(model_name=MODEL_NAME)
prompt = "List 10 popular cookie recipes"

output_parser = PydanticOutputParser(pydantic_object=Recipes)
format_instructions = output_parser.get_format_instructions()

Pass the format instructions to the model with a prompt template:

prompt_template = PromptTemplate(
    template="Answer the user query.\n{format_instructions}\n{query}\n",
    input_variables=["query"],
    partial_variables={"format_instructions": format_instructions},
)

chain = prompt_template | model | output_parser
response = chain.invoke({"query": prompt})

Run it:

python main.py with_pydantic

You should see the detailed format instructions printed:

Format instructions: The output should be formatted as a JSON instance that conforms to the JSON schema below.

As an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}
the object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.

Here is the output schema:

{"$defs": {"Recipe": {"properties": {"recipe_name": {"description": "Name of the recipe", "title": "Recipe Name", "type": "string"}, "calories": {"description": "Calories of the recipe", "title": "Calories", "type": "integer"}}, "required": ["recipe_name", "calories"], "title": "Recipe", "type": "object"}}, "properties": {"recipes": {"items": {"$ref": "#/$defs/Recipe"}, "title": "Recipes", "type": "array"}}, "required": ["recipes"]}

The response will include Recipe objects in a recipes list, automatically parsed and populated:

Prompt: List 10 popular cookie recipes
Response: recipes=[Recipe(recipe_name='Chocolate Chip Cookies', calories=78), Recipe(recipe_name='Oatmeal Raisin Cookies', calories=110), ...]

Conclusion

It can be challenging to parse LLM outputs if they’re just free-text. Thankfully, both LLM providers and frameworks like LangChain provide ways to structure the output.

Here are some references for further reading:

GenAI  GoogleAI  VertexAI  Gemini  Google Cloud Platform 

See also