Initial commit: Pegas Memory (PG + Qdrant + Memgraph, MCP server, 14 tools)

mcp_server.py

@@ -0,0 +1,509 @@
+#!/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()