diff --git a/server.py b/server.py new file mode 100644 index 0000000..2fb4864 --- /dev/null +++ b/server.py @@ -0,0 +1,121 @@ +import json +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 + +from fastmcp import MCPServer, MCPRequest, MCPResponse +from langchain.embeddings import OpenAIEmbeddings +from langchain.vectorstores import Qdrant +from langchain.schema import Document + +from rich.console import Console +from rich.table import Table + +# Настройки сервера +APP_NAME = "AgentMemoryMCP" +APP_VERSION = "1.0.0" +QDRANT_URL = "http://localhost:6333" +COLLECTION_NAME = "agent_memory" + +# Инициализация консоли rich +console = Console() + +# Создаём FastAPI приложение +app = FastAPI(title=APP_NAME, version=APP_VERSION) + +# CORS (если понадобится) +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# Инициализация 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__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file