MCP-сервер для управления памятью агента: memory_server.py
This commit is contained in:
+45
-92
@@ -1,109 +1,62 @@
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
import asyncio
|
||||
from fastmcp import MCPServer, Tool
|
||||
from typing import Dict, Any
|
||||
|
||||
from fastmcp import FastMCP, tool
|
||||
# Внутреннее хранилище памяти
|
||||
memory_store: Dict[str, Any] = {}
|
||||
namespace_store: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
"""
|
||||
Простое хранилище в памяти с поддержкой пространств имён.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# основной словарь: ключ → значение
|
||||
self._store: Dict[str, Any] = {}
|
||||
# словари по namespace: namespace → {key: value}
|
||||
self._namespaces: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
# ---------- Основные операции ----------
|
||||
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)
|
||||
|
||||
def delete(self, key: str) -> bool:
|
||||
"""Удаляет ключ. Возвращает True, если удалено."""
|
||||
if key in self._store:
|
||||
del self._store[key]
|
||||
return True
|
||||
return False
|
||||
|
||||
def list_keys(self) -> List[str]:
|
||||
"""Возвращает список всех ключей."""
|
||||
return list(self._store.keys())
|
||||
|
||||
# ---------- Операции с namespace ----------
|
||||
def save_with_namespace(self, namespace: str, key: str, value: Any) -> None:
|
||||
"""Сохраняет значение в указанном пространстве имён."""
|
||||
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)
|
||||
def save(key: str, value: Any) -> str:
|
||||
"""Сохраняет значение по ключу."""
|
||||
memory_store[key] = value
|
||||
return f"Значение сохранено под ключом '{key}'."
|
||||
|
||||
|
||||
# Создаём экземпляр хранилища
|
||||
memory_store = MemoryStore()
|
||||
def get(key: str) -> Any:
|
||||
"""Возвращает значение по ключу. Если ключ не найден, возвращается None."""
|
||||
return memory_store.get(key)
|
||||
|
||||
|
||||
@tool(name="save", description="Сохраняет значение по ключу")
|
||||
def tool_save(key: str, value: Any) -> None:
|
||||
memory_store.save(key, value)
|
||||
return f"Ключ '{key}' сохранён."
|
||||
def delete(key: str) -> str:
|
||||
"""Удаляет значение по ключу."""
|
||||
if key in memory_store:
|
||||
del memory_store[key]
|
||||
return f"Ключ '{key}' удалён."
|
||||
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
|
||||
def list_keys() -> Dict[str, Any]:
|
||||
"""Возвращает список всех сохранённых ключей."""
|
||||
return {"keys": list(memory_store.keys())}
|
||||
|
||||
|
||||
@tool(name="delete", description="Удаляет ключ из памяти")
|
||||
def tool_delete(key: str) -> bool:
|
||||
deleted = memory_store.delete(key)
|
||||
return deleted
|
||||
def save_with_namespace(namespace: str, key: str, value: Any) -> str:
|
||||
"""Сохраняет значение в указанном пространстве имён."""
|
||||
ns = namespace_store.setdefault(namespace, {})
|
||||
ns[key] = value
|
||||
return f"Значение сохранено под ключом '{key}' в пространстве имён '{namespace}'."
|
||||
|
||||
def get_by_namespace(namespace: str, key: str) -> Any:
|
||||
"""Получает значение из указанного пространства имён."""
|
||||
ns = namespace_store.get(namespace, {})
|
||||
return ns.get(key)
|
||||
|
||||
|
||||
@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 сервер для управления памятью агента",
|
||||
# Создаём сервер MCP
|
||||
server = MCPServer(
|
||||
name="memory_server",
|
||||
description="Сервер для управления памятью агента.",
|
||||
)
|
||||
|
||||
# Регистрируем все инструменты
|
||||
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)
|
||||
|
||||
# Регистрируем инструменты
|
||||
server.register_tool(Tool(name="save", func=save))
|
||||
server.register_tool(Tool(name="get", func=get))
|
||||
server.register_tool(Tool(name="delete", func=delete))
|
||||
server.register_tool(Tool(name="list_keys", func=list_keys))
|
||||
server.register_tool(Tool(name="save_with_namespace", func=save_with_namespace))
|
||||
server.register_tool(Tool(name="get_by_namespace", func=get_by_namespace))
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Запускаем сервер на порту 8000
|
||||
app.run(host="0.0.0.0", port=8000)
|
||||
# Запускаем сервер
|
||||
asyncio.run(server.serve())
|
||||
Reference in New Issue
Block a user