src/tools.py

"""Tools for the agent."""

import math
from collections.abc import Callable
from typing import ClassVar

from langchain_community.agent_toolkits.load_tools import load_tools
from langchain_community.retrievers import BM25Retriever
from smolagents import Tool, tool

from src.documents import get_documents


# Tool to list the available occasions
@tool
def list_occasions() -> str:
    """List the available occasions."""
    return "casual, formal, superhero, custom"


# Tool to suggest a menu based on the occasion
@tool
def suggest_menu(occasion: str) -> str:
    """
    Suggest a menu based on the occasion.

    Args:
        occasion (str): The type of occasion for the party. Allowed values are:
                        - "casual": Menu for casual party.
                        - "formal": Menu for formal party.
                        - "superhero": Menu for superhero party.
                        - "custom": Custom menu.

    """
    return {
        "casual": "Pizza, snacks, and drinks.",
        "formal": "3-course dinner with wine and dessert.",
        "superhero": "Buffet with high-energy and healthy food.",
    }.get(occasion, "Custom menu for the butler.")


# function that fetches the highest-rated catering services.
@tool
def catering_service_tool() -> str:
    """
    Tool that returns the highest-rated catering service in Gotham City.

    Args:
        query: A search term for finding catering services.

    """
    # Example list of catering services and their ratings
    services = {
        "Gotham Catering Co.": 4.9,
        "Wayne Manor Catering": 4.8,
        "Gotham City Events": 4.7,
    }

    # Find the highest rated catering service (simulating search query filtering)
    return max(services, key=services.get)


class SuperheroPartyThemeTool(Tool):
    """Super Hero Party Theme Generator - Tool for agent."""

    name = "superhero_party_theme_generator"
    description = """
    This tool suggests creative superhero-themed party ideas based on a category.
    It returns a unique party theme idea."""

    inputs: ClassVar[dict] = {
        "category": {
            "type": "string",
            "description": """
            The type of superhero party
            (e.g., 'classic heroes','villain masquerade', 'futuristic Gotham').
            """,
        },
    }

    output_type = "string"

    def forward(self, category: str) -> str:
        """Generate a superhero-themed party theme based on the category."""
        themes = {
            "classic heroes": """
            Justice League Gala: Guests come dressed as their favorite DC heroes
            with themed cocktails like 'The Kryptonite Punch'.
            """,
            "villain masquerade": """
            Gotham Rogues' Ball: A mysterious masquerade
            where guests dress as classic Batman villains.
            """,
            "futuristic Gotham": """
            Neo-Gotham Night: A cyberpunk-style party inspired
            by Batman Beyond, with neon decorations and futuristic gadgets.
            """,
        }

        return themes.get(
            category.lower(),
            """
            Themed party idea not found. Try 'classic heroes', 'villain masquerade',
            or 'futuristic Gotham'.
            """,
        )


def get_langchain_serpapi_tool() -> Callable:
    """Return a tool that uses the SerpAPI to search the web."""
    serpapi_search_tool = Tool.from_langchain(load_tools(["serpapi"])[0])
    serpapi_search_tool.name = "serpapi_search_tool"
    return serpapi_search_tool


def get_image_generation_tool() -> Callable:
    """Return a tool that generates an image based on a prompt."""
    return Tool.from_space(
        "black-forest-labs/FLUX.1-schnell",
        name="image_generator",
        description="Generate an image from a prompt",
    )


class PartyPlanningRetrieverTool(Tool):
    """Retrieve relevant party planning ideas for Alfred's party."""

    name = "party_planning_retriever"
    description = """
        Uses semantic search to retrieve relevant party planning ideas for
        Alfred's superhero-themed party at Wayne Manor.
        """

    inputs: ClassVar[dict] = {
        "query": {
            "type": "string",
            "description": """
                The query to perform. This should be a query related to
                party planning or superhero themes.
            """,
        },
    }
    output_type = "string"

    def __init__(self, **kwargs) -> None:  # noqa: ANN003
        """Initialize the retriever tool."""
        super().__init__(**kwargs)
        docs = get_documents()
        self.retriever = BM25Retriever.from_documents(
            docs,
            k=5,  # Retrieve the top 5 documents
        )

    def forward(self, query: str) -> str:
        """Retrieve relevant party planning ideas for Alfred's party."""
        docs = self.retriever.invoke(
            query,
        )
        return "\nRetrieved ideas:\n" + "".join(
            [
                f"\n\n===== Idea {i!s} =====\n" + doc.page_content
                for i, doc in enumerate(docs)
            ],
        )


@tool
def calculate_cargo_travel_time(
    origin_coords: tuple[float, float],
    destination_coords: tuple[float, float],
    cruising_speed_kmh: float | None = 750.0,  # Average speed for cargo planes
) -> float:
    """
    Calculate the travel time for a cargo plane between two points on Earth.

    It great-circle distance.

    Args:
        origin_coords: Tuple of (latitude, longitude) for the starting point.
        destination_coords: Tuple of (latitude, longitude) for the destination.
        cruising_speed_kmh: Optional cruising speed in km/h (defaults to 750 km/h for
            typical cargo planes).

    Returns:
        float: The estimated travel time in hours

    Example:
        >>> # Chicago (41.8781° N, 87.6298° W) to Sydney (33.8688° S, 151.2093° E)
        >>> result = calculate_cargo_travel_time((41.878, -87.629), (-33.868, 151.209))

    """

    def to_radians(degrees: float) -> float:
        return degrees * (math.pi / 180)

    # Extract coordinates
    lat1, lon1 = map(to_radians, origin_coords)
    lat2, lon2 = map(to_radians, destination_coords)

    # Earth's radius in kilometers
    earth_radius_km = 6371.0

    # Calculate great-circle distance using the haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = (
        math.sin(dlat / 2) ** 2
        + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2
    )
    c = 2 * math.asin(math.sqrt(a))
    distance = earth_radius_km * c

    # Add 10% to account for non-direct routes and air traffic controls
    actual_distance = distance * 1.1

    # Calculate flight time
    # Add 1 hour for takeoff and landing procedures
    flight_time = (actual_distance / cruising_speed_kmh) + 1.0

    # Format the results
    return round(flight_time, 2)