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

2026-05-27 08:04:18 +00:00
parent 50bf3f76f7
commit aba250b750
1 changed files with 97 additions and 108 deletions
@@ -1,115 +1,104 @@
-import json
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Optional
+<|channel|>final code<|message|>import json
+from typing import Dict, List, Optional

 from fastmcp import FastMCP
+from fastmcp.exceptions import RPCError
+from pydantic import BaseModel


-class MemoryServer:
-    def __init__(self):
-        self.mcp = FastMCP("Memory-Server")
-        self.storage_path = Path("./memory_data.json")
-
-    # ------------------------------------------------------------------
-    # Вспомогательные методы для работы с файлом
-    # ------------------------------------------------------------------
-    def _load_memory(self) -> dict[str, Any]:
-        """Загружает память из JSON-файла."""
-        if not self.storage_path.exists():
-            return {}
-        with open(self.storage_path, "r", encoding="utf-8") as f:
-            try:
-                data = json.load(f)
-                if isinstance(data, dict):
-                    return data
-            except json.JSONDecodeError:
-                pass
-        return {}
-
-    def _save_memory(self, data: dict[str, Any]) -> None:
-        """Сохраняет память в JSON-файл."""
-        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)
-
-    # ------------------------------------------------------------------
-    # Инструменты без namespace
-    # ------------------------------------------------------------------
-    @self.mcp.tool()
-    def save(key: str, value: Any) -> bool:
-        """Сохраняет значение по ключу в память сервера."""
-        if not isinstance(key, str) or not key.strip():
-            return False
-        data = self._load_memory()
-        data[key] = {"value": value, "timestamp": datetime.utcnow().isoformat()}
-        self._save_memory(data)
-        return True
-
-    @self.mcp.tool()
-    def get(key: str) -> Optional[dict]:
-        """Возвращает значение по ключу с метаданными."""
-        if not isinstance(key, str):
-            return None
-        data = self._load_memory()
-        entry = data.get(key)
-        if entry is None:
-            return None
-        return {"key": key, "value": entry["value"], "timestamp": entry["timestamp"]}
-
-    @self.mcp.tool()
-    def delete(key: str) -> bool:
-        """Удаляет ключ из памяти сервера."""
-        if not isinstance(key, str):
-            return False
-        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(pattern: str = "*") -> list[str]:
-        """Возвращает список всех ключей с поддержкой wildcard-паттерна."""
-        if not isinstance(pattern, str):
-            pattern = "*"
-        data = self._load_memory()
-        import fnmatch
-
-        return [k for k in data.keys() if fnmatch.fnmatch(k, pattern)]
-
-    # ------------------------------------------------------------------
-    # Инструменты с namespace
-    # ------------------------------------------------------------------
-    @self.mcp.tool()
-    def save_with_namespace(key: str, value: Any, namespace: str = "default") -> bool:
-        """Сохраняет значение с указанием пространства имён."""
-        if not isinstance(namespace, str) or not namespace.strip():
-            return False
-        full_key = f"{namespace}:{key}"
-        return self.save(full_key, value)
-
-    @self.mcp.tool()
-    def get_by_namespace(namespace: str = "default") -> list[dict]:
-        """Возвращает все ключи из указанного namespace."""
-        if not isinstance(namespace, str) or not namespace.strip():
-            return []
-        data = self._load_memory()
-        prefix = f"{namespace}:"
-        result = []
-        for k, v in data.items():
-            if k.startswith(prefix):
-                key_part = k[len(prefix) :]
-                result.append(
-                    {"key": key_part, "value": v["value"], "timestamp": v["timestamp"]}
-                )
-        return result
+class SaveRequest(BaseModel):
+    key: str
+    value: str
+
+
+class GetRequest(BaseModel):
+    key: str
+
+
+class DeleteRequest(BaseModel):
+    key: str
+
+
+class ListKeysResponse(BaseModel):
+    keys: List[str]
+
+
+class NamespaceSaveRequest(BaseModel):
+    namespace: str
+    key: str
+    value: str
+
+
+class NamespaceGetRequest(BaseModel):
+    namespace: str
+    key: str
+
+
+# In‑memory store.  Keys are strings, values are arbitrary JSON serialisable objects.
+_memory_store: Dict[str, any] = {}
+# Namespaces are dictionaries inside the main store.
+_namespaces: Dict[str, Dict[str, any]] = {}
+
+
+def _ensure_namespace(namespace: str):
+    if namespace not in _namespaces:
+        _namespaces[namespace] = {}
+
+
+app = FastMCP()
+
+
+@app.rpc("save")
+async def save(request: SaveRequest) -> None:
+    """Сохраняет значение по ключу."""
+    _memory_store[request.key] = request.value
+
+
+@app.rpc("get")
+async def get(request: GetRequest):
+    """Возвращает сохранённое значение. Если ключа нет – ошибка."""
+    if request.key not in _memory_store:
+        raise RPCError(f"Key '{request.key}' not found")
+    return _memory_store[request.key]
+
+
+@app.rpc("delete")
+async def delete(request: DeleteRequest) -> None:
+    """Удаляет ключ из памяти."""
+    if request.key in _memory_store:
+        del _memory_store[request.key]
+    else:
+        raise RPCError(f"Key '{request.key}' not found")
+
+
+@app.rpc("list_keys")
+async def list_keys() -> ListKeysResponse:
+    """Возвращает список всех ключей."""
+    return ListKeysResponse(keys=list(_memory_store.keys()))
+
+
+@app.rpc("save_with_namespace")
+async def save_with_namespace(request: NamespaceSaveRequest) -> None:
+    """Сохраняет значение в указанном пространстве имён."""
+    _ensure_namespace(request.namespace)
+    _namespaces[request.namespace][request.key] = request.value
+
+
+@app.rpc("get_by_namespace")
+async def get_by_namespace(request: NamespaceGetRequest):
+    """Получает значение из пространства имён. Ошибка, если не найдено."""
+    if request.namespace not in _namespaces:
+        raise RPCError(f"Namespace '{request.namespace}' does not exist")
+    ns = _namespaces[request.namespace]
+    if request.key not in ns:
+        raise RPCError(
+            f"Key '{request.key}' not found in namespace '{request.namespace}'"
+        )
+    return ns[request.key]


-# ----------------------------------------------------------------------
-# Запуск сервера
-# ----------------------------------------------------------------------
 if __name__ == "__main__":
-    server = MemoryServer()
-    server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
+    # Запускаем сервер FastMCP
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)