tools/tool_manager.py

"""
Tool manager for the GAIA Agent.

This handles the coordination between different tools and provides them to the agent.
"""

from smolagents import Tool
from typing import Dict, List, Any
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.retrievers import BM25Retriever
import functools

from config import load_knowledge_base
from tools.web_tools import WebSearchTool, WebContentTool
from tools.youtube_tool import YoutubeVideoTool
from tools.wikipedia_tool import WikipediaTool
from tools.knowledge_tool import GaiaRetrieverTool

class ToolManager:
    """
    Manages and initializes all available tools for the GAIA agent.
    Also provides tool selection logic based on question analysis.
    """

    def __init__(self):
        # Load and process knowledge base
        knowledge_text = load_knowledge_base()
        self.knowledge_docs = self._create_knowledge_documents(knowledge_text)
        
        # Initialize tools
        self.tools = self._initialize_tools()

    def _create_knowledge_documents(self, text: str) -> List[Document]:
        """Create searchable documents from knowledge base text."""
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
        )
        knowledge_chunks = text_splitter.split_text(text)
        return [Document(page_content=chunk) for chunk in knowledge_chunks]

    def _initialize_tools(self) -> List[Tool]:
        """Initialize all available tools."""
        return [
            GaiaRetrieverTool(self.knowledge_docs),
            WebSearchTool(),
            WebContentTool(),
            YoutubeVideoTool(),
            WikipediaTool(),
        ]

    def get_tools(self) -> List[Tool]:
        """Return all available tools."""
        return self.tools