AlexGrig23/Unit_3_Agentic_RAG
0
1from smolagents import Tool2from langchain_community.retrievers import BM25Retriever3from langchain.docstore.document import Document4import datasets5 6 7class GuestInfoRetrieverTool(Tool):8 name = "guest_info_retriever"9 description = "Retrieves detailed information about gala guests based on their name or relation."10 inputs = {11 "query": {12 "type": "string",13 "description": "The name or relation of the guest you want information about."14 }15 }16 output_type = "string"17 18 def __init__(self, docs):19 self.is_initialized = False20 self.retriever = BM25Retriever.from_documents(docs)21 22 23 def forward(self, query: str):24 results = self.retriever.get_relevant_documents(query)25 if results:26 return "\n\n".join([doc.page_content for doc in results[:3]])27 else:28 return "No matching guest information found."29 30 31def load_guest_dataset():32 # Load the dataset33 guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")34 35 # Convert dataset entries into Document objects36 docs = [37 Document(38 page_content="\n".join([39 f"Name: {guest['name']}",40 f"Relation: {guest['relation']}",41 f"Description: {guest['description']}",42 f"Email: {guest['email']}"43 ]),44 metadata={"name": guest["name"]}45 )46 for guest in guest_dataset47 ]48 49 # Return the tool50 return GuestInfoRetrieverTool(docs)51 52 53 54 