llama_index_openai_pydantic_program_structured_extraction.py

1import os
2from pydantic import BaseModel
3from llama_index.program.openai import OpenAIPydanticProgram
4
5# Define the output schema
6class Song(BaseModel):
7    """Data model for a song."""
8    title: str
9    length_seconds: int
10
11# Initialize the program
12prompt_template_str = "Generate a song with the title {song_title}"
13program = OpenAIPydanticProgram.from_defaults(
14    output_cls=Song, 
15    prompt_template_str=prompt_template_str, 
16    verbose=True
17)
18
19# Run the program to get structured output
20output = program(song_title="Happy Birthday")
21
22print(f"Title: {output.title}")
23print(f"Length: {output.length_seconds} seconds")