Обновить memory_server.py
This commit is contained in:
+109
-150
@@ -1,162 +1,121 @@
|
||||
# memory_server.py
|
||||
|
||||
"""
|
||||
MCP-сервер для управления памятью с поддержкой namespaces.
|
||||
"""
|
||||
|
||||
from fastmcp import FastMCP
|
||||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, List, Dict
|
||||
import uuid
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.logger import logger as fastapi_logger
|
||||
|
||||
class MemoryServer:
|
||||
"""
|
||||
Сервер памяти. Хранит данные в JSON‑файле и предоставляет инструменты
|
||||
для сохранения, чтения, удаления и поиска ключей.
|
||||
"""
|
||||
from fastmcp import MCPServer, MCPRequest, MCPResponse
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from langchain.vectorstores import Qdrant
|
||||
from langchain.schema import Document
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Инициализируем FastMCP с именем сервера
|
||||
self.mcp = FastMCP("Memory-Server")
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
# Путь к файлу хранения данных
|
||||
self.storage_path: Path = Path("./memory_data.json")
|
||||
self.memory: Dict[str, Dict[str, Any]] = self._load_memory()
|
||||
# Настройки сервера
|
||||
APP_NAME = "AgentMemoryMCP"
|
||||
APP_VERSION = "1.0.0"
|
||||
QDRANT_URL = "http://localhost:6333"
|
||||
COLLECTION_NAME = "agent_memory"
|
||||
|
||||
def _load_memory(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Загружает память из JSON‑файла."""
|
||||
if not self.storage_path.exists():
|
||||
return {}
|
||||
with open(self.storage_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
# Инициализация консоли rich
|
||||
console = Console()
|
||||
|
||||
def _save_memory(self) -> 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(self.memory, f, indent=2, ensure_ascii=False)
|
||||
# Создаём FastAPI приложение
|
||||
app = FastAPI(title=APP_NAME, version=APP_VERSION)
|
||||
|
||||
@self.mcp.tool()
|
||||
def save(self, key: str, value: Any) -> bool:
|
||||
"""
|
||||
Сохраняет значение по ключу в память сервера.
|
||||
|
||||
Args:
|
||||
key: Идентификатор для сохранения (уникальный ключ).
|
||||
value: Любое сериализуемое значение.
|
||||
|
||||
Returns:
|
||||
True при успешном сохранении.
|
||||
"""
|
||||
self.memory[key] = {
|
||||
"value": value,
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
}
|
||||
self._save_memory()
|
||||
return True
|
||||
|
||||
@self.mcp.tool()
|
||||
def get(self, key: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Возвращает значение по ключу с метаданными.
|
||||
|
||||
Args:
|
||||
key: Идентификатор для поиска.
|
||||
|
||||
Returns:
|
||||
Словарь с полями {"key": ..., "value": ..., "timestamp": ...}
|
||||
или None, если ключ не найден.
|
||||
"""
|
||||
entry = self.memory.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
return {
|
||||
"key": key,
|
||||
"value": entry["value"],
|
||||
"timestamp": entry["timestamp"],
|
||||
}
|
||||
|
||||
@self.mcp.tool()
|
||||
def delete(self, key: str) -> bool:
|
||||
"""
|
||||
Удаляет ключ из памяти сервера.
|
||||
|
||||
Args:
|
||||
key: Идентификатор для удаления.
|
||||
|
||||
Returns:
|
||||
True при успешном удалении, False если ключ не найден.
|
||||
"""
|
||||
if key in self.memory:
|
||||
del self.memory[key]
|
||||
self._save_memory()
|
||||
return True
|
||||
return False
|
||||
|
||||
@self.mcp.tool()
|
||||
def list_keys(self, pattern: str = "*") -> List[str]:
|
||||
"""
|
||||
Возвращает список всех ключей с поддержкой wildcard‑паттерна.
|
||||
|
||||
Args:
|
||||
pattern: Паттерн для фильтрации (поддерживает * и ?).
|
||||
|
||||
Returns:
|
||||
Список совпадающих ключей.
|
||||
"""
|
||||
import fnmatch
|
||||
|
||||
return [k for k in self.memory.keys() if fnmatch.fnmatch(k, pattern)]
|
||||
|
||||
@self.mcp.tool()
|
||||
def save_with_namespace(
|
||||
self,
|
||||
key: str,
|
||||
value: Any,
|
||||
namespace: str = "default",
|
||||
) -> bool:
|
||||
"""
|
||||
Сохраняет значение с указанием пространства имён.
|
||||
|
||||
Args:
|
||||
key: Идентификатор.
|
||||
value: Значение для сохранения.
|
||||
namespace: Пространство имён (по умолчанию 'default').
|
||||
|
||||
Returns:
|
||||
True при успехе, False иначе.
|
||||
"""
|
||||
ns_key = f"{namespace}:{key}"
|
||||
return self.save(ns_key, value)
|
||||
|
||||
@self.mcp.tool()
|
||||
def get_by_namespace(self, namespace: str = "default") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Возвращает все ключи из указанного namespace.
|
||||
|
||||
Args:
|
||||
namespace: Пространство имён для чтения.
|
||||
|
||||
Returns:
|
||||
Список словарей с метаданными всех ключей namespace.
|
||||
"""
|
||||
result: List[Dict[str, Any]] = []
|
||||
prefix = f"{namespace}:"
|
||||
for k, entry in self.memory.items():
|
||||
if k.startswith(prefix):
|
||||
result.append(
|
||||
{
|
||||
"key": k,
|
||||
"value": entry["value"],
|
||||
"timestamp": entry["timestamp"],
|
||||
}
|
||||
# CORS (если понадобится)
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
return result
|
||||
|
||||
# Инициализация Qdrant
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = Qdrant(
|
||||
client=qdrant_client := Qdrant(
|
||||
url=QDRANT_URL,
|
||||
collection_name=COLLECTION_NAME,
|
||||
embeddings=embeddings,
|
||||
),
|
||||
embeddings=embeddings,
|
||||
collection_name=COLLECTION_NAME,
|
||||
)
|
||||
|
||||
# Создаём MCP сервер
|
||||
mcp_server = MCPServer(app=app, name=APP_NAME, version=APP_VERSION)
|
||||
|
||||
|
||||
# Вспомогательные функции
|
||||
def _log_request(req: MCPRequest):
|
||||
table = Table(title="MCP Request", show_header=True, header_style="bold magenta")
|
||||
table.add_column("Method")
|
||||
table.add_column("Path")
|
||||
table.add_column("Body")
|
||||
table.add_row(req.method, req.path, json.dumps(req.body, indent=2))
|
||||
console.print(table)
|
||||
|
||||
|
||||
def _log_response(res: MCPResponse):
|
||||
table = Table(title="MCP Response", show_header=True, header_style="bold green")
|
||||
table.add_column("Status")
|
||||
table.add_column("Body")
|
||||
table.add_row(str(res.status), json.dumps(res.body, indent=2))
|
||||
console.print(table)
|
||||
|
||||
|
||||
# Обработчики MCP
|
||||
@mcp_server.on("memory.add")
|
||||
async def handle_memory_add(req: MCPRequest) -> MCPResponse:
|
||||
_log_request(req)
|
||||
content = req.body.get("content")
|
||||
if not content:
|
||||
raise HTTPException(status_code=400, detail="Missing 'content' field")
|
||||
# Создаём документ и сохраняем в Qdrant
|
||||
doc = Document(page_content=content, metadata={"id": str(uuid.uuid4())})
|
||||
vectorstore.add_documents([doc])
|
||||
res = MCPResponse(status=200, body={"id": doc.metadata["id"]})
|
||||
_log_response(res)
|
||||
return res
|
||||
|
||||
|
||||
@mcp_server.on("memory.get")
|
||||
async def handle_memory_get(req: MCPRequest) -> MCPResponse:
|
||||
_log_request(req)
|
||||
query = req.body.get("query")
|
||||
if not query:
|
||||
raise HTTPException(status_code=400, detail="Missing 'query' field")
|
||||
# Поиск похожих документов
|
||||
results = vectorstore.similarity_search(query, k=5)
|
||||
res_body = [
|
||||
{"id": doc.metadata.get("id"), "content": doc.page_content} for doc in results
|
||||
]
|
||||
res = MCPResponse(status=200, body=res_body)
|
||||
_log_response(res)
|
||||
return res
|
||||
|
||||
|
||||
@mcp_server.on("memory.delete")
|
||||
async def handle_memory_delete(req: MCPRequest) -> MCPResponse:
|
||||
_log_request(req)
|
||||
doc_id = req.body.get("id")
|
||||
if not doc_id:
|
||||
raise HTTPException(status_code=400, detail="Missing 'id' field")
|
||||
# Удаляем документ по id
|
||||
deleted = vectorstore.delete(ids=[doc_id])
|
||||
res = MCPResponse(status=200, body={"deleted": deleted})
|
||||
_log_response(res)
|
||||
return res
|
||||
|
||||
|
||||
# Запуск через uvicorn
|
||||
if __name__ == "__main__":
|
||||
server = MemoryServer()
|
||||
# Запускаем сервер с транспортом stdio
|
||||
server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
Reference in New Issue
Block a user