File size: 2,391 Bytes
beee8e9 634d586 beee8e9 99f2881 beee8e9 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
from langchain.docstore.document import Document
class VegetableInfoRetrieverTool(Tool):
name = "vegetable_info_retriever"
description = "Retrieve information about vegetables based on their name or characteristics."
inputs = {
"query": {
"type": "string",
"description": "The name or description of the vegetable 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 vegetable information found."
def load_vegetable_dataset():
# Manually define the vegetable data
vegetable_data = [
{
"name": "Sweet Potatoes",
"type": "root",
"description": "A starchy root vegetable with sweet flavor, rich in fiber and beta-carotene."
},
{
"name": "Fresh Basil",
"type": "herb",
"description": "An aromatic green herb often used fresh in Mediterranean cuisine."
},
{
"name": "Broccoli",
"type": "cruciferous",
"description": "Green vegetable with edible flowering heads and stalks, high in fiber and vitamin C."
},
{
"name": "Celery",
"type": "stem",
"description": "Crisp green stalks often eaten raw or cooked, known for high water content and fiber."
},
{
"name": "Lettuce",
"type": "leaf",
"description": "Leafy green used in salads, comes in various varieties like romaine or butterhead."
},
]
# Convert each entry to a Document
docs = [
Document(
page_content="\n".join([
f"Name: {veg['name']}",
f"Type: {veg['type']}",
f"Description: {veg['description']}"
]),
metadata={"name": veg["name"]}
)
for veg in vegetable_data
]
# Return the tool
return VegetableInfoRetrieverTool(docs)
|