Files
pegas-memory/mcp_server.py
T

595 lines
28 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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):
"""Row factory: dict по именам колонок.
cursor.description может быть None (UPDATE/DELETE без RETURNING) —
в этом случае возвращаем строку как есть, иначе падаем на итерации None.
"""
if cursor.description is None:
return lambda row: row
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
try:
gap_id = research.add_gap(_conn(), question,
reason=str(args.get("reason") or "user_query"),
priority=float(args.get("priority") or 0.5))
except ValueError as e:
return json.dumps({"error": str(e)}, ensure_ascii=False), True
return json.dumps({"gap_id": gap_id, "question": question, "status": "open"},
ensure_ascii=False), False
def _tool_list_gaps(args):
from pegas_memory import research
status = args.get("status") or None
limit = int(args.get("limit") or 50)
try:
gaps = research.list_gaps(_conn(), status=status, limit=limit)
except ValueError as e:
return json.dumps({"error": str(e)}, ensure_ascii=False), True
return json.dumps({"status": status or "all", "count": len(gaps), "gaps": gaps},
ensure_ascii=False, indent=2, default=str), False
def _tool_get_gap(args):
from pegas_memory import research
gap_id = args.get("gap_id")
if gap_id is None:
return json.dumps({"error": "нужен gap_id"}, ensure_ascii=False), True
gap = research.get_gap(_conn(), int(gap_id))
if gap is None:
return json.dumps({"error": f"пробел {gap_id} не найден"}, ensure_ascii=False), True
return json.dumps(gap, ensure_ascii=False, indent=2, default=str), False
def _tool_update_gap(args):
from pegas_memory import research
gap_id = args.get("gap_id")
if gap_id is None:
return json.dumps({"error": "нужен gap_id"}, ensure_ascii=False), True
fields = {k: args[k] for k in ("question", "reason", "priority", "status")
if args.get(k) is not None}
if not fields:
return json.dumps({"error": "нечего обновлять (укажи question/reason/priority/status)"},
ensure_ascii=False), True
try:
ok = research.update_gap(_conn(), int(gap_id), **fields)
except ValueError as e:
return json.dumps({"error": str(e)}, ensure_ascii=False), True
if not ok:
return json.dumps({"error": f"пробел {gap_id} не найден"}, ensure_ascii=False), True
return json.dumps({"updated": int(gap_id), "gap": research.get_gap(_conn(), int(gap_id))},
ensure_ascii=False, indent=2, default=str), False
def _tool_delete_gap(args):
from pegas_memory import research
gap_id = args.get("gap_id")
if gap_id is None:
return json.dumps({"error": "нужен gap_id"}, ensure_ascii=False), True
if not research.delete_gap(_conn(), int(gap_id)):
return json.dumps({"error": f"пробел {gap_id} не найден"}, ensure_ascii=False), True
return json.dumps({"deleted": int(gap_id)}, ensure_ascii=False), 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": {
"status": {"type": "string", "description": "Фильтр: open | researching | resolved | dismissed (опционально)."},
"limit": {"type": "integer", "default": 50, "description": "Сколько вернуть (1-200)."},
},
"required": [],
},
},
{
"name": "memory_get_gap",
"description": "Получить один пробел по id (все поля).",
"inputSchema": {
"type": "object",
"properties": {"gap_id": {"type": "integer", "description": "ID пробела."}},
"required": ["gap_id"],
},
},
{
"name": "memory_update_gap",
"description": "Отредактировать пробел: текст вопроса, причину, приоритет и/или статус. "
"Статусы: open (открыт) | researching (взят в работу) | resolved (закрыт) | dismissed (не актуален).",
"inputSchema": {
"type": "object",
"properties": {
"gap_id": {"type": "integer", "description": "ID пробела."},
"question": {"type": "string", "description": "Новая формулировка вопроса."},
"reason": {"type": "string", "description": "low_coverage|low_confidence|stale|contested|user_query."},
"priority": {"type": "number", "description": "Новый приоритет 0-1."},
"status": {"type": "string", "description": "open | researching | resolved | dismissed."},
},
"required": ["gap_id"],
},
},
{
"name": "memory_delete_gap",
"description": "Удалить пробел по id (если он больше не нужен). "
"Обычно вместо удаления лучше перевести в status=dismissed или resolved.",
"inputSchema": {
"type": "object",
"properties": {"gap_id": {"type": "integer", "description": "ID пробела."}},
"required": ["gap_id"],
},
},
{
"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_get_gap": _tool_get_gap,
"memory_update_gap": _tool_update_gap,
"memory_delete_gap": _tool_delete_gap,
"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()