Alfred_RAG / tools /search_tool.py
chughtaihamad's picture
Main commit
b9744a5
raw
history blame contribute delete
746 Bytes
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]])