Add agent.py

agent.py

@@ -0,0 +1,136 @@
+import os
+import json
+import datetime
+from pathlib import Path
+from dotenv import load_dotenv
+from fastmcp import FastMCP, mcp
+from pydantic import BaseModel, Field
+
+# Load environment variables
+load_dotenv()
+
+# Configuration
+DATA_FILE = Path(os.getenv("MEMORY_DATA_FILE", "memory.json"))
+
+# Pydantic model for stored value
+class StoredItem(BaseModel):
+    key: str = Field(..., description="The key for the item")
+    value: str = Field(..., description="The stored value")
+    timestamp: str = Field(..., description="ISO 8601 timestamp when stored")
+
+# In-memory cache
+_memory_cache: dict[str, StoredItem] = {}
+
+# Load persisted data on startup
+def load_data() -> None:
+    global _memory_cache
+    if DATA_FILE.exists():
+        try:
+            with DATA_FILE.open("r", encoding="utf-8") as f:
+                data = json.load(f)
+            for key, item in data.items():
+                try:
+                    _memory_cache[key] = StoredItem(**item)
+                except Exception as exc:
+                    print(f"[WARN] Failed to parse item {key}: {exc}")
+        except Exception as exc:
+            print(f"[ERROR] Unable to read data file {DATA_FILE}: {exc}")
+
+# Persist data to disk
+def persist_data() -> None:
+    try:
+        data = {k: v.dict() for k, v in _memory_cache.items()}
+        with DATA_FILE.open("w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2)
+    except Exception as exc:
+        print(f"[ERROR] Unable to write data file {DATA_FILE}: {exc}")
+
+# Tool definitions
+@mcp.tool("save")
+def save(key: str, value: str) -> str:
+    """Save a key-value pair with timestamp."""
+    try:
+        item = StoredItem(key=key, value=value, timestamp=datetime.datetime.utcnow().isoformat())
+        _memory_cache[key] = item
+        persist_data()
+        return f"Saved key '{key}'."
+    except Exception as exc:
+        return f"Error saving key '{key}': {exc}"
+
+@mcp.tool("get")
+def get(key: str) -> str:
+    """Retrieve value for a key."""
+    try:
+        item = _memory_cache[key]
+        return f"{item.value} (stored at {item.timestamp})"
+    except KeyError:
+        return f"Key '{key}' not found."
+    except Exception as exc:
+        return f"Error retrieving key '{key}': {exc}"
+
+@mcp.tool("delete")
+def delete(key: str) -> str:
+    """Delete a key."""
+    try:
+        del _memory_cache[key]
+        persist_data()
+        return f"Deleted key '{key}'."
+    except KeyError:
+        return f"Key '{key}' not found."
+    except Exception as exc:
+        return f"Error deleting key '{key}': {exc}"
+
+@mcp.tool("list_keys")
+def list_keys() -> str:
+    """List all keys."""
+    try:
+        keys = ", ".join(_memory_cache.keys()) or "No keys"
+        return f"Keys: {keys}"
+    except Exception as exc:
+        return f"Error listing keys: {exc}"
+
+# Namespace tools
+@mcp.tool("save_with_namespace")
+def save_with_namespace(namespace: str, key: str, value: str) -> str:
+    """Save a key-value pair under a namespace."""
+    namespaced_key = f"{namespace}:{key}"
+    return save(namespaced_key, value)
+
+@mcp.tool("get_by_namespace")
+def get_by_namespace(namespace: str, key: str) -> str:
+    """Retrieve a value from a namespace."""
+    namespaced_key = f"{namespace}:{key}"
+    return get(namespaced_key)
+
+# Initialize server
+mcpServer = FastMCP()
+# Register tools automatically via decorators
+
+# Load existing data
+load_data()
+
+# Demo client if invoked with --client
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) > 1 and sys.argv[1] == "--client":
+        # Simple interactive client using stdio transport
+        print("Running client demo. Type 'exit' to quit.")
+        print("Available tools: save, get, delete, list_keys, save_with_namespace, get_by_namespace")
+        while True:
+            try:
+                user_input = input("> ")
+            except EOFError:
+                break
+            if user_input.strip().lower() == "exit":
+                break
+            # Send command via mcp.run() with input string
+            try:
+                response = mcp.run(user_input)
+                print(response)
+            except Exception as e:
+                print(f"Error: {e}")
+    else:
+        # Run server with stdio transport
+        print("Starting FastMCP server with stdio transport...")
+        mcp.run()