diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..062c402 --- /dev/null +++ b/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()