ATK20 commited on
Commit
8488c37
·
verified ·
1 Parent(s): 8ff4014

Create Create the Retriever Tool

Browse files
Files changed (1) hide show
  1. Create the Retriever Tool +27 -0
Create the Retriever Tool ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from smolagents import Tool
2
+ from langchain_community.retrievers import BM25Retriever
3
+
4
+ class GuestInfoRetrieverTool(Tool):
5
+ name = "guest_info_retriever"
6
+ description = "Retrieves detailed information about gala guests based on their name or relation."
7
+ inputs = {
8
+ "query": {
9
+ "type": "string",
10
+ "description": "The name or relation of the guest you want information about."
11
+ }
12
+ }
13
+ output_type = "string"
14
+
15
+ def __init__(self, docs):
16
+ self.is_initialized = False
17
+ self.retriever = BM25Retriever.from_documents(docs)
18
+
19
+ def forward(self, query: str):
20
+ results = self.retriever.get_relevant_documents(query)
21
+ if results:
22
+ return "\n\n".join([doc.page_content for doc in results[:3]])
23
+ else:
24
+ return "No matching guest information found."
25
+
26
+ # Initialize the tool
27
+ guest_info_tool = GuestInfoRetrieverTool(docs)