agents/assistant_tools.py

# =============================================================================
# MEETUP DEMO - AwesomeAgent (LangGraph + LLM Tools)
#
# TABLE OF CONTENTS
# [1] Import Built-in LangChain tools
# [2] Create Custom Tools
# [3] Init Built-in LangChain tools
# [4] Return Complete list of tools
# =============================================================================

import base64
import math
import os
from typing import Optional

import pandas as pd
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI

# [1] Import Built-in LangChain tools
# ---

from langchain_community.tools import ShellTool
from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
from langchain_community.document_loaders.assemblyai import TranscriptFormat

# Youtube related tools
from agents.youtube_transcript_tool import (
    get_youtube_transcript_tool,
    get_youtube_title_description,
)

# Web search subagent
from agents.web_search_subagent import web_search_agent

load_dotenv()
vision_llm = ChatGoogleGenerativeAI(model=os.getenv("GOOGLE_VISION_MODEL"))

# ============== CUSTOM TOOLS (not available in LangChain) ==============

# [2] Create Custom Tools
# ---


@tool
def calculator(expression: str) -> str:
    """
    Perform mathematical calculations safely. Supports basic arithmetic operations.

    Example:
    Input: "10 + 2 * 3 - 4 / 2"
    Output: "14.0"
    """
    try:
        # Safe evaluation - only allow basic math operations
        allowed_chars = set("0123456789+-*/.() ")
        if all(c in allowed_chars for c in expression):
            result = eval(expression)
            return str(result)
        else:
            return "Error: Invalid characters in expression"
    except Exception as e:
        return f"Calculation error: {str(e)}"


@tool
def advanced_math(operation: str, num1: float, num2: Optional[float] = None) -> str:
    """
    Perform advanced math operations.
    Supported: sqrt, log, sin, cos, tan, power.

    Examples:
    Input: operation="sqrt", num1=16
    Output: "4.0"

    Input: operation="power", num1=2, num2=3
    Output: "8.0"
    """
    try:
        if operation == "sqrt":
            return str(math.sqrt(num1))
        elif operation == "log":
            return str(math.log(num1))
        elif operation == "sin":
            return str(math.sin(num1))
        elif operation == "cos":
            return str(math.cos(num1))
        elif operation == "tan":
            return str(math.tan(num1))
        elif operation == "power":
            if num2 is None:
                return "power operation requires two numbers"
            return str(math.pow(num1, num2))
        else:
            return f"Unknown operation: {operation}"
    except Exception as e:
        return f"Math error: {str(e)}"


@tool
def ask_question_on_image_content(question: str, img_path: str) -> str:
    """Ask specific question on image content."""
    try:
        if "vision_llm" not in globals():
            return "Error: Vision LLM not configured. Please uncomment and configure vision_llm."

        with open(img_path, "rb") as image_file:
            image_bytes = image_file.read()

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

        message = [
            HumanMessage(
                content=[
                    {
                        "type": "text",
                        "text": question,
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{image_base64}"},
                    },
                ]
            )
        ]

        response_content = vision_llm.invoke(message).content.strip()
        return response_content
    except Exception as e:
        return f"Multimodal text extraction error: {str(e)}"


@tool
def read_excel_file(file_path: str, sheet_name: Optional[str] = None) -> str:
    """
    Read an Excel file and return its content.
    """
    try:
        if sheet_name:
            json_str = pd.read_excel(file_path, sheet_name=sheet_name).to_json(
                orient="records", date_format="iso", force_ascii=False
            )
        else:
            json_str = pd.read_excel(file_path).to_json(
                orient="records", date_format="iso", force_ascii=False
            )
        return json_str
    except Exception as e:
        return f"Excel reading error: {str(e)}"


@tool
def transcribe_mp3(
    file_path: str,
) -> str:
    """
    Transcribe an MP3 (local path or URL) using AssemblyAI and return the result.
    """
    try:
        tfmt = getattr(TranscriptFormat, "TEXT", TranscriptFormat.TEXT)

        loader = AssemblyAIAudioTranscriptLoader(
            file_path=file_path,
            transcript_format=tfmt,
        )
        docs = loader.load()

        # Concatenate documents
        text = "\n\n".join(d.page_content for d in docs)
        MAX = 40000
        if len(text) > MAX:
            return text[:MAX] + f"\n\n...[truncated {len(text) - MAX} chars]..."
        return text
    except Exception as e:
        return f"Transcription error: {str(e)}"


@tool
def execute_python_file(file_path: str) -> str:
    """
    Execute Python code from the provided file path and return its stdout/stderr output.
    """
    import subprocess
    import sys

    if not isinstance(file_path, str) or not file_path.strip():
        return "Error: file_path must be a non-empty string."

    file_path = file_path.strip()

    if not os.path.isfile(file_path):
        return f"Error: file not found at '{file_path}'."

    if not file_path.endswith(".py"):
        return "Error: only '.py' files are supported."

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            f.read()
    except Exception as e:
        return f"Error reading file: {str(e)}"

    try:
        completed = subprocess.run(
            [sys.executable, file_path],
            capture_output=True,
            text=True,
            check=False,
        )
    except Exception as e:
        return f"Execution error: {str(e)}"

    stdout_content = (completed.stdout or "").strip()
    stderr_content = (completed.stderr or "").strip()

    if stdout_content and stderr_content:
        return f"STDOUT:\n{stdout_content}\n\nSTDERR:\n{stderr_content}"
    if stdout_content:
        return stdout_content
    if stderr_content:
        return f"STDERR:\n{stderr_content}"
    return "Execution completed with no output."


def build_tools():
    """
    Initialize and return a list of built-in and custom LangChain tools.
    """

    # [3] Init Built-in LangChain tools
    # ---

    # Initialize built-in LangChain tools

    # SECURITY: Shell Tool
    # Should implement:
    # - Command whitelist: Only allow safe commands (ls, cat, grep, find, etc.)
    # - Block dangerous commands: rm, sudo, chmod, kill, curl, wget, etc.
    # - Block command chaining: Prevent ; && || | operators
    # - Block redirections: Prevent > >> < operators
    # - Block subshells: Prevent $() and backticks
    # - Path restrictions: Limit file operations to specific directories
    # - Validate command before execution
    shell_tool = ShellTool()

    # [4] Return Complete list of tools
    # ---

    # Combine built-in tools with custom tools
    all_tools = [
        # Built-in LangChain tools
        shell_tool,
        # Web search subagent (replaces individual web/search tools for isolated context)
        web_search_agent,
        # Custom tools for specialized tasks
        execute_python_file,
        read_excel_file,
        get_youtube_transcript_tool,
        get_youtube_title_description,
        calculator,
        advanced_math,
        ask_question_on_image_content,
        transcribe_mp3,
    ]
    return all_tools


if __name__ == "__main__":
    from pprint import pprint

    print("\n--- calculator ---")
    pprint(calculator.invoke({"expression": "10 + 2 * 3 - 4 / 2"}))

    print("\n--- advanced_math ---")
    pprint(advanced_math.invoke({"operation": "sqrt", "num1": 16}))
    pprint(advanced_math.invoke({"operation": "power", "num1": 2, "num2": 3}))

    print("\n--- read_excel_file ---")
    pprint(read_excel_file.invoke({"file_path": "examples/file_example_XLS_10.xls"}))
    pprint(
        read_excel_file.invoke(
            {"file_path": "examples/file_example_XLS_10.xls", "sheet_name": "Sheet2"}
        )
    )

    print("\n--- ask_question_on_image_content ---")
    pprint(
        ask_question_on_image_content.invoke(
            {
                "question": "How many apples are in the basket?",
                "img_path": "examples/apples.png",
            }
        )
    )

    print("\n--- transcribe_mp3 ---")
    pprint(transcribe_mp3.invoke({"file_path": "./examples/sample.mp3"}))