fix: main.py — MCP-сервер для управления памятью агента

main.py

@@ -1,14 +1,15 @@
 import os
 import asyncio
-from typing import Any, Dict
+from typing import Any
 
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
-from langchain.tools import tool
+from fastmcp import Client
 from deepagents import create_deep_agent
 from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
+from langchain_openai import ChatOpenAI
+from langchain.tools import tool
 
-# LLM via OpenRouter
+
+# LLM configuration (OpenRouter)
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -16,6 +17,7 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
+
 backend = CompositeBackend(
     [
         LocalShellBackend(workspace_dir="./workspace"),
@@ -23,48 +25,58 @@ backend = CompositeBackend(
     ]
 )
 
-# Tool that forwards calls to the MCP memory server
+
 @tool
-def memory_action(action: str, params: Dict[str, Any]) -> str:
-    """
-    Perform a memory operation via the MCP server.
-
-    action: one of save_with_namespace, get_by_namespace, list_keys, delete, get, save
-    params: dictionary of parameters required by the chosen action.
-    Returns a JSON string with the server response.
-    """
-    import json
-    import asyncio
-    from fastmcp import Client
-
+def memory_save(key: str, value: Any, namespace: str = "default") -> bool:
+    """Save a value in the remote memory server."""
     async def _call():
         client = Client("python memory_server.py")
         await client.connect()
         try:
-            result = await client.call_tool(action, params)
-            return json.dumps(result, ensure_ascii=False)
+            result = await client.call_tool(
+                "save_with_namespace",
+                {"key": key, "value": value, "namespace": namespace},
+            )
+            return result
         finally:
             await client.close()
+    return asyncio.run(_call())
 
-    loop = asyncio.get_event_loop()
-    return loop.run_until_complete(_call())
+
+@tool
+def memory_get(namespace: str = "default") -> str:
+    """Retrieve all key-value pairs from a namespace as a formatted string."""
+    async def _call():
+        client = Client("python memory_server.py")
+        await client.connect()
+        try:
+            data = await client.call_tool("get_by_namespace", {"namespace": namespace})
+            if not data:
+                return "No data."
+            lines = [f"{item['key']}: {item['value']}" for item in data]
+            return "\\n".join(lines)
+        finally:
+            await client.close()
+    return asyncio.run(_call())
 
 
 agent = create_deep_agent(
     model=llm,
-    tools=[memory_action],
+    tools=[memory_save, memory_get],
     backend=backend,
-    system_prompt="You are an assistant that can store and retrieve data using a remote memory server.",
+    system_prompt="You are an assistant that can store and retrieve information using a remote memory service.",
 )
 
+
 async def demo():
-    # Example: store a name and then retrieve it
-    query = (
-        "Save the user name 'Алексей' in the default namespace using the memory_action tool. "
-        "Then read back all entries from the default namespace and return them."
+    # Store a fact
+    await agent.ainvoke(
+        {"messages": [{"role": "user", "content": "Запомни, что мой любимый цвет - синий."}]},
+        {"configurable": {"thread_id": "demo-1"}},
     )
+    # Retrieve stored facts
     result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=query)]},
+        {"messages": [{"role": "user", "content": "Что я просил запомнить?"}]},
         {"configurable": {"thread_id": "demo-1"}},
     )
     print(result["messages"][-1].content)