ContextQMD
Back to Llama Index

OpenAI Pydantic Program

docs/examples/output_parsing/openai_pydantic_program.ipynb

0.14.216.4 KB
Original Source

<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/output_parsing/openai_pydantic_program.ipynb" target="_parent"></a>

OpenAI Pydantic Program

This guide shows you how to generate structured data with new OpenAI API via LlamaIndex. The user just needs to specify a Pydantic object.

We demonstrate two settings:

  • Extraction into an Album object (which can contain a list of Song objects)
  • Extraction into a DirectoryTree object (which can contain recursive Node objects)

Extraction into Album

This is a simple example of parsing an output into an Album schema, which can contain multiple songs.

If you're opening this Notebook on colab, you will probably need to install LlamaIndex šŸ¦™.

python
%pip install llama-index-llms-openai
%pip install llama-index-program-openai
python
%pip install llama-index
python
from pydantic import BaseModel
from typing import List

from llama_index.program.openai import OpenAIPydanticProgram

Without docstring in Model

Define output schema (without docstring)

python
class Song(BaseModel):
    title: str
    length_seconds: int


class Album(BaseModel):
    name: str
    artist: str
    songs: List[Song]

Define openai pydantic program

python
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album, prompt_template_str=prompt_template_str, verbose=True
)

Run program to get structured output.

python
output = program(
    movie_name="The Shining", description="Data model for an album."
)

With docstring in Model

python
class Song(BaseModel):
    """Data model for a song."""

    title: str
    length_seconds: int


class Album(BaseModel):
    """Data model for an album."""

    name: str
    artist: str
    songs: List[Song]
python
prompt_template_str = """\
Generate an example album, with an artist and a list of songs. \
Using the movie {movie_name} as inspiration.\
"""
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album, prompt_template_str=prompt_template_str, verbose=True
)

Run program to get structured output.

python
output = program(movie_name="The Shining")

The output is a valid Pydantic object that we can then use to call functions/APIs.

python
output

Stream partial intermediate Pydantic Objects

Instead of waiting for the Function Call to generate the entire JSON, we can use the stream_partial_objects() method of the program to stream valid intermediate instances of the Pydantic Output class as soon as they're available šŸ”„

First let's define the Output Pydantic class

python
from pydantic import BaseModel, Field


class CharacterInfo(BaseModel):
    """Information about a character."""

    character_name: str
    name: str = Field(..., description="Name of the actor/actress")
    hometown: str


class Characters(BaseModel):
    """List of characters."""

    characters: list[CharacterInfo] = Field(default_factory=list)

Now we'll initialilze the program with prompt template

python
from llama_index.program.openai import OpenAIPydanticProgram

prompt_template_str = "Information about 3 characters from the movie: {movie}"

program = OpenAIPydanticProgram.from_defaults(
    output_cls=Characters, prompt_template_str=prompt_template_str
)

Finally we stream the partial objects using the stream_partial_objects() method

python
for partial_object in program.stream_partial_objects(movie="Harry Potter"):
    # send the partial object to the frontend for better user experience
    print(partial_object)

Extracting List of Album (with Parallel Function Calling)

With the latest parallel function calling feature from OpenAI, we can simultaneously extract multiple structured data from a single prompt!

To do this, we need to:

  1. pick one of the latest models (e.g. gpt-3.5-turbo-1106), and
  2. set allow_multiple to True in our OpenAIPydanticProgram (if not, it will only return the first object, and raise a warning).
python
from llama_index.llms.openai import OpenAI

prompt_template_str = """\
Generate 4 albums about spring, summer, fall, and winter.
"""
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    llm=OpenAI(model="gpt-3.5-turbo-1106"),
    prompt_template_str=prompt_template_str,
    allow_multiple=True,
    verbose=True,
)
python
output = program()

The output is a list of valid Pydantic object.

python
output

Extraction into Album (Streaming)

We also support streaming a list of objects through our stream_list function.

Full credits to this idea go to openai_function_call repo: https://github.com/jxnl/openai_function_call/tree/main/examples/streaming_multitask

python
prompt_template_str = "{input_str}"
program = OpenAIPydanticProgram.from_defaults(
    output_cls=Album,
    prompt_template_str=prompt_template_str,
    verbose=False,
)

output = program.stream_list(
    input_str="make up 5 random albums",
)
for obj in output:
    print(obj.json(indent=2))

Extraction into DirectoryTree object

This is directly inspired by jxnl's awesome repo here: https://github.com/jxnl/openai_function_call.

That repository shows how you can use OpenAI's function API to parse recursive Pydantic objects. The main requirement is that you want to "wrap" a recursive Pydantic object with a non-recursive one.

Here we show an example in a "directory" setting, where a DirectoryTree object wraps recursive Node objects, to parse a file structure.

python
# NOTE: defining recursive objects in a notebook causes errors
from directory import DirectoryTree, Node
python
DirectoryTree.schema()
python
program = OpenAIPydanticProgram.from_defaults(
    output_cls=DirectoryTree,
    prompt_template_str="{input_str}",
    verbose=True,
)
python
input_str = """
root
ā”œā”€ā”€ folder1
ā”‚   ā”œā”€ā”€ file1.txt
ā”‚   ā””ā”€ā”€ file2.txt
ā””ā”€ā”€ folder2
    ā”œā”€ā”€ file3.txt
    ā””ā”€ā”€ subfolder1
        ā””ā”€ā”€ file4.txt
"""

output = program(input_str=input_str)

The output is a full DirectoryTree structure with recursive Node objects.

python
output