tools.py

import base64
import datetime
import math
import os
import urllib.parse
from pathlib import Path
from typing import Dict, Union

import pandas
import pytz
import requests
from bs4 import BeautifulSoup
from langchain_community.document_loaders import (
    ArxivLoader, AssemblyAIAudioTranscriptLoader, WikipediaLoader)
from langchain_community.document_loaders.generic import GenericLoader
from langchain_community.document_loaders.parsers import LanguageParser
from langchain_core.messages import HumanMessage
# from langchain_community.tools import DuckDuckGoSearchRun
# from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_tavily import TavilySearch


@tool
def download_file(url: str, filename: str = None) -> str:
    """
    Download a file from a URL and save it locally for analysis.
    Analysis can be done using another tool.
    
    Args:
        url: The URL of the file to download
        filename: Optional custom filename. If not provided, will extract from URL
        
    Returns:
        The local file path where the file was saved
    """
    try:
        # Clean and validate URL
        url = url.strip()
        if not url.startswith(('http://', 'https://')):
            return f"Error: Invalid URL format: {url}"
        
        # Create downloads directory if it doesn't exist
        download_dir = Path("downloads")
        download_dir.mkdir(exist_ok=True)
        
        # Determine filename
        if not filename:
            # Extract filename from URL
            parsed_url = urllib.parse.urlparse(url)
            filename = os.path.basename(parsed_url.path)
            
            # If no filename in URL, create one based on URL hash
            if not filename or '.' not in filename:
                url_hash = str(hash(url))[-8:]  # Last 8 characters of hash
                filename = f"downloaded_file_{url_hash}"
        
        # Ensure filename has proper extension if we can detect content type
        filepath = download_dir / filename
        
        # Download the file
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        
        print(f"Downloading file from: {url}")
        response = requests.get(url, headers=headers, timeout=30, stream=True)
        response.raise_for_status()
        
        # Try to get better filename from Content-Disposition header
        if 'content-disposition' in response.headers:
            import re
            cd = response.headers['content-disposition']
            filename_match = re.findall('filename=(.+)', cd)
            if filename_match:
                suggested_filename = filename_match[0].strip('"')
                if suggested_filename:
                    filepath = download_dir / suggested_filename
        
        # Write file to disk
        with open(filepath, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        file_size = os.path.getsize(filepath)
        print(f"Successfully downloaded {filename} ({file_size} bytes)")
        
        return str(filepath)
        
    except requests.exceptions.RequestException as e:
        return f"Error downloading file: {str(e)}"
    except Exception as e:
        return f"Error saving file: {str(e)}"


@tool
def multiply(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """Multiplies two numbers and returns the product.

    Args:
        a: The first number.
        b: The second number.

    Returns:
        The product of the two input numbers.
    """
    try:
        result = a * b
        return int(result) if isinstance(a, int) and isinstance(b, int) else result
    except Exception as e:
        return f"Error in multiplication: {str(e)}"


@tool
def add(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """Adds two numbers and returns the sum.

    Args:
        a: The first number.
        b: The second number.

    Returns:
        The sum of the two input numbers.
    """
    try:
        result = a + b
        return int(result) if isinstance(a, int) and isinstance(b, int) else result
    except Exception as e:
        return f"Error in addition: {str(e)}"


@tool
def power(a: Union[int, float], b: Union[int, float]) -> float:
    """Raises a number to the power of another.

    Args:
        a: The base number.
        b: The exponent.

    Returns:
        The result of raising `a` to the power of `b`.
    """
    try:
        if a == 0 and b < 0:
            return "Error: Cannot raise 0 to a negative power"
        result = a ** b
        return result
    except OverflowError:
        return "Error: Result too large to compute"
    except Exception as e:
        return f"Error in power calculation: {str(e)}"


@tool
def subtract(a: Union[int, float], b: Union[int, float]) -> Union[int, float]:
    """Subtracts the second number from the first.

    Args:
        a: The number from which to subtract.
        b: The number to subtract.

    Returns:
        The result of `a` minus `b`.
    """
    try:
        result = a - b
        return int(result) if isinstance(a, int) and isinstance(b, int) else result
    except Exception as e:
        return f"Error in subtraction: {str(e)}"


@tool
def divide(a: Union[int, float], b: Union[int, float]) -> float:
    """Divides one number by another.

    Args:
        a: The numerator.
        b: The denominator.

    Returns:
        The result of `a` divided by `b`.
    """
    try:
        if b == 0:
            return "Error: Division by zero is not allowed"
        return a / b
    except Exception as e:
        return f"Error in division: {str(e)}"


@tool
def modulus(a: int, b: int) -> Union[int, str]:
    """Returns the remainder of the division of two integers.

    Args:
        a: The dividend.
        b: The divisor.

    Returns:
        The remainder when `a` is divided by `b`.
    """
    try:
        if b == 0:
            return "Error: Modulus by zero is not allowed"
        return a % b
    except Exception as e:
        return f"Error in modulus operation: {str(e)}"


@tool
def square_root(x: Union[int, float]) -> Union[float, str]:
    """Returns the square root of a number.

    Args:
        x: The input number. Must be non-negative.

    Returns:
        The square root of `x`.
    """
    try:
        if x < 0:
            return "Error: Square root of negative number is not allowed"
        return math.sqrt(x)
    except Exception as e:
        return f"Error in square root calculation: {str(e)}"


@tool
def floor_divide(a: int, b: int) -> Union[int, str]:
    """Performs integer division (floor division) of two numbers.

    Args:
        a: The dividend.
        b: The divisor.

    Returns:
        The floor of the quotient.
    """
    try:
        if b == 0:
            return "Error: Division by zero is not allowed"
        return a // b
    except Exception as e:
        return f"Error in floor division: {str(e)}"


@tool
def absolute(x: Union[int, float]) -> Union[int, float]:
    """Returns the absolute value of a number.

    Args:
        x: The input number.

    Returns:
        The absolute value of `x`.
    """
    try:
        result = abs(x)
        return int(result) if isinstance(x, int) else result
    except Exception as e:
        return f"Error in absolute value calculation: {str(e)}"


@tool
def logarithm(x: Union[int, float], base: Union[int, float] = math.e) -> Union[float, str]:
    """Returns the logarithm of a number with a given base.

    Args:
        x: The number to take the logarithm of. Must be positive.
        base: The logarithmic base. Must be positive and not equal to 1.

    Returns:
        The logarithm of `x` to the given base.
    """
    try:
        if x <= 0:
            return "Error: Logarithm input must be positive"
        if base <= 0 or base == 1:
            return "Error: Logarithm base must be positive and not equal to 1"
        return math.log(x, base)
    except Exception as e:
        return f"Error in logarithm calculation: {str(e)}"


@tool
def exponential(x: Union[int, float]) -> Union[float, str]:
    """Returns e raised to the power of `x`.

    Args:
        x: The exponent.

    Returns:
        The value of e^x.
    """
    try:
        if x > 700:  # Prevent overflow
            return "Error: Exponent too large, would cause overflow"
        return math.exp(x)
    except OverflowError:
        return "Error: Result too large to compute"
    except Exception as e:
        return f"Error in exponential calculation: {str(e)}"


@tool
def roman_calculator_converter(value1: int, value2: int, oper: str) -> str:
    """Performs an operation on 2 numbers and returns the result as a Roman numeral.
    
    Args:
        value1: The first value
        value2: The second value
        oper: Operator for the calculation ("add", "subtract", "multiply", "divide")
    
    Returns:
        The result as a Roman numeral string.
    """
    try:
        # Input validation
        if not isinstance(value1, int) or not isinstance(value2, int):
            return "Error: Both values must be integers"
        
        if oper not in ["add", "subtract", "multiply", "divide"]:
            return "Error: Operator must be 'add', 'subtract', 'multiply', or 'divide'"
        
        # Roman numeral mapping
        roman_numerals = [
            (1000, "M"), (900, "CM"), (500, "D"), (400, "CD"),
            (100, "C"), (90, "XC"), (50, "L"), (40, "XL"),
            (10, "X"), (9, "IX"), (5, "V"), (4, "IV"), (1, "I")
        ]
        
        # Perform calculation
        if oper == "add":
            result = value1 + value2
        elif oper == "subtract":
            result = value1 - value2
        elif oper == "multiply":
            result = value1 * value2
        elif oper == "divide":
            if value2 == 0:
                return "Error: Division by zero is not allowed"
            result = int(value1 / value2)  # Integer division for Roman numerals
        
        # Handle invalid results for Roman numerals
        if result <= 0:
            return f"Error: Roman numerals cannot represent zero or negative numbers. Result was: {result}"
        
        if result > 3999:  # Roman numerals traditionally don't go beyond this
            return f"Error: Result ({result}) is too large for standard Roman numeral representation"
        
        # Convert to Roman numeral
        roman_string = ""
        for value, numeral in roman_numerals:
            count = result // value
            if count:
                roman_string += numeral * count
                result -= value * count
        
        return f"The result of {oper}ing {value1} and {value2} is: {roman_string}"
        
    except Exception as e:
        return f"Error in Roman calculator: {str(e)}"


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York', 'Europe/London').
    
    Returns:
        The current time in the specified timezone.
    """
    try:
        if not timezone or not timezone.strip():
            return "Error: Timezone cannot be empty"
        
        # Clean the timezone string
        timezone = timezone.strip()
        
        # Handle common timezone aliases
        timezone_aliases = {
            'EST': 'America/New_York',
            'PST': 'America/Los_Angeles',
            'MST': 'America/Denver',
            'CST': 'America/Chicago',
            'GMT': 'GMT',
            'UTC': 'UTC',
            'CET': 'Europe/Berlin',
            'JST': 'Asia/Tokyo',
        }
        
        if timezone.upper() in timezone_aliases:
            timezone = timezone_aliases[timezone.upper()]
        
        # Create timezone object
        tz = pytz.timezone(timezone)
        
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz)
        formatted_time = local_time.strftime("%Y-%m-%d %H:%M:%S %Z")
        
        return f"The current local time in {timezone} is: {formatted_time}"
        
    except pytz.exceptions.UnknownTimeZoneError:
        return f"Error: Unknown timezone '{timezone}'. Please use a valid timezone like 'America/New_York' or 'Europe/London'"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


@tool
def factorial(n: int) -> Union[int, str]:
    """Calculates the factorial of a non-negative integer.
    
    Args:
        n: A non-negative integer.
        
    Returns:
        The factorial of n.
    """
    try:
        if not isinstance(n, int):
            return "Error: Input must be an integer"
        if n < 0:
            return "Error: Factorial is not defined for negative numbers"
        if n > 170:  # Prevent overflow
            return "Error: Number too large for factorial calculation"
        
        result = math.factorial(n)
        return result
    except Exception as e:
        return f"Error calculating factorial: {str(e)}"


@tool
def greatest_common_divisor(a: int, b: int) -> Union[int, str]:
    """Finds the greatest common divisor of two integers.
    
    Args:
        a: First integer.
        b: Second integer.
        
    Returns:
        The greatest common divisor of a and b.
    """
    try:
        if not isinstance(a, int) or not isinstance(b, int):
            return "Error: Both inputs must be integers"
        return math.gcd(abs(a), abs(b))
    except Exception as e:
        return f"Error calculating GCD: {str(e)}"


@tool
def least_common_multiple(a: int, b: int) -> Union[int, str]:
    """Finds the least common multiple of two integers.
    
    Args:
        a: First integer.
        b: Second integer.
        
    Returns:
        The least common multiple of a and b.
    """
    try:
        if not isinstance(a, int) or not isinstance(b, int):
            return "Error: Both inputs must be integers"
        if a == 0 or b == 0:
            return 0
        return abs(a * b) // math.gcd(abs(a), abs(b))
    except Exception as e:
        return f"Error calculating LCM: {str(e)}"


@tool
def is_prime(n: int) -> Union[bool, str]:
    """Checks if a number is prime.
    
    Args:
        n: The number to check.
        
    Returns:
        True if n is prime, False otherwise.
    """
    try:
        if not isinstance(n, int):
            return "Error: Input must be an integer"
        if n < 2:
            return False
        if n == 2:
            return True
        if n % 2 == 0:
            return False
        
        # Check odd divisors up to sqrt(n)
        for i in range(3, int(math.sqrt(n)) + 1, 2):
            if n % i == 0:
                return False
        return True
    except Exception as e:
        return f"Error checking if prime: {str(e)}"


@tool
def percentage_calculator(part: Union[int, float], whole: Union[int, float]) -> Union[float, str]:
    """Calculates what percentage 'part' is of 'whole'.
    
    Args:
        part: The part value.
        whole: The whole value.
        
    Returns:
        The percentage as a float.
    """
    try:
        if whole == 0:
            return "Error: Cannot calculate percentage when whole is zero"
        percentage = (part / whole) * 100
        return round(percentage, 2)
    except Exception as e:
        return f"Error calculating percentage: {str(e)}"


@tool
def compound_interest(principal: Union[int, float], rate: Union[int, float], 
                     time: Union[int, float], compound_frequency: int = 1) -> Union[float, str]:
    """Calculates compound interest.
    
    Args:
        principal: The initial amount of money.
        rate: The annual interest rate (as a percentage, e.g., 5 for 5%).
        time: The time period in years.
        compound_frequency: How many times per year the interest is compounded (default: 1).
        
    Returns:
        The final amount after compound interest.
    """
    try:
        if principal <= 0:
            return "Error: Principal must be positive"
        if rate < 0:
            return "Error: Interest rate cannot be negative"
        if time < 0:
            return "Error: Time cannot be negative"
        if compound_frequency <= 0:
            return "Error: Compound frequency must be positive"
        
        # Convert percentage to decimal
        rate_decimal = rate / 100
        
        # Compound interest formula: A = P(1 + r/n)^(nt)
        amount = principal * (1 + rate_decimal / compound_frequency) ** (compound_frequency * time)
        
        return round(amount, 2)
    except Exception as e:
        return f"Error calculating compound interest: {str(e)}"


@tool
def convert_temperature(value: Union[int, float], from_unit: str, to_unit: str) -> Union[float, str]:
    """Converts temperature between Celsius, Fahrenheit, and Kelvin.
    
    Args:
        value: The temperature value to convert.
        from_unit: The source unit ('C', 'F', or 'K').
        to_unit: The target unit ('C', 'F', or 'K').
        
    Returns:
        The converted temperature value.
    """
    try:
        from_unit = from_unit.upper().strip()
        to_unit = to_unit.upper().strip()
        
        valid_units = ['C', 'F', 'K', 'CELSIUS', 'FAHRENHEIT', 'KELVIN']
        
        # Normalize unit names
        unit_map = {
            'CELSIUS': 'C', 'FAHRENHEIT': 'F', 'KELVIN': 'K'
        }
        
        from_unit = unit_map.get(from_unit, from_unit)
        to_unit = unit_map.get(to_unit, to_unit)
        
        if from_unit not in ['C', 'F', 'K'] or to_unit not in ['C', 'F', 'K']:
            return "Error: Units must be 'C' (Celsius), 'F' (Fahrenheit), or 'K' (Kelvin)"
        
        if from_unit == to_unit:
            return float(value)
        
        # Convert to Celsius first
        if from_unit == 'F':
            celsius = (value - 32) * 5/9
        elif from_unit == 'K':
            celsius = value - 273.15
        else:  # from_unit == 'C'
            celsius = value
        
        # Convert from Celsius to target unit
        if to_unit == 'F':
            result = celsius * 9/5 + 32
        elif to_unit == 'K':
            result = celsius + 273.15
        else:  # to_unit == 'C'
            result = celsius
        
        return round(result, 2)
    except Exception as e:
        return f"Error converting temperature: {str(e)}"

@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia for a query and return maximum 2 results.
    
    Args:
        query: The search query.
    """
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"wiki_results": formatted_search_docs}

@tool
def arvix_search(query: str) -> str:
    """
    Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query.
    """
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arvix_results": formatted_search_docs}


@tool
def analyze_excel_file(file_path: str, query: str) -> str:
    """
    Analyze an Excel file using pandas and answer a question about it.
    Args:
        file_path (str): the path to the Excel file.
        query (str): Question about the data
    """
    try:
        file = pandas.read_excel(file_path)

        result = (
            f"Excel file loaded with {len(file)} rows and {len(file.columns)} columns.\n"
        )
        result += f"Columns: {', '.join(file.columns)}\n\n"

        result += "Summary statistics:\n"
        result += str(file.describe())

        return result

    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"
    

@tool
def python_code_parser(file_path: str) -> str:
    """
    Parse Python code to extract function names and their docstrings.

    Args:
        file_path: The path to the Python file to parse.

    Returns:
        Interpreted Python code as a string.
    """
    if not os.path.exists(file_path):
        return "0"


    loader = GenericLoader.from_filesystem(
        file_path,
        glob="**/*",
        suffixes=[".py"],
        parser=LanguageParser()
    )
    search_docs = loader.load()

    formatted_search_docs = "\n\n---\n\n".join(
    [
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
        for doc in search_docs
    ])

    return {"audio_results": formatted_search_docs}


@tool
def audio_transcription(file_path: str) -> str:
    """
    Transcribe an audio file to text using AssemblyAI.
    
    Args:
        file_path: The path to the audio file.
    
    Returns:
        The transcribed text from the audio file.
    """
    search_docs = AssemblyAIAudioTranscriptLoader(file_path=file_path).load()

    formatted_search_docs = "\n\n---\n\n".join(
    [
        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
        for doc in search_docs
    ])

    return {"audio_results": formatted_search_docs}


@tool
def analyze_csv_file(file_path: str, query: str) -> str:
    """
    Analyze a CSV file using pandas and answer a question about it.
    Args:
        file_path (str): the path to the CSV file.
        query (str): Question about the data
    """
    try:
        file = pandas.read_csv(file_path)

        result = f"CSV file loaded with {len(file)} rows and {len(file.columns)} columns.\n"
        result += f"Columns: {', '.join(file.columns)}\n\n"

        result += "Summary statistics:\n"
        result += str(file.describe())

        return result

    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"


# Extract Text Tool
vision_llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro",temperature=0)

@tool
def extract_text_from_image(img_path: str) -> str:
    """
    Extract text from an image file using a multimodal model.

    Args:
        img_path: A local image file path (strings).

    Returns:
        A single string containing the concatenated text extracted from each image.
    """
    all_text = ""
    try:

        # Read image and encode as base64
        with open(img_path, "rb") as image_file:
            image_bytes = image_file.read()

        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        # Prepare the prompt including the base64 image data
        message = [
            HumanMessage(
                content=[
                    {
                        "type": "text",
                        "text": (
                            "Extract all the text from this image. "
                            "Return only the extracted text, no explanations."
                        ),
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{image_base64}"
                        },
                    },
                ]
            )
        ]

        # Call the vision-capable model
        response = vision_llm.invoke(message)

        # Append extracted text
        all_text += response.content + "\n\n"

        return all_text.strip()
    except Exception as e:
        # You can choose whether to raise or just return an empty string / error message
        error_msg = f"Error extracting text: {str(e)}"
        print(error_msg)
        return ""


@tool
def reverse_sentence(text: str) -> str:
    """
    Reverses the input text.
    Args:
        text (str): The input string to be reversed.
    Returns:
        str: The reversed string.
    """
    return text[::-1]

@tool
def web_search(query: str) -> str:
    """
    Searches the web and returns a list of the most relevant URLs.
    Use this FIRST for complex queries, metadata questions, or to find the right sources.
    Then follow up with web_content_extract on the most promising URL.
    """
    try:
        tavily_search = TavilySearch(
            max_results=5,
            topic="general",
            search_depth="advanced",
            include_raw_content=False,  # Just URLs and snippets
        )

        results = tavily_search.invoke(query)
        # Format results to show URLs and brief descriptions
        web_search_results = "Search Results:\n"
        for i, result in enumerate(results["results"], 1):
            web_search_results += f"{i}. {result['title']}: {result['url']}\n   {result['content'][:150]}...\n\n"

        return web_search_results
    except Exception as e:
        return f"web_search tool error: {str(e)}"

@tool
def web_content_extract(url: str) -> str:
    """
    Extracts and analyzes specific content from a URL using BeautifulSoup.
    Particularly effective for Wikipedia metadata pages, discussion pages, 
    and structured web content.
    Can be used after web_search to get detailed information.
    """
    try:

        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
        }

        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()  # Raise exception for 4XX/5XX responses

        soup = BeautifulSoup(response.text, 'html.parser')
        for element in soup.select('script, style, footer, nav, header'):
            if element:
                element.decompose()          
        text = soup.body.get_text(separator='\n', strip=True) if soup.body else soup.get_text(separator='\n', strip=True)

        # Limit content length for response
        return f"Content extracted from {url}:\n\n{text[:10000]}..." if len(text) > 10000 else text

    except Exception as e:
        return f"web_content_extract tool error: {str(e)}"