Обновить memory_server.py

memory_server.py

@@ -1,162 +1,121 @@
-# memory_server.py
-
-"""
-MCP-сервер для управления памятью с поддержкой namespaces.
-"""
-
-from fastmcp import FastMCP
 import json
-from datetime import datetime
+import uuid
-from pathlib import Path
+from typing import List, Optional
-from typing import Any, Optional, List, Dict
 
+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:
+from fastmcp import MCPServer, MCPRequest, MCPResponse
-    """
+from langchain.embeddings import OpenAIEmbeddings
-    Сервер памяти. Хранит данные в JSON‑файле и предоставляет инструменты
+from langchain.vectorstores import Qdrant
-    для сохранения, чтения, удаления и поиска ключей.
+from langchain.schema import Document
-    """
 
-    def __init__(self) -> None:
+from rich.console import Console
-        # Инициализируем FastMCP с именем сервера
+from rich.table import Table
-        self.mcp = FastMCP("Memory-Server")
 
-        # Путь к файлу хранения данных
+# Настройки сервера
-        self.storage_path: Path = Path("./memory_data.json")
+APP_NAME = "AgentMemoryMCP"
-        self.memory: Dict[str, Dict[str, Any]] = self._load_memory()
+APP_VERSION = "1.0.0"
+QDRANT_URL = "http://localhost:6333"
+COLLECTION_NAME = "agent_memory"
 
-    def _load_memory(self) -> Dict[str, Dict[str, Any]]:
+# Инициализация консоли rich
-        """Загружает память из JSON‑файла."""
+console = Console()
-        if not self.storage_path.exists():
-            return {}
-        with open(self.storage_path, "r", encoding="utf-8") as f:
-            return json.load(f)
 
-    def _save_memory(self) -> None:
+# Создаём FastAPI приложение
-        """Сохраняет память в JSON‑файл."""
+app = FastAPI(title=APP_NAME, version=APP_VERSION)
-        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)
 
-    @self.mcp.tool()
+# CORS (если понадобится)
-    def save(self, key: str, value: Any) -> bool:
+app.add_middleware(
-        """
+    CORSMiddleware,
-        Сохраняет значение по ключу в память сервера.
+    allow_origins=["*"],
-
+    allow_credentials=True,
-        Args:
+    allow_methods=["*"],
-            key: Идентификатор для сохранения (уникальный ключ).
+    allow_headers=["*"],
-            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"],
-                    }
 )
-        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()
+    import uvicorn
-    # Запускаем сервер с транспортом stdio
+
-    server.mcp.run(transport="stdio", show_banner=False, log_level="ERROR")
+    uvicorn.run(app, host="0.0.0.0", port=8000)