# Knowledge retrieval tool using BM25
from langchain_community.retrievers import BM25Retriever
from smolagents import Tool

class GaiaRetrieverTool(Tool):
    name = "gaia_retriever"
    description = "Semantic search for retrieving relevant information for GaiaAgent."
    inputs = {
        "query": {
            "type": "string",
            "description": "Query for semantic search."
        }
    }
    output_type = "string"

    def __init__(self, docs, **kwargs):
        super().__init__(**kwargs)
        self.retriever = BM25Retriever.from_documents(docs, k=3)
        self.docs = docs  # Store docs for fallback

    def forward(self, query: str) -> str:
        assert isinstance(query, str), "Query must be a string."
        try:
            docs = self.retriever.invoke(query)
            if not docs:
                return "\nNo specific information found. Here's some general knowledge:\n" + "".join([
                    f"\n- {self.docs[i].page_content}" for i in range(min(3, len(self.docs)))
                ])
            return "\nRetrieved Information:\n" + "".join([
                f"\n- {doc.page_content}" for doc in docs
            ])
        except Exception as e:
            print(f"Error in retriever: {str(e)}")
            return f"Unable to retrieve specific information. The agent will rely on its general knowledge."