from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
from langchain.schema import Document
from typing import List


class SearchTool(Tool):
    name = "search"
    description = "Performs a general search over provided documents. Use `/search <query>` to call."
    inputs = {"query": {"type": "string", "description": "Search query"}}
    output_type = "string"

    def __init__(self, docs: List[Document]):
        self.retriever = BM25Retriever.from_documents(docs)

    def forward(self, query: str) -> str:
        results = self.retriever.get_relevant_documents(query)
        if not results:
            return "No results found."

        return "\n\n".join([f"- {r.page_content}" for r in results[:5]])