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

solutions/69f8e929da860fb4533faa2a_MCP-сервер_для_управления_памятью_агента/memory_server.py

@@ -1,104 +1,109 @@
 import json
-from typing import Dict, List, Optional
+from typing import Any, Dict, List
 
-from fastmcp import FastMCP
-from fastmcp.exceptions import RPCError
-from pydantic import BaseModel
+from fastmcp import FastMCP, tool
 
 
-class SaveRequest(BaseModel):
-    key: str
-    value: str
+class MemoryStore:
+    """
+    Простое хранилище в памяти с поддержкой пространств имён.
+    """
 
+    def __init__(self) -> None:
+        # основной словарь: ключ → значение
+        self._store: Dict[str, Any] = {}
+        # словари по namespace: namespace → {key: value}
+        self._namespaces: Dict[str, Dict[str, Any]] = {}
 
-class GetRequest(BaseModel):
-    key: str
+    # ---------- Основные операции ----------
+    def save(self, key: str, value: Any) -> None:
+        """Сохраняет значение под ключом."""
+        self._store[key] = value
 
+    def get(self, key: str) -> Any | None:
+        """Возвращает значение по ключу или None, если его нет."""
+        return self._store.get(key)
 
-class DeleteRequest(BaseModel):
-    key: str
+    def delete(self, key: str) -> bool:
+        """Удаляет ключ. Возвращает True, если удалено."""
+        if key in self._store:
+            del self._store[key]
+            return True
+        return False
 
-
-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:
+    def list_keys(self) -> List[str]:
         """Возвращает список всех ключей."""
-    return ListKeysResponse(keys=list(_memory_store.keys()))
+        return list(self._store.keys())
 
-
-@app.rpc("save_with_namespace")
-async def save_with_namespace(request: NamespaceSaveRequest) -> None:
+    # ---------- Операции с namespace ----------
+    def save_with_namespace(self, namespace: str, key: str, value: Any) -> None:
         """Сохраняет значение в указанном пространстве имён."""
-    _ensure_namespace(request.namespace)
-    _namespaces[request.namespace][request.key] = request.value
+        ns = self._namespaces.setdefault(namespace, {})
+        ns[key] = value
+
+    def get_by_namespace(self, namespace: str, key: str) -> Any | None:
+        """Получает значение из пространства имён."""
+        return self._namespaces.get(namespace, {}).get(key)
 
 
-@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}'"
+# Создаём экземпляр хранилища
+memory_store = MemoryStore()
+
+
+@tool(name="save", description="Сохраняет значение по ключу")
+def tool_save(key: str, value: Any) -> None:
+    memory_store.save(key, value)
+    return f"Ключ '{key}' сохранён."
+
+
+@tool(name="get", description="Получает значение по ключу")
+def tool_get(key: str) -> Any | None:
+    result = memory_store.get(key)
+    if result is None:
+        return f"Ключ '{key}' не найден."
+    return result
+
+
+@tool(name="delete", description="Удаляет ключ из памяти")
+def tool_delete(key: str) -> bool:
+    deleted = memory_store.delete(key)
+    return deleted
+
+
+@tool(name="list_keys", description="Возвращает список всех ключей в памяти")
+def tool_list_keys() -> List[str]:
+    return memory_store.list_keys()
+
+
+@tool(name="save_with_namespace", description="Сохраняет значение под ключом внутри namespace")
+def tool_save_with_namespace(namespace: str, key: str, value: Any) -> None:
+    memory_store.save_with_namespace(namespace, key, value)
+    return f"Ключ '{key}' в namespace '{namespace}' сохранён."
+
+
+@tool(name="get_by_namespace", description="Получает значение по ключу из namespace")
+def tool_get_by_namespace(namespace: str, key: str) -> Any | None:
+    result = memory_store.get_by_namespace(namespace, key)
+    if result is None:
+        return f"Ключ '{key}' в namespace '{namespace}' не найден."
+    return result
+
+
+# Создаём FastMCP приложение
+app = FastMCP(
+    name="MemoryServer",
+    description="FastMCP сервер для управления памятью агента",
 )
-    return ns[request.key]
+
+# Регистрируем все инструменты
+app.register_tool(tool_save)
+app.register_tool(tool_get)
+app.register_tool(tool_delete)
+app.register_tool(tool_list_keys)
+app.register_tool(tool_save_with_namespace)
+app.register_tool(tool_get_by_namespace)
 
 
 if __name__ == "__main__":
-    # Запускаем сервер FastMCP
-    import uvicorn
-
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    # Запускаем сервер на порту 8000
+    app.run(host="0.0.0.0", port=8000)