tools.py

import os
import ast
import io
import sys
import numpy as np
import pandas as pd
import scipy

from pathlib import Path
import mimetypes
import base64

from google import genai

ALLOWED_MODULES = {"numpy", "pandas", "scipy"}
GEMINI_API_KEY = os.getenv("GEMINI_TOKEN")
GEMINI_MODEL_NAME = "gemini-2.0-flash"

def interpret_python_math_code(python_code: str) -> str:
    """
    Interprets a string of Python code to perform math calculations.

    Security Note: This function uses exec(). While it attempts to restrict
    imports to numpy, pandas, and scipy, and runs with a restricted
    global scope, executing arbitrary code always carries risks. Ensure
    that input code is from a trusted source or properly sanitized.

    The code must only import modules from the allowed list: numpy, pandas, scipy.
    Submodules of these (e.g., numpy.linalg, scipy.stats) are permitted.
    For example:
    'import numpy as np' is allowed.
    'from scipy.stats import norm' is allowed.
    'import os' is NOT allowed.

    To return a result, the code should either:
    1. End with an expression (e.g., '1 + 1' or 'np.array([1,2,3]).sum()').
    2. Assign the result to a variable named '_result' (e.g., '_result = my_calculation').
    
    Print statements will also be captured and returned along with the result.
    """
    # 1. Validate imports using AST
    try:
        tree = ast.parse(python_code)
        for node in tree.body:
            if isinstance(node, ast.Import):
                for alias in node.names:
                    root_module = alias.name.split('.')[0]
                    if root_module not in ALLOWED_MODULES:
                        return (f"Error: Import of '{alias.name}' is not allowed. "
                                f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
            elif isinstance(node, ast.ImportFrom):
                if node.module:  # Handles cases like 'from . import something' where module is None
                    root_module = node.module.split('.')[0]
                    if root_module not in ALLOWED_MODULES:
                        return (f"Error: Import from '{node.module}' is not allowed. "
                                f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
    except SyntaxError as e:
        return f"Syntax Error in input code: {e}"

    # 2. Prepare execution environment
    restricted_globals = {
        "__builtins__": {
            "print": print,
            "abs": abs, "round": round, "min": min, "max": max, "sum": sum, "len": len,
            "range": range, "zip": zip, "enumerate": enumerate,
            "int": int, "float": float, "str": str, "list": list, "dict": dict, "tuple": tuple, "set": set,
            "True": True, "False": False, "None": None,
            "__import__": __import__, # Add this line
        }
        # numpy, pandas, scipy are NOT pre-loaded here.
        # The user's code `import numpy` will use Python's import mechanism.
        # The AST check above is the primary guard.
    }
    local_vars = {}

    # 3. Capture stdout
    old_stdout = sys.stdout
    redirected_output = io.StringIO()
    sys.stdout = redirected_output

    # 4. Execute code and retrieve result
    calculated_value = None
    result_source = ""
    output_str = ""

    try:
        compiled_code = compile(python_code, '<string>', 'exec')
        exec(compiled_code, restricted_globals, local_vars)
        
        # Priority 1: Check for '_result' variable
        if "_result" in local_vars:
            calculated_value = local_vars["_result"]
            result_source = "variable '_result'"
        # Priority 2: If no _result, and the last AST node was an expression, evaluate it.
        elif tree.body and isinstance(tree.body[-1], ast.Expr):
            # Ensure the expression node's value is a valid AST object for ast.Expression
            if isinstance(tree.body[-1].value, ast.AST):
                last_expr_ast = ast.Expression(body=tree.body[-1].value)
                # Compile the expression in 'eval' mode
                compiled_expr = compile(last_expr_ast, '<string>', 'eval')
                # Evaluate in the context of restricted_globals and local_vars (which holds state from exec)
                calculated_value = eval(compiled_expr, restricted_globals, local_vars)
                result_source = "last expression"
            
        sys.stdout = old_stdout # Restore stdout before getting its value
        output_str = redirected_output.getvalue()

        if calculated_value is not None:
            return f"Result (from {result_source}):\n{calculated_value}\n\nCaptured Output:\n{output_str}".strip()
        else:
            return f"Executed successfully.\n\nCaptured Output:\n{output_str}\n(No specific result value found via '_result' variable or last expression evaluation.)".strip()

    except Exception as e:
        if sys.stdout == redirected_output: # Ensure stdout is restored on error too
            sys.stdout = old_stdout
        output_str = redirected_output.getvalue() # Get any output captured before the error
        return f"Execution Error: {type(e).__name__}: {e}\n\nCaptured Output:\n{output_str}".strip()
    finally:
        # Ensure stdout is always restored
        if sys.stdout == redirected_output:
            sys.stdout = old_stdout


# STT tool
def convert_audio_to_text(path_to_audio: str) -> str:
    """
    Converts speech from an audio file into text.
    Args:
        path_to_audio (str): The path to the audio file to be transcribed.
    Returns:
        str: The transcribed text content of the audio file.
    """
    
    client = genai.Client(api_key="GOOGLE_API_KEY")

    myfile = client.files.upload(file=path_to_audio)

    transcription = client.models.generate_content(
        model=GEMINI_MODEL_NAME, contents=["Provide a transcription of this audio file.", myfile]
    )
    

    return transcription.text

# Analyze image tool
def image_understanding(path_to_image: str, question: str) -> str:
    """
    Analyzes an image and generates a response to a given question based on the image's content.
    
    Args:
        path_to_image (str): The path to the image file to be analyzed.
        question (str): The question to be answered, based on the contents of the image.
        
    Returns:
        str: The response from a VLM, typically a textual analysis or description based on the image.
    """

    client = genai.Client(api_key=GEMINI_API_KEY)

    my_file = client.files.upload(file=path_to_image)

    response = client.models.generate_content(
    model=GEMINI_MODEL_NAME,
    contents=[my_file, question],
    )

    return response.text

# Analyze video tool
def video_understanding(path_to_video: str, question: str) -> str:
    """
    Analyzes a video and generates a response to a given question based on the video's content.
    
    Args:
        path_to_video (str): The path to the video file to be analyzed.
        question (str): The question to be answered, based on the contents of the video.
        
    Returns:
        str: The response from a VLM, typically a textual analysis or description based on the video.
    """

    client = genai.Client(api_key=GEMINI_API_KEY)

    my_file = client.files.upload(file=path_to_video)

    response = client.models.generate_content(
    model=GEMINI_MODEL_NAME,
    contents=[my_file, question],
    )

    return response.text
    

## Read .csv file tool
def read_csv_file(path_to_csv: str) -> str:
    """
    Reads a CSV file from the specified path and returns its content as plain text.
    
    Args:
    path_to_csv (str): The file path to the CSV file.
    
    Returns:
    str: Content of the CSV file as plain text.
    """
    try:
        # Read the CSV file using pandas
        df = pd.read_csv(path_to_csv)
        
        # Return df as plain tect
        return df.to_string(index=False)
    except Exception as e:
        return f"Error reading the CSV file: {e}"

## Read .xlsx file tool
def read_xlsx_file(path_to_xlsx: str) -> str:
    """
    Reads a XLSX file from the specified path and returns its content as plain text.
    
    Args:
    path_to_xlsx (str): The file path to the XLSX file.
    
    Returns:
    str: Content of the XLSX file as plain text.
    """
    try:
        # Read the XLSX file using pandas
        df = pd.read_excel(path_to_xlsx)
        
        # Return df as plain tect
        return df.to_string(index=False)
    except Exception as e:
        return f"Error reading the XLSX file: {e}"