Spaces:
Sleeping
Sleeping
File size: 917 Bytes
8488c37 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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) |