Create retriever.py

retriever.py

@@ -0,0 +1,87 @@
+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 Potato",
+            "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": "Green Beans",
+            "type": "legume",
+            "description": "Long green pods with small seeds inside, eaten whole and rich in vitamins."
+        },
+        {
+            "name": "Bell Pepper",
+            "type": "fruit vegetable",
+            "description": "Crunchy and colorful, bell peppers are technically fruits, but used as vegetables in cooking."
+        },
+        {
+            "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": "Zucchini",
+            "type": "squash",
+            "description": "Soft-skinned green vegetable, technically a fruit, often grilled or sautéed."
+        },
+        {
+            "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)