llamaindex_openai_pydantic_program_structured_extraction.py

1import os
2from pydantic import BaseModel
3from llama_index.program.openai import OpenAIPydanticProgram
4from llama_index.llms.openai import OpenAI
5
6# Define the output schema
7class Album(BaseModel):
8    """Data model for an album."""
9    album_title: str
10    artist: str
11    year: int
12
13# Initialize the program
14prompt_template_str = "Generate an example album, with an artist and a release year, for the genre {genre}."
15
16program = OpenAIPydanticProgram.from_defaults(
17    output_cls=Album,
18    prompt_template_str=prompt_template_str,
19    llm=OpenAI(model="gpt-3.5-turbo-0613"),
20    verbose=True,
21)
22
23# Run the program to get structured output
24output = program(genre="punk rock")
25
26print(f"Album Title: {output.album_title}")
27print(f"Artist: {output.artist}")
28print(f"Year: {output.year}")