156 lines
5.5 KiB
Python
156 lines
5.5 KiB
Python
import json
|
|
from typing import List, Dict, Any
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, HTTPException, Body
|
|
from pydantic import BaseModel
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.http.models import PointStruct, VectorParams
|
|
from langchain.memory import ConversationBufferMemory
|
|
from langchain.embeddings.openai import OpenAIEmbeddings
|
|
|
|
# ------------------------------------------------------------------
|
|
# Конфигурация сервера и клиента
|
|
# ------------------------------------------------------------------
|
|
|
|
QDRANT_HOST = "localhost"
|
|
QDRANT_PORT = 6333
|
|
COLLECTION_NAME = "agent_memory"
|
|
|
|
OPENAI_API_KEY = "YOUR_OPENAI_API_KEY" # замените на свой ключ
|
|
|
|
# ------------------------------------------------------------------
|
|
# Модели данных для API
|
|
# ------------------------------------------------------------------
|
|
|
|
class MemoryItem(BaseModel):
|
|
text: str
|
|
metadata: Dict[str, Any] | None = None
|
|
|
|
|
|
class QueryParams(BaseModel):
|
|
query: str
|
|
top_k: int = 5
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Инициализация Qdrant и LangChain
|
|
# ------------------------------------------------------------------
|
|
|
|
client_qdrant = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
|
|
|
# Создаём коллекцию при первом запуске (если её нет)
|
|
if COLLECTION_NAME not in client_qdrant.get_collections().collections:
|
|
client_qdrant.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=1536, distance="Cosine"),
|
|
)
|
|
|
|
embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
|
|
memory_buffer = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# FastAPI приложение
|
|
# ------------------------------------------------------------------
|
|
|
|
app = FastAPI(title="MCP Memory Server")
|
|
|
|
|
|
@app.post("/memory/add")
|
|
def add_memory(item: MemoryItem):
|
|
"""
|
|
Добавляет новый элемент памяти в Qdrant.
|
|
"""
|
|
vector = embeddings.embed_query(item.text)
|
|
point_id = client_qdrant.upload_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
points=[PointStruct(id=None, vector=vector, payload=item.metadata or {})],
|
|
).points[0].id
|
|
return {"status": "ok", "point_id": point_id}
|
|
|
|
|
|
@app.post("/memory/query")
|
|
def query_memory(params: QueryParams):
|
|
"""
|
|
Выполняет поиск похожих элементов памяти.
|
|
"""
|
|
query_vector = embeddings.embed_query(params.query)
|
|
search_result = client_qdrant.search(
|
|
collection_name=COLLECTION_NAME,
|
|
query_vector=query_vector,
|
|
limit=params.top_k,
|
|
)
|
|
results = [
|
|
{"id": hit.id, "score": hit.score, "payload": hit.payload}
|
|
for hit in search_result
|
|
]
|
|
return {"results": results}
|
|
|
|
|
|
@app.post("/memory/clear")
|
|
def clear_memory():
|
|
"""
|
|
Очищает всю коллекцию памяти.
|
|
"""
|
|
client_qdrant.delete_collection(collection_name=COLLECTION_NAME)
|
|
# Пересоздаём пустую коллекцию
|
|
client_qdrant.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=1536, distance="Cosine"),
|
|
)
|
|
return {"status": "cleared"}
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Клиентская часть для взаимодействия с сервером
|
|
# ------------------------------------------------------------------
|
|
|
|
class MemoryClient:
|
|
"""
|
|
Простой клиент для работы с MCP-сервером памяти.
|
|
"""
|
|
|
|
def __init__(self, base_url: str = "http://localhost:8000"):
|
|
self.base_url = base_url.rstrip("/")
|
|
|
|
def add(self, text: str, metadata: Dict[str, Any] | None = None) -> int:
|
|
payload = {"text": text, "metadata": metadata}
|
|
resp = httpx.post(f"{self.base_url}/memory/add", json=payload)
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"Failed to add memory: {resp.text}")
|
|
return resp.json()["point_id"]
|
|
|
|
def query(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
|
|
payload = {"query": query, "top_k": top_k}
|
|
resp = httpx.post(f"{self.base_url}/memory/query", json=payload)
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"Failed to query memory: {resp.text}")
|
|
return resp.json()["results"]
|
|
|
|
def clear(self) -> None:
|
|
resp = httpx.post(f"{self.base_url}/memory/clear")
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"Failed to clear memory: {resp.text}")
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Пример использования (можно запустить как скрипт)
|
|
# ------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
console.print("[bold green]Запуск MCP Memory Server...[/]")
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
|
|
# Пример клиента (не будет выполнен при запуске сервера)
|
|
client = MemoryClient()
|
|
point_id = client.add("Привет, как дела?", {"source": "user"})
|
|
console.print(f"Добавлен пункт памяти с id={point_id}")
|
|
results = client.query("привет")
|
|
console.print("[bold]Результаты поиска:[/]")
|
|
for r in results:
|
|
console.print(r) |