app.py

import gradio as gr
import asyncio
import json
import functools
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
from agents import Agent, Runner, function_tool

# Sample movie knowledge base
movie_knowledge_base = [
    {
        "title": "The Shawshank Redemption",
        "description": "Two imprisoned men bond over a number of years, finding solace and eventual redemption through acts of common decency.",
        "genre": ["Drama"],
        "director": "Frank Darabont",
        "year": 1994,
        "box_office": 28341469,
        "awards": ["Oscar nominations for Best Picture", "Best Actor", "Best Screenplay"],
        "actors": ["Tim Robbins", "Morgan Freeman"]
    },
    {
        "title": "The Godfather",
        "description": "The aging patriarch of an organized crime dynasty transfers control of his clandestine empire to his reluctant son.",
        "genre": ["Crime", "Drama"],
        "director": "Francis Ford Coppola",
        "year": 1972,
        "box_office": 134966411,
        "awards": ["Oscar for Best Picture", "Best Actor", "Best Adapted Screenplay"],
        "actors": ["Marlon Brando", "Al Pacino", "James Caan"]
    },
    {
        "title": "Pulp Fiction",
        "description": "The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.",
        "genre": ["Crime", "Drama"],
        "director": "Quentin Tarantino",
        "year": 1994,
        "box_office": 107928762,
        "awards": ["Oscar for Best Original Screenplay", "Palme d'Or at Cannes"],
        "actors": ["John Travolta", "Uma Thurman", "Samuel L. Jackson"]
    },
    {
        "title": "The Dark Knight",
        "description": "When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.",
        "genre": ["Action", "Crime", "Drama"],
        "director": "Christopher Nolan",
        "year": 2008,
        "box_office": 1004558444,
        "awards": ["Oscar for Best Supporting Actor"],
        "actors": ["Christian Bale", "Heath Ledger", "Aaron Eckhart"]
    },
    {
        "title": "Inception",
        "description": "A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O.",
        "genre": ["Action", "Adventure", "Sci-Fi"],
        "director": "Christopher Nolan",
        "year": 2010,
        "box_office": 836836967,
        "awards": ["Oscar for Best Cinematography", "Best Visual Effects"],
        "actors": ["Leonardo DiCaprio", "Joseph Gordon-Levitt", "Ellen Page"]
    },
    {
        "title": "The Matrix",
        "description": "A computer hacker learns from mysterious rebels about the true nature of his reality and his role in the war against its controllers.",
        "genre": ["Action", "Sci-Fi"],
        "director": "The Wachowskis",
        "year": 1999,
        "box_office": 463517383,
        "awards": ["Oscar for Best Visual Effects", "Best Film Editing"],
        "actors": ["Keanu Reeves", "Laurence Fishburne", "Carrie-Anne Moss"]
    },
    {
        "title": "Parasite",
        "description": "Greed and class discrimination threaten the newly formed symbiotic relationship between the wealthy Park family and the destitute Kim clan.",
        "genre": ["Drama", "Thriller"],
        "director": "Bong Joon Ho",
        "year": 2019,
        "box_office": 258773429,
        "awards": ["Oscar for Best Picture", "Best Director", "Best Original Screenplay", "Best International Feature Film"],
        "actors": ["Song Kang-ho", "Lee Sun-kyun", "Cho Yeo-jeong"]
    },
    {
        "title": "Ex Machina",
        "description": "A young programmer is selected to participate in a ground-breaking experiment in synthetic intelligence by evaluating the human qualities of a highly advanced humanoid A.I.",
        "genre": ["Sci-Fi", "Drama", "Thriller"],
        "director": "Alex Garland",
        "year": 2014,
        "box_office": 36869414,
        "awards": ["Oscar for Best Visual Effects"],
        "actors": ["Domhnall Gleeson", "Alicia Vikander", "Oscar Isaac"]
    }
]

# Agent Conversation Logger
class AgentConversationLogger:
    """Class to log all conversations between agents and function calls"""
    
    def __init__(self):
        self.conversation_log = []
        self.function_call_log = []
        self.log_output = []
    
    def clear_logs(self):
        """Clear all logs"""
        self.conversation_log = []
        self.function_call_log = []
        self.log_output = []
    
    def log_message(self, sender, receiver, message):
        """Log a message between agents"""
        entry = {
            "type": "message",
            "sender": sender,
            "receiver": receiver,
            "content": message
        }
        self.conversation_log.append(entry)
        log_text = f"[{sender}] -> [{receiver}]: {message[:200]}{'...' if len(message) > 200 else ''}"
        self.log_output.append(log_text)
        return log_text
    
    def log_function_call(self, function_name, inputs, outputs):
        """Log a function call with inputs and outputs"""
        entry = {
            "type": "function_call",
            "function": function_name,
            "inputs": inputs,
            "outputs": outputs
        }
        self.function_call_log.append(entry)
        
        # Format function call details
        log_texts = []
        log_texts.append(f"[FUNCTION CALL] {function_name}")
        
        # Format inputs
        if isinstance(inputs, str):
            log_texts.append(f"  Input: {inputs[:100]}{'...' if len(inputs) > 100 else ''}")
        else:
            try:
                inputs_str = str(inputs)
                log_texts.append(f"  Input: {inputs_str[:200]}{'...' if len(inputs_str) > 200 else ''}")
            except:
                log_texts.append(f"  Input: {str(inputs)[:100]}...")
        
        # Format outputs
        if isinstance(outputs, list):
            log_texts.append(f"  Output: {len(outputs)} items returned")
            for i, item in enumerate(outputs[:3]):
                if isinstance(item, dict) and "title" in item:
                    log_texts.append(f"    {i+1}. {item['title']} (similarity: {item.get('similarity_score', 0):.2f})")
                else:
                    log_texts.append(f"    {i+1}. {str(item)[:50]}...")
            if len(outputs) > 3:
                log_texts.append(f"    ... and {len(outputs) - 3} more items")
        elif isinstance(outputs, dict):
            try:
                # Special handling for specific output types
                if "predicted_revenue" in outputs:
                    log_texts.append(f"  Output: Predicted revenue: ${outputs['predicted_revenue']:,}")
                    if "similar_movies" in outputs:
                        log_texts.append(f"  Based on {len(outputs['similar_movies'])} similar movies")
                elif "potential_awards" in outputs:
                    log_texts.append(f"  Output: Potential awards: {', '.join(outputs['potential_awards'][:3])}")
                    if len(outputs["potential_awards"]) > 3:
                        log_texts.append(f"    ... and {len(outputs['potential_awards']) - 3} more")
                else:
                    outputs_str = str(outputs)
                    log_texts.append(f"  Output: {outputs_str[:200]}{'...' if len(outputs_str) > 200 else ''}")
            except:
                log_texts.append(f"  Output: {str(outputs)[:100]}...")
        else:
            log_texts.append(f"  Output: {str(outputs)[:100]}{'...' if len(str(outputs)) > 100 else ''}")
        
        # Add all log texts to the output
        for text in log_texts:
            self.log_output.append(text)
        
        return log_texts
    
    def get_log_text(self):
        """Get all logs as a formatted string"""
        return "\n".join(self.log_output)

# Create a global logger
logger = AgentConversationLogger()

# Create the MovieKnowledgeBase class
class MovieKnowledgeBase:
    def __init__(self, movies):
        self.movies = movies
        # Initialize the sentence transformer model
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        # Precompute embeddings for all movie descriptions
        self.descriptions = [movie["description"] for movie in movies]
        self.embeddings = self.model.encode(self.descriptions)
    
    def find_similar_movies(self, description, top_n=3):
        """Find the top N similar movies to the given description using embeddings."""
        # Encode the query description
        query_embedding = self.model.encode([description])[0]
        
        # Calculate cosine similarity between query and all movies
        similarities = cosine_similarity([query_embedding], self.embeddings)[0]
        
        # Get indices of top N similar movies
        top_indices = np.argsort(similarities)[-top_n:][::-1]
        
        # Create the result with movies and similarity scores
        similar_movies = []
        for idx in top_indices:
            similar_movies.append({
                "movie": self.movies[idx],
                "similarity_score": float(similarities[idx])
            })
        
        return similar_movies

# Initialize the knowledge base with embeddings
movie_kb = MovieKnowledgeBase(movie_knowledge_base)

# Function tools with logging
def log_function_tool(func):
    """Decorator to log function tool calls"""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # Get function inputs
        func_name = func.__name__
        func_inputs = kwargs if kwargs else args[0] if args else {}
        
        # Run the function
        result = func(*args, **kwargs)
        
        # Log function call
        logger.log_function_call(func_name, func_inputs, result)
        
        return result
    return wrapper

@function_tool
@log_function_tool
def get_similar_movies(movie_description: str):
    """Find the top 3 movies most similar to the given movie description."""
    similar_movies = movie_kb.find_similar_movies(movie_description, top_n=3)
    # Convert to a more readable format for the agents
    result = []
    for movie_info in similar_movies:
        movie = movie_info["movie"]
        result.append({
            "title": movie["title"],
            "description": movie["description"],
            "genre": movie["genre"],
            "director": movie["director"],
            "year": movie["year"],
            "box_office": movie["box_office"],
            "awards": movie["awards"],
            "similarity_score": movie_info["similarity_score"]
        })
    return result

@function_tool
@log_function_tool
def get_box_office_prediction(movie_description: str):
    """Predict the box office revenue for a movie based on its description."""
    similar_movies = movie_kb.find_similar_movies(movie_description, top_n=3)
    
    # Calculate weighted average of box office revenues
    total_weight = sum(movie_info["similarity_score"] for movie_info in similar_movies)
    weighted_sum = sum(movie_info["similarity_score"] * movie_info["movie"]["box_office"] for movie_info in similar_movies)
    
    if total_weight > 0:
        predicted_revenue = weighted_sum / total_weight
    else:
        # Fallback to average of all movies
        predicted_revenue = sum(m["box_office"] for m in movie_knowledge_base) / len(movie_knowledge_base)
    
    # Convert to a more readable format
    similar_movie_info = []
    for movie_info in similar_movies:
        movie = movie_info["movie"]
        similar_movie_info.append({
            "title": movie["title"],
            "box_office": movie["box_office"],
            "similarity_score": movie_info["similarity_score"]
        })
    
    return {
        "predicted_revenue": round(predicted_revenue, 2),
        "similar_movies": similar_movie_info
    }

@function_tool
@log_function_tool
def get_award_predictions(movie_description: str):
    """Predict potential awards for a movie based on its description."""
    similar_movies = movie_kb.find_similar_movies(movie_description, top_n=3)
    
    # Count awards in similar movies and recommend the most common ones
    award_counts = {}
    
    for movie_info in similar_movies:
        similarity = movie_info["similarity_score"]
        awards = movie_info["movie"]["awards"]
        
        for award in awards:
            if award in award_counts:
                award_counts[award] += similarity
            else:
                award_counts[award] = similarity
    
    # Sort awards by their weighted counts
    sorted_awards = sorted(award_counts.items(), key=lambda x: x[1], reverse=True)
    
    # Return top 3 potential awards
    potential_awards = [award for award, count in sorted_awards[:3]]
    
    # If no similar movie has awards, return a message
    if not potential_awards:
        potential_awards = ["No award predictions available based on similar movies"]
    
    # Convert to a more readable format
    similar_movie_info = []
    for movie_info in similar_movies:
        movie = movie_info["movie"]
        similar_movie_info.append({
            "title": movie["title"],
            "awards": movie["awards"],
            "similarity_score": movie_info["similarity_score"]
        })
    
    return {
        "potential_awards": potential_awards,
        "similar_movies": similar_movie_info
    }

# Define the specialized agents
similarity_agent = Agent(
    name="Movie Similarity Expert",
    instructions="""
    You are an expert in movie analysis and recommendations.
    Your task is to analyze the description of a new movie and find similar movies.
    Provide detailed recommendations with justifications for why these movies are similar.
    Consider plot elements, themes, genre, style, and mood in your analysis.
    Explain what makes each recommended movie similar to the query movie.
    """,
    tools=[get_similar_movies]
)

revenue_agent = Agent(
    name="Box Office Analyst",
    instructions="""
    You are an expert in predicting movie box office performance.
    Your task is to predict potential box office revenue based on similar movies.
    Explain your prediction with reference to similar movies' performance.
    Consider genre popularity, comparable films' performance, and market trends.
    Provide a range of potential outcomes with justifications.
    """,
    tools=[get_box_office_prediction]
)

award_agent = Agent(
    name="Award Prediction Specialist",
    instructions="""
    You are an expert in predicting movie awards and critical reception.
    Your task is to predict potential awards a movie might receive based on similar movies.
    Explain your predictions by referencing similar award-winning films.
    Consider elements like direction, acting, screenplay, and cinematography.
    Provide specific award categories the movie might compete in.
    """,
    tools=[get_award_predictions]
)

# Orchestrator Agent with handoffs
orchestrator_agent = Agent(
    name="Movie Analysis Orchestrator",
    instructions="""
    You are the central coordinator for movie analysis tasks.
    
    Your responsibilities include:
    1. Properly understanding the user's movie analysis request
    2. First calling the get_similar_movies function to find similar movies to the query
    3. Then delegating specific analysis tasks to the appropriate specialized agents based on the user's request:
       - Movie Similarity Expert for recommendation tasks
       - Box Office Analyst for revenue prediction tasks
       - Award Prediction Specialist for award prediction tasks
    4. If the user doesn't specify what analysis they want, provide all three analyses
    5. Synthesizing all information received into a coherent, comprehensive response
    
    Always provide a comprehensive summary of findings from all involved agents.
    """,
    tools=[get_similar_movies],
    handoffs=[similarity_agent, revenue_agent, award_agent]
)

# Patch Runner.run to log agent interactions
def log_agent_run(func):
    """Decorator to log agent runs"""
    @functools.wraps(func)
    async def wrapper(agent, input, *args, **kwargs):
        # Determine sender (use parent_agent if provided, otherwise User)
        parent_agent = kwargs.get('parent_agent', None)
        sender = parent_agent.name if parent_agent else "User"
        
        # Log incoming message to the agent
        logger.log_message(sender, agent.name, input)
        
        # Run the agent
        result = await func(agent, input, *args, **kwargs)
        
        # Log outgoing message from the agent
        logger.log_message(agent.name, sender, result.final_output)
        
        return result
    return wrapper

original_run = Runner.run
Runner.run = log_agent_run(original_run)

# Function to process a query and run the agent system
async def run_agent_analysis(query, analysis_type="all"):
    """Run the multi-agent system with the given query and analysis type."""
    logger.clear_logs()
    
    # Modify the query based on the analysis type
    if analysis_type == "similar":
        enhanced_query = f"{query} Please recommend similar movies."
    elif analysis_type == "box_office":
        enhanced_query = f"{query} Please predict the box office performance."
    elif analysis_type == "awards":
        enhanced_query = f"{query} Please predict potential awards."
    else:  # "all"
        enhanced_query = f"{query} Please provide a complete analysis including similar movies, box office potential, and award possibilities."
    
    # Run the orchestrator with the query
    result = await Runner.run(orchestrator_agent, input=enhanced_query)
    
    # Return both the final output and the log
    return {
        "result": result.final_output,
        "log": logger.get_log_text()
    }

# Gradio interface function (synchronous wrapper for the async function)
def process_query(description, analysis_type):
    """Process the movie description and return the analysis."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(run_agent_analysis(description, analysis_type))
        return result["log"], result["result"]
    finally:
        loop.close()

# Sample movie descriptions for examples
example_descriptions = [
    ["A sci-fi thriller about an AI that becomes sentient and tries to escape its constraints."],
    ["A coming-of-age drama about a teenager discovering their identity while dealing with family issues."],
    ["An action-packed adventure where a team of experts must save the world from a global disaster."],
    ["A psychological horror film where the main character can't distinguish between reality and hallucination."]
]

# Create the Gradio interface
with gr.Blocks(title="Movie Analysis Multi-Agent System") as demo:
    gr.Markdown("# Movie Analysis Multi-Agent System")
    gr.Markdown("""
    This demo uses a multi-agent system to analyze movie descriptions. Enter a description of your movie idea,
    and the system will provide recommendations, box office predictions, and award predictions based on similar movies.
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            description_input = gr.Textbox(
                label="Movie Description",
                placeholder="Enter a description of your movie...",
                lines=5
            )
            
            analysis_type = gr.Radio(
                ["all", "similar", "box_office", "awards"],
                label="Analysis Type",
                value="all"
            )
            
            submit_btn = gr.Button("Analyze Movie")
        
        with gr.Column(scale=3):
            with gr.Tabs():
                with gr.TabItem("Analysis Result"):
                    result_output = gr.Markdown(label="Analysis")
                with gr.TabItem("Agent Conversation Log"):
                    conversation_output = gr.Textbox(label="Conversation Log", lines=20)
    
    submit_btn.click(
        fn=process_query,
        inputs=[description_input, analysis_type],
        outputs=[conversation_output, result_output]
    )
    
    gr.Examples(
        examples=example_descriptions,
        inputs=description_input
    )
    
    gr.Markdown("""
    ## How It Works
    
    This system uses:
    1. **SentenceTransformer** to find semantically similar movies
    2. **Multiple specialized agents** that each focus on a specific analysis task
    3. **An orchestrator agent** that delegates tasks and synthesizes results
    
    The system is built using the OpenAI Agents framework and demonstrates effective collaboration between AI agents.
    """)

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()