from typing import List
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI

# Define our Movie and FilmFestival models
class Movie(BaseModel):
    title: str = Field(..., description="The title of the movie")
    director: str = Field(..., description="The director of the movie")
    runtime: int = Field(..., description="The runtime of the movie in minutes")

class FilmFestival(BaseModel):
    movies: List[Movie]

# Set up the parser
parser = PydanticOutputParser(pydantic_object=FilmFestival)

# Create the prompt template
prompt = ChatPromptTemplate.from_messages([
    ("system", "Answer the user query about movies in the film festival. Wrap the output in `json` format following the schema below\n{format_instructions}"),
    ("human", "{query}"),
]).partial(format_instructions=parser.get_format_instructions())

# Set up the LLM and chain
llm = ChatOpenAI(model="gpt-4", temperature=0)
chain = prompt | llm | parser

# Generate movie information
query = "Please provide details about the movies 'Inception' directed by Christopher Nolan with a runtime of 148 minutes and 'Parasite' directed by Bong Joon-ho with a runtime of 132 minutes."
result = chain.invoke({"query": query})

print(result)
# Expected output: FilmFestival(movies=[Movie(title='Inception', director='Christopher Nolan', runtime=148), Movie(title='Parasite', director='Bong Joon-ho', runtime=132)])

        __
        __