#!/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 (Qdrant, EmbeddingGemma 2). " "Возвращает типизированные атомы (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 или все), счётчики узлов.", "inputSchema": { "type": "object", "properties": { "kind": {"type": "string", "description": "Фильтр по kind ребра (опционально)."}, "top_k": {"type": "integer", "default": 50, "description": "Число рёбер."}, }, "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()