121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
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) |