Spaces:
Sleeping
Sleeping
File size: 2,747 Bytes
e40cdd9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
"""
Requires:
/usr/local/bin/python3.13 -m pip install python-dotenv pydantic-ai imdbpy
OpenAI API Key
"""
from dataclasses import dataclass
import imdb
import asyncio
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
import gradio as gr
from pathlib import Path
from dotenv import load_dotenv
import os
script_folder = Path(__file__).parent
dotenv_path = script_folder/"Open_AI.env"
load_dotenv(dotenv_path=dotenv_path)
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
# A pydantic basemodel for our output schema
class Review(BaseModel):
why: str
rating: int
recommended: bool
# A class to externally connect to the IMDB API to fetch movie information
class IMDbConnection:
def __init__(self):
self.ia = imdb.IMDb()
async def get_movie_info(self,title: str) -> dict:
# Search for the movie using IMDB API
print(f"Searching for movie: {title}")
movies = self.ia.search_movie(title)
movie = self.ia.get_movie(movies[0].movieID)
system_prompt = f"""
Title: {movie.get("title")}\n
Rating: {movie.get("rating")}\n
Plot: {movie.get("plot")[0] if movie.get("plot") else None}
"""
print(f"System prompt: {system_prompt}\n\n")
return system_prompt
# A container containing dependencies for the agent
@dataclass
class MovieData:
title: str
imdb_conn: IMDbConnection
agent = Agent(
"openai:gpt-4o-mini",
deps_type=MovieData,
output_type=Review,
)
# Dynamically generate the system prompt
@agent.system_prompt
async def get_movie_info(ctx: RunContext[Review]):
# Given an input, the system prompt will be generated by querying the IMDB API
return await ctx.deps.imdb_conn.get_movie_info(ctx.deps.title) #Fetch movie info
# create a gradio wrapper
async def run_agent(user_query: str, movie_title: str):
deps = MovieData(title=movie_title, imdb_conn=IMDbConnection())
result = await agent.run(user_query, deps=deps)
return {
"Why": result.output.why,
"Rating": result.output.rating,
"Recommend": result.output.recommended
}
#run agent in gradio
def run_gradio(user_query, movie_title):
return asyncio.run(run_agent(user_query, movie_title))
# build gradio ui
with gr.Blocks() as demo:
gr.Markdown("## Movie Recommender Agent")
user_query = gr.Textbox(label="What is your preference? (e.g. 'I like supernatural horror movies')")
movie_title = gr.Textbox(label="Movie name (e.g. 'The Conjuring')")
output = gr.JSON(label = "Agent Recommendation")
submit_btn = gr.Button("Check Movie")
submit_btn.click(fn=run_gradio, inputs=[user_query, movie_title], outputs=output)
#launch
if __name__ == "__main__":
demo.launch(share=True) |