Unit_3_Agentic_RAG / Create the Retriever Tool
ATK20's picture
Create Create the Retriever Tool
8488c37 verified
raw
history blame contribute delete
917 Bytes
from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
class GuestInfoRetrieverTool(Tool):
name = "guest_info_retriever"
description = "Retrieves detailed information about gala guests based on their name or relation."
inputs = {
"query": {
"type": "string",
"description": "The name or relation of the guest you want information about."
}
}
output_type = "string"
def __init__(self, docs):
self.is_initialized = False
self.retriever = BM25Retriever.from_documents(docs)
def forward(self, query: str):
results = self.retriever.get_relevant_documents(query)
if results:
return "\n\n".join([doc.page_content for doc in results[:3]])
else:
return "No matching guest information found."
# Initialize the tool
guest_info_tool = GuestInfoRetrieverTool(docs)