add: main.py

main.py

@@ -0,0 +1,138 @@
+import asyncio
+import json
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional
+from fnmatch import fnmatch
+
+from fastmcp import FastMCP, Client
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from langchain.tools import tool
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+
+# ------------------- Memory Server -------------------
+class MemoryServer:
+    def __init__(self):
+        self.mcp = FastMCP("Memory-Server")
+        self.storage_path = Path("./memory_data.json")
+
+    def _load_memory(self) -> dict:
+        if not self.storage_path.exists():
+            return {}
+        with open(self.storage_path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+
+    def _save_memory(self, data: dict):
+        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.storage_path, 'w', encoding='utf-8') as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+
+    @self.mcp.tool()
+    def save(self, key: str, value: Any) -> bool:
+        data = self._load_memory()
+        data[key] = {
+            "value": value,
+            "timestamp": datetime.utcnow().isoformat()
+        }
+        self._save_memory(data)
+        return True
+
+    @self.mcp.tool()
+    def get(self, key: str) -> Optional[dict]:
+        data = self._load_memory()
+        if key in data:
+            return {"key": key, "value": data[key]["value"], "timestamp": data[key]["timestamp"]}
+        return None
+
+    @self.mcp.tool()
+    def delete(self, key: str) -> bool:
+        data = self._load_memory()
+        if key in data:
+            del data[key]
+            self._save_memory(data)
+            return True
+        return False
+
+    @self.mcp.tool()
+    def list_keys(self, pattern: str = "*") -> list[str]:
+        data = self._load_memory()
+        return [k for k in data.keys() if fnmatch(k, pattern)]
+
+    @self.mcp.tool()
+    def save_with_namespace(self, key: str, value: Any, namespace: str = "default") -> bool:
+        full_key = f"{namespace}:{key}"
+        return self.save(full_key, value)
+
+    @self.mcp.tool()
+    def get_by_namespace(self, namespace: str = "default") -> list[dict]:
+        data = self._load_memory()
+        result = []
+        prefix = f"{namespace}:"
+        for k, v in data.items():
+            if k.startswith(prefix):
+                result.append({"key": k, "value": v["value"], "timestamp": v["timestamp"]})
+        return result
+
+# ------------------- DeepAgent that uses Memory Server -------------------
+# The agent will use a simple tool that calls the memory server via MCP client
+@tool
+def memory_tool(action: str, params: dict) -> str:
+    """Calls the memory server tool specified by action.
+    Parameters must include the tool name and its arguments.
+    """
+    client = Client("python memory_server.py")
+    asyncio.run(client.connect())
+    try:
+        result = asyncio.run(client.call_tool(action, params))
+        return str(result)
+    finally:
+        asyncio.run(client.close())
+
+# LLM and backend setup
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
+)
+
+backend = CompositeBackend([
+    LocalShellBackend(workspace_dir="./workspace"),
+    FilesystemBackend(),
+])
+
+agent = create_deep_agent(
+    model=llm,
+    tools=[memory_tool],
+    backend=backend,
+    system_prompt="You are an agent that can store and retrieve data using the memory server.",
+)
+
+async def run_agent_demo():
+    response = await agent.ainvoke(
+        {"messages": [HumanMessage(content="Store key 'greeting' with value 'Hello' in default namespace.")]},
+        {"configurable": {"thread_id": "demo-1"}},
+    )
+    print("Agent response:", response["messages"][-1].content)
+
+# ------------------- Server run -------------------
+if __name__ == "__main__":
+    # Start server in background if called directly
+    server = MemoryServer()
+    # Run server in a separate thread so that the agent demo can also run
+    import threading
+    server_thread = threading.Thread(target=lambda: server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR"), daemon=True)
+    server_thread.start()
+    # Give server a moment to start
+    asyncio.run(asyncio.sleep(0.5))
+    # Run agent demo
+    asyncio.run(run_agent_demo())
+    # Keep server alive for a short while to allow manual testing
+    try:
+        while True:
+            pass
+    except KeyboardInterrupt:
+        print("Shutting down.")