Update app.py

app.py

@@ -4,6 +4,8 @@ from smolagents import LiteLLMModel, CodeAgent, ToolCallingAgent, Tool
 import wikipedia
 import gradio as gr
 import pandas as pd
+import json
+from datetime import datetime
 
 # Optional: Hugging Face token (for private models)
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -25,6 +27,59 @@ class WebSearchTool(Tool):
         tool = DuckDuckGoSearchTool()
         return tool.forward(query)
 
+class MemoryTool(Tool):
+    name = "memory_store"
+    description = "Store and retrieve persistent agent memory."
+    inputs = {
+        "action": {
+            "type": "string",
+            "description": "Either 'write' or 'read'"
+        },
+        "key": {
+            "type": "string",
+            "description": "Memory key"
+        },
+        "value": {
+            "type": "string",
+            "description": "Memory value (required for write)"
+        }
+    }
+    output_type = "string"
+
+    def __init__(self, memory_path: str = "/app/memory.json"):
+        self.memory_path = memory_path
+        if not os.path.exists(self.memory_path):
+            with open(self.memory_path, "w") as f:
+                json.dump([], f)
+
+    def forward(self, action: str, key: str, value: str = "") -> str:
+        try:
+            with open(self.memory_path, "r") as f:
+                memory = json.load(f)
+
+            if action == "write":
+                memory.append({
+                    "timestamp": datetime.utcnow().isoformat(),
+                    "key": key,
+                    "value": value
+                })
+                with open(self.memory_path, "w") as f:
+                    json.dump(memory, f, indent=2)
+                return "Memory stored successfully."
+
+            elif action == "read":
+                results = [m for m in memory if m["key"] == key]
+                if not results:
+                    return "No memory found for this key."
+                return json.dumps(results, indent=2)
+
+            else:
+                return "Invalid action. Use 'write' or 'read'."
+
+        except Exception as e:
+            return f"Memory error: {str(e)}"
+
+
 class LoadCsvTool(Tool):
     name = "load_csv"
     description = "Load and analyze CSV file. Returns summary statistics and first few rows. Input: file path."
@@ -123,7 +178,7 @@ You can send JSON payloads to external webhook URLs using the webhook_post tool.
 
 # --- Initialize Tool-Calling Agent ---
 agent = ToolCallingAgent(
-    tools=[WebSearchTool(), WikipediaTool(), WeatherTool(), LoadCsvTool(), WebhookPostTool()],
+    tools=[WebSearchTool(), WikipediaTool(), WeatherTool(), LoadCsvTool(), WebhookPostTool(), MemoryTool()],
     model=model,
     max_steps=10,
 )