515 lines
24 KiB
Python
515 lines
24 KiB
Python
#!/usr/bin/env python3
|
||
"""MCP-сервер Pegas Memory — память и база знаний агента Alicorn.
|
||
|
||
Оборачивает pegas_memory (PostgreSQL = источник истины, Qdrant = векторы,
|
||
Memgraph = граф) в MCP-инструменты для DeepSeek Harness (dsh).
|
||
|
||
Транспорт — stdio (JSON-RPC 2.0, MCP 2024-11-05). Только stdlib + pegas_memory.
|
||
Запускается из venv проекта: .venv/bin/python mcp_server.py
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
PROTOCOL_VERSION = "2024-11-05"
|
||
SERVER_NAME = "pegas-memory"
|
||
SERVER_VERSION = "0.1.0"
|
||
|
||
# --------------------------------------------------------------------------
|
||
# Ленивые глобальные ресурсы (соединения открываются при первом обращении)
|
||
# --------------------------------------------------------------------------
|
||
_state = {"conn": None}
|
||
|
||
|
||
def _conn():
|
||
import psycopg
|
||
if _state["conn"] is None or _state["conn"].closed:
|
||
from pegas_memory import config
|
||
_state["conn"] = psycopg.connect(config.PG_DSN, row_factory=dict_row)
|
||
return _state["conn"]
|
||
|
||
|
||
def dict_row(cursor):
|
||
cols = [d.name for d in cursor.description]
|
||
return lambda row: dict(zip(cols, row))
|
||
|
||
|
||
def _sync():
|
||
"""Прогнать outbox → Qdrant + Memgraph (ошибки не фатальны)."""
|
||
from pegas_memory import outbox_sync
|
||
try:
|
||
outbox_sync.sync_all(_conn())
|
||
except Exception as exc: # Qdrant/Memgraph недоступны — не роняем инструмент
|
||
print(f"[pegas-memory] sync skipped: {exc}", file=sys.stderr)
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# Инструменты
|
||
# --------------------------------------------------------------------------
|
||
def _tool_search(args):
|
||
from pegas_memory import qdrant_store
|
||
query = str(args.get("query") or "").strip()
|
||
if not query:
|
||
return json.dumps({"error": "пустой query"}, ensure_ascii=False), True
|
||
top_k = min(int(args.get("top_k") or 5), 20)
|
||
type_filter = args.get("type") or None
|
||
status_filter = args.get("status") or None
|
||
hits = qdrant_store.search(query, top_k=top_k,
|
||
type_filter=type_filter, status_filter=status_filter)
|
||
return json.dumps({"query": query, "hits": hits}, ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_store(args):
|
||
from pegas_memory import db, dedup
|
||
type_ = str(args.get("type") or "fact").strip()
|
||
claim = args.get("claim")
|
||
content = str(args.get("content") or "").strip()
|
||
if not content:
|
||
return json.dumps({"error": "пустой content"}, ensure_ascii=False), True
|
||
if type_ not in ("text", "definition", "logical_rule", "code", "table", "fact"):
|
||
return json.dumps({"error": f"недопустимый type: {type_}"}, ensure_ascii=False), True
|
||
confidence = float(args.get("confidence") or 0.5)
|
||
# семантическая дедупликация (опционально)
|
||
if args.get("dedup", True):
|
||
dup_id = dedup.find_duplicate((claim or "") + " " + content[:800])
|
||
if dup_id is not None:
|
||
return json.dumps({"atom_id": dup_id, "created": False, "duplicate_of": dup_id},
|
||
ensure_ascii=False), False
|
||
atom_id, created = db.upsert_atom(
|
||
_conn(), type=type_, claim=claim, content=content,
|
||
confidence=confidence, status="candidate",
|
||
metadata=(args.get("metadata") or None),
|
||
)
|
||
if created:
|
||
db.enqueue_outbox(_conn(), entity_type="atom", entity_id=atom_id)
|
||
_conn().commit()
|
||
_sync()
|
||
return json.dumps({"atom_id": atom_id, "created": created}, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_ingest(args):
|
||
from pegas_memory import pipeline
|
||
text = str(args.get("text") or "").strip()
|
||
url = args.get("url")
|
||
path = args.get("path")
|
||
if text:
|
||
stats = pipeline.ingest_text(_conn(), text, title=args.get("title"), min_confidence=0.4)
|
||
elif url:
|
||
stats = pipeline.ingest_url(_conn(), url, min_confidence=0.4)
|
||
elif path:
|
||
stats = pipeline.ingest_file(_conn(), path, min_confidence=0.4)
|
||
else:
|
||
return json.dumps({"error": "нужен text|url|path"}, ensure_ascii=False), True
|
||
_sync()
|
||
return json.dumps(stats, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_add_gap(args):
|
||
from pegas_memory import research
|
||
question = str(args.get("question") or "").strip()
|
||
if not question:
|
||
return json.dumps({"error": "пустой question"}, ensure_ascii=False), True
|
||
research.add_gap(_conn(), question,
|
||
reason=str(args.get("reason") or "user_query"),
|
||
priority=float(args.get("priority") or 0.5))
|
||
return json.dumps({"added": question}, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_list_gaps(args):
|
||
from pegas_memory import research
|
||
gaps = research.detect_gaps(_conn())
|
||
return json.dumps({"gaps": gaps}, ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_research(args):
|
||
from pegas_memory import research
|
||
question = str(args.get("question") or "").strip()
|
||
if not question:
|
||
return json.dumps({"error": "пустой question"}, ensure_ascii=False), True
|
||
result = research.research_question(_conn(), question)
|
||
_sync()
|
||
return json.dumps(result, ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_stats(args):
|
||
c = _conn()
|
||
with c.cursor() as cur:
|
||
cur.execute("SELECT count(*) AS n FROM pm_atom")
|
||
atoms = cur.fetchone()["n"]
|
||
cur.execute("SELECT count(*) AS n FROM pm_document")
|
||
docs = cur.fetchone()["n"]
|
||
cur.execute("SELECT count(*) AS n FROM pm_research_gap WHERE status='open'")
|
||
gaps = cur.fetchone()["n"]
|
||
return json.dumps({"atoms": atoms, "documents": docs, "open_gaps": gaps},
|
||
ensure_ascii=False), False
|
||
|
||
|
||
def _tool_get_context(args):
|
||
"""Собрать контекст для агента: поиск + связанные атомы."""
|
||
from pegas_memory import qdrant_store
|
||
query = str(args.get("query") or "").strip()
|
||
if not query:
|
||
return json.dumps({"error": "пустой query"}, ensure_ascii=False), True
|
||
top_k = min(int(args.get("top_k") or 5), 20)
|
||
hits = qdrant_store.search(query, top_k=top_k)
|
||
parts = [f"- [{h['score']:.3f}] ({h['type']}) {h['claim']}" for h in hits]
|
||
return json.dumps({"query": query, "context": parts, "hits": hits},
|
||
ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_store_code(args):
|
||
from pegas_memory import code_symbols
|
||
language = str(args.get("language") or "").strip()
|
||
source = str(args.get("source") or "")
|
||
path = args.get("path") or ""
|
||
if not language or not source:
|
||
return json.dumps({"error": "нужны language и source"}, ensure_ascii=False), True
|
||
repo = args.get("repo") or ""
|
||
n = code_symbols.store_symbols(_conn(), source, language, repo, path)
|
||
return json.dumps({"symbols_stored": n}, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_search_code(args):
|
||
query = str(args.get("query") or "").strip()
|
||
if not query:
|
||
return json.dumps({"error": "пустой query"}, ensure_ascii=False), True
|
||
limit = min(int(args.get("top_k") or 10), 50)
|
||
c = _conn()
|
||
with c.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT id, repo, path, language, symbol_name, kind, signature, line_start "
|
||
"FROM pm_code_symbol WHERE symbol_name ILIKE %s "
|
||
"ORDER BY similarity(symbol_name, %s) DESC LIMIT %s",
|
||
(f"%{query}%", query, limit),
|
||
)
|
||
rows = cur.fetchall()
|
||
return json.dumps({"query": query, "symbols": rows}, ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_code_graph(args):
|
||
from pegas_memory import code_symbols, code_graph
|
||
language = str(args.get("language") or "").strip()
|
||
source = str(args.get("source") or "")
|
||
path = args.get("path") or ""
|
||
if not language or not source:
|
||
return json.dumps({"error": "нужны language и source"}, ensure_ascii=False), True
|
||
repo = args.get("repo") or ""
|
||
ns = code_symbols.store_symbols(_conn(), source, language, repo, path)
|
||
nd = code_graph.store_dependencies(_conn(), source, language, repo, path)
|
||
return json.dumps({"symbols_stored": ns, "dependencies_stored": nd}, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_infer(args):
|
||
from pegas_memory import inference, memgraph_store
|
||
if args.get("if") and args.get("then"):
|
||
# сохранить правило и сразу применить
|
||
rule_id = inference.store_rule(
|
||
_conn(), claim=args.get("claim") or "rule",
|
||
if_=args.get("if"), then=args.get("then"),
|
||
)
|
||
_conn().commit()
|
||
result = inference.infer(_conn())
|
||
result["edges_total"] = memgraph_store.count_edges()
|
||
result["code_symbols"] = memgraph_store.count_code_symbols()
|
||
return json.dumps(result, ensure_ascii=False), False
|
||
|
||
|
||
def _tool_query_graph(args):
|
||
from pegas_memory import memgraph_store
|
||
kind = args.get("kind") or None
|
||
limit = min(int(args.get("top_k") or 50), 200)
|
||
edges = memgraph_store.query_edges(kind=kind, limit=limit)
|
||
return json.dumps({"kind": kind, "edges": edges,
|
||
"nodes": memgraph_store.count(),
|
||
"code_symbols": memgraph_store.count_code_symbols()},
|
||
ensure_ascii=False, indent=2), False
|
||
|
||
|
||
def _tool_verify(args):
|
||
from pegas_memory import verify
|
||
atom_id = args.get("atom_id")
|
||
if atom_id is None:
|
||
# верифицировать все candidate
|
||
result = verify.verify_all(_conn(), self_critique_enabled=bool(args.get("self_critique", True)))
|
||
else:
|
||
result = verify.verify_atom(_conn(), int(atom_id), self_critique_enabled=bool(args.get("self_critique", True)))
|
||
return json.dumps(result, ensure_ascii=False, indent=2), False
|
||
|
||
|
||
TOOLS = [
|
||
{
|
||
"name": "memory_search",
|
||
"description": "Семантический поиск ВНУТРИ твоей базы знаний Pegas Memory (не в интернете!). "
|
||
"Возвращает типизированные атомы (text/definition/logical_rule/code/table/fact) "
|
||
"с оценкой сходства. Фильтры: type, status (candidate/verified/stale/contested).",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "Поисковый запрос."},
|
||
"top_k": {"type": "integer", "default": 5, "description": "Число результатов (1-20)."},
|
||
"type": {"type": "string", "description": "Фильтр по типу атома (опционально)."},
|
||
"status": {"type": "string", "description": "Фильтр по статусу (опционально)."},
|
||
},
|
||
"required": ["query"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_store",
|
||
"description": "Сохранить атом знания в базу (PostgreSQL + Qdrant + Memgraph). "
|
||
"type: text|definition|logical_rule|code|table|fact. Для code передавай исходник, "
|
||
"для logical_rule — JSON {if, then}, для table — JSON или markdown-таблицу.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"type": {"type": "string", "description": "Тип атома."},
|
||
"claim": {"type": "string", "description": "Короткое каноническое утверждение."},
|
||
"content": {"type": "string", "description": "Полное значение атома."},
|
||
"confidence": {"type": "number", "description": "Уверенность 0-1 (default 0.5)."},
|
||
"metadata": {"type": "object", "description": "Дополнительные метаданные (опционально)."},
|
||
"dedup": {"type": "boolean", "description": "Семантическая дедупликация (default true)."},
|
||
},
|
||
"required": ["type", "content"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_ingest",
|
||
"description": "Загрузить документ в базу знаний: сырой текст, URL (web) или путь к файлу (pdf/epub/text). "
|
||
"Дистиллирует содержимое в атомы через локальную LLM. Возвращает статистику.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"text": {"type": "string", "description": "Сырой текст для загрузки."},
|
||
"url": {"type": "string", "description": "URL веб-страницы."},
|
||
"path": {"type": "string", "description": "Путь к файлу (pdf/epub/text)."},
|
||
"title": {"type": "string", "description": "Заголовок документа (для text)."},
|
||
},
|
||
"required": [],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_add_gap",
|
||
"description": "Зарегистрировать пробел в знаниях для последующего исследования.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"question": {"type": "string", "description": "Вопрос/тема пробела."},
|
||
"reason": {"type": "string", "description": "low_coverage|low_confidence|stale|contested|user_query."},
|
||
"priority": {"type": "number", "description": "Приоритет 0-1 (default 0.5)."},
|
||
},
|
||
"required": ["question"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_list_gaps",
|
||
"description": "Список открытых пробелов (вопросов, на которые пока нет ответа) в базе, "
|
||
"отсортированный по приоритету. Вызывай, чтобы узнать, что ещё нужно исследовать.",
|
||
"inputSchema": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
{
|
||
"name": "memory_research",
|
||
"description": "Полный цикл исследования вопроса: поиск в интернете (SearXNG), "
|
||
"загрузка найденных источников, дистилляция в атомы, сохранение.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"question": {"type": "string", "description": "Исследовательский вопрос."},
|
||
},
|
||
"required": ["question"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_stats",
|
||
"description": "Сколько всего в твоей базе знаний: атомов, документов, открытых пробелов. "
|
||
"Вызывай, когда пользователь спрашивает, ЧТО хранится в базе или её состояние.",
|
||
"inputSchema": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
{
|
||
"name": "memory_store_code",
|
||
"description": "Сохранить исходник кода: извлекает символы (функции/классы) через tree-sitter "
|
||
"и кладёт в pm_code_symbol. language: python|javascript|typescript|rust.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"language": {"type": "string", "description": "Язык: python|javascript|typescript|rust."},
|
||
"source": {"type": "string", "description": "Исходный код."},
|
||
"repo": {"type": "string", "description": "Имя репозитория (опционально)."},
|
||
"path": {"type": "string", "description": "Путь файла (опционально)."},
|
||
},
|
||
"required": ["language", "source"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_search_code",
|
||
"description": "Лексический поиск по символам кода (имена функций/классов) в базе.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "Имя символа или его часть."},
|
||
"top_k": {"type": "integer", "default": 10, "description": "Число результатов."},
|
||
},
|
||
"required": ["query"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_code_graph",
|
||
"description": "Построить граф зависимостей кода: извлекает символы (tree-sitter) и зависимости "
|
||
"(calls/imports/inherits), сохраняет в PostgreSQL и Memgraph.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"language": {"type": "string", "description": "Язык: python|javascript|typescript|rust."},
|
||
"source": {"type": "string", "description": "Исходный код."},
|
||
"repo": {"type": "string", "description": "Имя репозитория (опционально)."},
|
||
"path": {"type": "string", "description": "Путь файла (опционально)."},
|
||
},
|
||
"required": ["language", "source"],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_infer",
|
||
"description": "Запустить логическую инференцию в Memgraph. Правила — атомы logical_rule "
|
||
"с JSON {if, then}. Можно передать if/then, чтобы добавить правило на лету.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"if": {"type": "string", "description": "Условие (kind ребра) или список для композиции."},
|
||
"then": {"type": "string", "description": "Вывод (kind ребра)."},
|
||
"claim": {"type": "string", "description": "Описание правила."},
|
||
},
|
||
"required": [],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_query_graph",
|
||
"description": "Запросить граф связей: рёбра между знаниями/кодом и счётчики узлов. "
|
||
"ВАЖНО: kind — это тип РЕБРА (связи), НЕ тип атома (text/fact/...) и НЕ поисковый запрос. "
|
||
"Если не уверен в типе ребра — оставь kind пустым, вернутся все рёбра.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"kind": {"type": "string", "description": "Тип РЕБРА (не тип атома!): imports | depends_on | calls | inherits | related. Оставь пустым — вернутся все рёбра."},
|
||
"top_k": {"type": "integer", "default": 50, "description": "Сколько рёбер вернуть (1-200)."},
|
||
},
|
||
"required": [],
|
||
},
|
||
},
|
||
{
|
||
"name": "memory_verify",
|
||
"description": "Верифицировать атом знания: self-critique (LLM) + кросс-подтверждение "
|
||
"по эмбеддингам из нескольких источников. Статус: candidate→verified/contested/stale.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"atom_id": {"type": "integer", "description": "ID атома (если не указан — все candidate)."},
|
||
"self_critique": {"type": "boolean", "description": "Включить LLM self-critique."},
|
||
},
|
||
"required": [],
|
||
},
|
||
},
|
||
|
||
|
||
|
||
{
|
||
"name": "memory_get_context",
|
||
"description": "Собрать релевантный контекст из твоей базы для ответа (поиск по атомам + форматирование). "
|
||
"Используй перед ответом на вопрос, чтобы опереться на уже известные факты.",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string", "description": "Запрос контекста."},
|
||
"top_k": {"type": "integer", "default": 5, "description": "Число результатов."},
|
||
},
|
||
"required": ["query"],
|
||
},
|
||
},
|
||
]
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# JSON-RPC 2.0 (MCP 2024-11-05)
|
||
# --------------------------------------------------------------------------
|
||
def handle(msg):
|
||
method = msg.get("method", "")
|
||
msg_id = msg.get("id")
|
||
|
||
if method == "initialize":
|
||
return {"jsonrpc": "2.0", "id": msg_id, "result": {
|
||
"protocolVersion": PROTOCOL_VERSION,
|
||
"capabilities": {"tools": {}},
|
||
"serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
|
||
}}
|
||
|
||
if method == "notifications/initialized":
|
||
return None
|
||
|
||
if method in ("ping", "notifications/cancelled"):
|
||
return {"jsonrpc": "2.0", "id": msg_id, "result": {}} if msg_id is not None else None
|
||
|
||
if method == "tools/list":
|
||
return {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": TOOLS}}
|
||
|
||
if method == "tools/call":
|
||
params = msg.get("params") or {}
|
||
tool = params.get("name", "")
|
||
args = params.get("arguments") or {}
|
||
handlers = {
|
||
"memory_search": _tool_search,
|
||
"memory_store": _tool_store,
|
||
"memory_ingest": _tool_ingest,
|
||
"memory_add_gap": _tool_add_gap,
|
||
"memory_list_gaps": _tool_list_gaps,
|
||
"memory_research": _tool_research,
|
||
"memory_stats": _tool_stats,
|
||
"memory_store_code": _tool_store_code,
|
||
"memory_search_code": _tool_search_code,
|
||
"memory_code_graph": _tool_code_graph,
|
||
"memory_infer": _tool_infer,
|
||
"memory_query_graph": _tool_query_graph,
|
||
"memory_verify": _tool_verify,
|
||
"memory_get_context": _tool_get_context,
|
||
}
|
||
fn = handlers.get(tool)
|
||
if fn is None:
|
||
return {"jsonrpc": "2.0", "id": msg_id,
|
||
"error": {"code": -32602, "message": f"Unknown tool: {tool}"}}
|
||
try:
|
||
text, failed = fn(args)
|
||
except Exception as exc: # сервер не должен падать от одного вызова
|
||
text = json.dumps({"error": f"{type(exc).__name__}: {exc}"}, ensure_ascii=False)
|
||
failed = True
|
||
return {"jsonrpc": "2.0", "id": msg_id,
|
||
"result": {"content": [{"type": "text", "text": text}], "isError": failed}}
|
||
|
||
return {"jsonrpc": "2.0", "id": msg_id,
|
||
"error": {"code": -32601, "message": f"Method not found: {method}"}}
|
||
|
||
|
||
def main():
|
||
for line in sys.stdin:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
msg = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
response = {"jsonrpc": "2.0", "id": None,
|
||
"error": {"code": -32700, "message": "Parse error"}}
|
||
else:
|
||
try:
|
||
response = handle(msg)
|
||
except Exception as exc:
|
||
response = {"jsonrpc": "2.0", "id": msg.get("id"),
|
||
"error": {"code": -32000, "message": f"{type(exc).__name__}: {exc}"}}
|
||
if response is not None:
|
||
sys.stdout.write(json.dumps(response, ensure_ascii=False) + "\n")
|
||
sys.stdout.flush()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|