src/tools.py

import json
import os

from langchain_community.document_loaders import \
    AssemblyAIAudioTranscriptLoader, WikipediaLoader, YoutubeLoader
from langchain_core.tools import tool


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.

    Args:
        a: first int
        b: second int
    """
    return a * b


@tool
def add(a: int, b: int) -> int:
    """
    Add two numbers.

    Args:
        a: first int
        b: second int
    """
    return a + b


@tool
def subtract(a: int, b: int) -> int:
    """
    Subtract two numbers.

    Args:
        a: first int
        b: second int
    """
    return a - b


@tool
def divide(a: int, b: int) -> int:
    """
    Divide two numbers.

    Args:
        a: numerator
        b: denominator
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b


@tool
def modulus(a: int, b: int) -> int:
    """
    Get the modulus of two numbers.

    Args:
        a: first int
        b: second int
    """
    return a % b


@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia for a query and return maximum 2 results. The results
    are formatted as a JSON string with the source and content of each
    document.

    Args:
        query: The search query."""
    search_docs = WikipediaLoader(
        query=query,
        load_max_docs=2,
    ).load()
    formatted_search_docs = [
        {
            "source": doc.metadata["source"],
            "content": doc.page_content,
        }
        for doc in search_docs
    ]
    return json.dumps({"wiki_results": formatted_search_docs})


@tool
def audio_transcript(
    audiofile_path: str
) -> str:
    """
    Transcribe an audio file to text.

    Args:
        audiofile_path: The path to the audio file to transcribe.
    """
    loader = AssemblyAIAudioTranscriptLoader(
        file_path=os.path.join('data', audiofile_path),
        api_key=os.environ.get("ASSEMBLY_API_KEY"),
    )
    docs = loader.load()
    return docs[0].page_content if docs else "No transcription available."


@tool
def youtube_transcript(
    video_url: str
) -> str:
    """
    Get the transcript of a YouTube video.

    Args:
        video_url: The URL of the YouTube video.
    """
    loader = YoutubeLoader.from_youtube_url(
        youtube_url=video_url,
        add_video_info=False
    )
    docs = loader.load()
    return docs[0].page_content if docs else "No transcript available."