basic_agent.py

from tools import (
    all_tools,
    analyze_youtube_video,
    extract_number_from_text,
    search_wikipedia,
    check_commutativity_table
)

class BasicAgent:
    def __init__(self):
        self.tool_registry = {
            "analyze_youtube_video": (analyze_youtube_video, "url"),
            "extract_number_from_text": (extract_number_from_text, "text"),
            "search_wikipedia": (search_wikipedia, "query"),
            "check_commutativity_table": (check_commutativity_table, "table_text"),
        }

    def __call__(self, question: str) -> str:
        try:
            tool, input_key = self.select_tool(question)
            return tool.invoke({input_key: question})
        except Exception as e:
            return f"Error: {e}"

    def select_tool(self, question: str):
        q = question.lower()

        if "youtube.com" in q or "youtu.be" in q or "in the video" in q:
            return self.tool_registry["analyze_youtube_video"]

        elif "how many" in q or "number of" in q:
            return self.tool_registry["extract_number_from_text"]

        elif "wikipedia" in q or "encyclopedia" in q or "who is" in q:
            return self.tool_registry["search_wikipedia"]

        elif "commutative" in q or "counter-examples" in q:
            return self.tool_registry["check_commutativity_table"]

        return self.tool_registry["extract_number_from_text"]  # fallback