PlantAI-rag-new/sdk/app/services/neo4j.py

"""services/neo4j.py —— Neo4j Bolt 直连（官方 neo4j 驱动，全异步）。

选型与凭据依据 DISCOVERY.md §2：
- 首选 neo4j-plant-srt-direct-v1-full（bolt 127.0.0.1:18788，db `neo4j`，
  用户 neo4j / 密码在 sdk/.env），schema canonical，QA 7920 生产使用。
- 回退 neo4j-plant-srt-pc2025-v3-20260817-normalized（bolt 127.0.0.1:19788），
  schema 不同（Entity/ASSERTS），本波未做 Cypher 适配（TODO）。
- 再次 plantai-protected-neo4j-1（bolt 127.0.0.1:17687），空库。

所有本模块异常统一抛 :class:`Neo4jUnavailable`，由路由层转 503 降级信封。
凭据只存在于本文件（或 env 覆盖），绝不下发浏览器。
"""
from __future__ import annotations

import asyncio
import logging
import os
import time
from typing import Any

from neo4j import AsyncDriver, AsyncGraphDatabase, exceptions as neo4j_exc

logger = logging.getLogger(__name__)

# 可选 env 覆盖（测试/部署）；未设置时使用 DISCOVERY.md 实测凭据。
_ENV_PREFIX = "BFF_NEO4J_"

PROBE_CACHE_TTL = 5.0

# 连接/查询超时（秒）。entity 详情与一跳邻域走全量属性，超时给得宽些。
CONNECT_TIMEOUT = 3.0
QUERY_TIMEOUT = 10.0
NEIGHBORHOOD_QUERY_TIMEOUT = 30.0
OVERVIEW_QUERY_TIMEOUT = 60.0  # 2026-08-25 PM 修复：1644 万边 COUNT 聚合冷缓存 >10s

NAME_INDEX = "srt2_canonical_entity_name_norm"
FULLTEXT_INDEX = None  # DISCOVERY.md §2：任何实例都不存在全文索引，且不得创建

# status 展示名（须与 API.md §6 fixture 一致）
DIRECT_ACTIVE_NAME = "neo4j-plant-srt-direct-v1-full"
NORMALIZED_ACTIVE_NAME = "neo4j-plant-srt-pc2025-v3-20260817-normalized"
PROTECTED_ACTIVE_NAME = "plantai-protected-neo4j-1"

_CANONICAL_LABELS = ["SRT2CanonicalEntity", "SRT2CanonicalAssertion", "SRT2Paper"]
_CANONICAL_REL_TYPES = ["SRT2_CANONICAL_ASSERTS", "SRT2_SUPPORTED_BY"]


class Neo4jUnavailable(Exception):
    """Neo4j 上游不可用（连接失败/认证失败/超时/Cypher 错误统一转此异常）。"""

    def __init__(self, message: str, *, cause: BaseException | None = None) -> None:
        super().__init__(message)
        self.cause = cause


class _Instance:
    """一个候选实例。本波只实现 canonical schema（direct 实例）的 Cypher。"""

    __slots__ = ("name", "uri", "auth", "database", "schema", "label", "rel")

    def __init__(
        self,
        name: str,
        uri: str,
        user: str,
        password: str,
        database: str,
        schema: str,
        label: str | None,
        rel: str | None,
    ) -> None:
        self.name = name
        self.uri = uri
        self.auth = (user, password)
        self.database = database
        self.schema = schema
        self.label = label
        self.rel = rel


def _env(name: str, default: str) -> str:
    return os.environ.get(f"{_ENV_PREFIX}{name}", default)


_INSTANCES: list[_Instance] = [
    _Instance(
        DIRECT_ACTIVE_NAME,
        _env("URI_DIRECT", "bolt://127.0.0.1:18788"),
        "neo4j",
        _env("PASSWORD_DIRECT", ""),
        "neo4j",
        "canonical",
        "SRT2CanonicalEntity",
        "SRT2_CANONICAL_ASSERTS",
    ),
    _Instance(
        NORMALIZED_ACTIVE_NAME,
        _env("URI_NORMALIZED", "bolt://127.0.0.1:19788"),
        "neo4j",
        _env("PASSWORD_NORMALIZED", ""),
        "neo4j",
        "normalized",
        None,
        None,
    ),
    _Instance(
        PROTECTED_ACTIVE_NAME,
        _env("URI_PROTECTED", "bolt://127.0.0.1:17687"),
        "neo4j",
        _env("PASSWORD_PROTECTED", ""),
        "neo4j",
        "empty",
        None,
        None,
    ),
]


class Neo4jService:
    """生命周期：app 启动时 create，关闭时 close()。health 路由与 kg 路由共用。

    内部为每个候选实例缓存一个 AsyncDriver（懒建）。实例选择只发生在
    显式调用 select()（probe/kg_status）或模块加载时：每个请求的处理函数
    首先调用 select()，得到当前实例（未探测时默认首选，惰性回退）。
    切换由下一次显式 probe 触发，探测结果缓存 5s。
    """

    def __init__(self, instances: list[_Instance] | None = None) -> None:
        self._instances = instances if instances is not None else _INSTANCES
        self._drivers: dict[str, AsyncDriver] = {}
        self._active: _Instance = self._instances[0]
        self._selected = False  # 是否已成功选出一个实例
        self._last_probe: float = 0.0
        self._last_probe_ok: bool = False
        self._lock = asyncio.Lock()

    # ---------- 生命周期 ----------

    async def close(self) -> None:
        for driver in self._drivers.values():
            await driver.close()
        self._drivers.clear()

    def _driver_for(self, instance: _Instance) -> AsyncDriver:
        driver = self._drivers.get(instance.name)
        if driver is None:
            driver = AsyncGraphDatabase.driver(
                instance.uri,
                auth=instance.auth,
                connection_timeout=CONNECT_TIMEOUT,
                connection_acquisition_timeout=CONNECT_TIMEOUT,
            )
            self._drivers[instance.name] = driver
        return driver

    # ---------- 实例选择 ----------

    async def select(self) -> _Instance:
        """返回当前可用实例。

        顺序探测（5s 结果缓存）：direct → normalized → protected。
        首次调用若未探测成功则依次尝试探测，全失败抛 Neo4jUnavailable。
        """
        if self._selected:
            return self._active
        now = time.monotonic()
        if now - self._last_probe < PROBE_CACHE_TTL:
            if self._last_probe_ok:
                return self._active
            # 探测失败结果同样缓存 5s，避免每个请求都打 3 个实例
            raise Neo4jUnavailable("Neo4j 不可用（最近探测失败，结果缓存 5s）")
        await self._probe()
        if self._selected:
            return self._active
        raise Neo4jUnavailable("所有 Neo4j 候选实例均不可达")

    async def _probe(self) -> None:
        async with self._lock:
            now = time.monotonic()
            if now - self._last_probe < PROBE_CACHE_TTL:
                if not self._last_probe_ok:
                    raise Neo4jUnavailable("Neo4j 不可用（最近探测失败，结果缓存 5s）")
                return
            self._selected = False
            last_error: BaseException | None = None
            for instance in self._instances:
                try:
                    await self._check_instance(instance)
                except Neo4jUnavailable as exc:
                    last_error = exc
                    continue
                self._active = instance
                self._selected = True
                break
            self._last_probe = now
            self._last_probe_ok = self._selected
            if not self._selected:
                raise Neo4jUnavailable(
                    f"所有 Neo4j 候选实例均不可达：{last_error}", cause=last_error
                )

    async def _check_instance(self, instance: _Instance) -> None:
        """连通性检查：TCP 预检 + `RETURN 1`，整体限时 CONNECT_TIMEOUT。

        先做 TCP 连通预检（毫秒级）：端口没开时直接判死，避免官方驱动
        的指数退避重试拖慢健康探测。TCP 通但 Bolt 协议失败时才走
        wait_for + 驱动（仍限时 3s）。
        """
        if not await self._tcp_reachable(instance.uri):
            raise Neo4jUnavailable(f"{instance.name} 连接失败（TCP 不可达）")
        driver = self._driver_for(instance)
        try:
            await asyncio.wait_for(
                driver.execute_query(
                    "RETURN 1 AS x", database_=instance.database, routing_="r"
                ),
                timeout=CONNECT_TIMEOUT,
            )
        except asyncio.TimeoutError as exc:
            raise Neo4jUnavailable(f"{instance.name} 连接超时", cause=exc) from exc
        except neo4j_exc.DriverError as exc:
            raise Neo4jUnavailable(
                f"{instance.name} 连接失败：{exc.__class__.__name__}", cause=exc
            ) from exc

    @staticmethod
    async def _tcp_reachable(uri: str) -> bool:
        """bolt://host:port 的 TCP 连通预检（超时 0.5s）。"""
        try:
            host, port = uri.split("://", 1)[1].rsplit(":", 1)
        except (ValueError, IndexError):
            return False
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(host, int(port)), timeout=0.5
            )
        except (OSError, asyncio.TimeoutError, ValueError):
            return False
        writer.close()
        try:
            await writer.wait_closed()
        except (OSError, AttributeError):
            pass
        return True

    # ---------- 通用查询 ----------

    async def _query(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
        *,
        timeout: float = QUERY_TIMEOUT,
    ) -> list[dict[str, Any]]:
        """按当前实例执行查询，返回 records 列表（每条为 dict）。"""
        instance = await self.select()
        driver = self._driver_for(instance)
        try:
            records, _, _ = await asyncio.wait_for(
                driver.execute_query(
                    cypher,
                    parameters_=params or {},
                    database_=instance.database,
                    routing_="r",
                ),
                timeout=timeout,
            )
        except asyncio.TimeoutError as exc:
            raise Neo4jUnavailable(
                f"Neo4j 查询超时（>{timeout}s）", cause=exc
            ) from exc
        except neo4j_exc.DriverError as exc:
            raise Neo4jUnavailable(
                f"Neo4j 查询失败：{exc.__class__.__name__}", cause=exc
            ) from exc
        return [dict(record) for record in records]

    @staticmethod
    def _canonical(instance: _Instance) -> _Instance:
        """当前仅 canonical schema 的 Cypher 就绪。"""
        if instance.schema != "canonical":
            # TODO(A3)：normalized 实例（Entity/ASSERTS，无 name_norm 索引，
            # 等值 6.81s）与空库实例的 Cypher 适配层，需 PM 决策后实现。
            raise Neo4jUnavailable(
                f"Neo4j 实例 {instance.name} 已连通，但其 schema（{instance.schema}）"
                "的 Cypher 适配尚未实现，暂不可用于图谱查询"
            )
        return instance

    # ---------- 健康探测 ----------

    async def probe(self) -> dict[str, Any]:
        """连通 + 计数 + 索引检查（/api/health 调用），结果缓存 5s。

        返回 {"ok": bool, "active": str|null, "detail": str|null}，
        主动探测时即使失败也返回 dict（health 聚合不因此抛错）。
        """
        now = time.monotonic()
        if now - self._last_probe < PROBE_CACHE_TTL and self._last_probe_ok:
            return {"ok": True, "active": self._active.name, "detail": None}
        try:
            await self._probe()
        except Neo4jUnavailable as exc:
            return {"ok": False, "active": None, "detail": str(exc)}

        instance = self._active
        detail: str | None = None
        if instance.schema == "canonical":
            try:
                counts = await self._counts(instance)
                indexes = await self._name_indexes(instance)
                # 成员判断而非全等：同一标签上存在唯一约束索引等其它索引是正常
                # 现象（DISCOVERY §2.2 实测有 srt2_canonical_entity_id），
                # 只有 name_norm 索引真正不在列表里才报缺失。
                if NAME_INDEX not in indexes:
                    detail = f"name_norm 索引缺失：{indexes or '(无任何索引)'}"
            except Neo4jUnavailable as exc:
                # 连通了但计数/索引检查失败：不视为整体失败，附 detail
                detail = f"计数/索引检查失败：{exc}"
        return {"ok": True, "active": instance.name, "detail": detail}

    async def _counts(self, instance: _Instance) -> dict[str, int]:
        rows = await self._query(
            """
            CALL {
                MATCH (e:SRT2CanonicalEntity) RETURN count(e) AS entity_count
            }
            CALL {
                MATCH (a:SRT2CanonicalAssertion) RETURN count(a) AS assertion_count
            }
            RETURN entity_count, assertion_count
            """,
            timeout=QUERY_TIMEOUT,
        )
        row = rows[0] if rows else {}
        return {
            "entity_count": row.get("entity_count", 0),
            "assertion_count": row.get("assertion_count", 0),
        }

    async def _name_indexes(self, instance: _Instance) -> list[str]:
        """该实体标签上的所有索引名（含约束索引），滤掉空名。

        SHOW INDEXES 的 labelsOrTypes 可以是字符串或字符串列表，用
        `ANY(x IN labelsOrTypes WHERE x = $label)` 兼容两者。
        """
        rows = await self._query(
            "SHOW INDEXES YIELD name, labelsOrTypes "
            "WHERE ANY(x IN labelsOrTypes WHERE x = $label) "
            "RETURN name",
            {"label": instance.label},
            timeout=QUERY_TIMEOUT,
        )
        return [r["name"] for r in rows if r.get("name")]

    # ---------- /api/kg/status ----------

    async def kg_status(self) -> dict[str, Any]:
        """labels / relationship_types / entity_count / assertion_count /
        name_index / fulltext_index（null，实例无全文索引）。"""
        instance = self._canonical(await self.select())
        counts = await self._counts(instance)
        indexes = await self._name_indexes(instance)
        name_index = NAME_INDEX if NAME_INDEX in indexes else None
        return {
            "ok": True,
            "active": instance.name,
            "schema": instance.schema,
            "labels": _CANONICAL_LABELS,
            "relationship_types": _CANONICAL_REL_TYPES,
            "entity_count": counts["entity_count"],
            "assertion_count": counts["assertion_count"],
            "name_index": name_index,
            "fulltext_index": FULLTEXT_INDEX,
        }

    # ---------- /api/kg/overview ----------

    async def kg_overview(self) -> dict[str, Any]:
        """按 entity_type 与 relation_category 计数。

        注意 1644 万边：只做 COUNT 聚合，绝不拉取全量边。
        """
        self._canonical(await self.select())
        rows = await self._query(
            """
            MATCH (e:SRT2CanonicalEntity)
            WITH e.entity_type AS t, count(*) AS c
            WHERE t IS NOT NULL
            RETURN t AS type, c AS count
            ORDER BY c DESC, t ASC
            LIMIT 200
            """,
            timeout=OVERVIEW_QUERY_TIMEOUT,
        )
        entities = [{"type": r["type"], "count": r["count"]} for r in rows]

        # 关系类别取自边上的 relation_category（属于关系属性，不在边集合里拉边）
        rows = await self._query(
            """
            MATCH ()-[r:SRT2_CANONICAL_ASSERTS]->()
            WITH r.relation_category AS cat, count(*) AS c
            WHERE cat IS NOT NULL
            RETURN cat AS type, c AS count
            ORDER BY c DESC, cat ASC
            LIMIT 200
            """,
            timeout=OVERVIEW_QUERY_TIMEOUT,
        )
        relations = [{"type": r["type"], "count": r["count"]} for r in rows]
        return {
            "entities_by_type": entities,
            "relations_by_category": relations,
        }

    # ---------- /api/kg/entities ----------

    async def search_entities(
        self, q: str, limit: int = 20
    ) -> list[dict[str, Any]]:
        """DISCOVERY.md §2.4 Cypher (a)：name_norm = 或 STARTS WITH（toLower 归一化）。

        走 name_norm RANGE 索引（实测 0.09-0.10s）。结果按 name ASC 排序保证稳定。
        返回 [{id, name, type, support_count}]。
        """
        self._canonical(await self.select())
        limit = min(max(int(limit), 1), 20)
        q_norm = q.lower()
        rows = await self._query(
            """
            MATCH (e:SRT2CanonicalEntity)
            WHERE e.name_norm = $q_norm OR e.name_norm STARTS WITH $q_norm
            RETURN e.canonical_id AS canonical_id, e.name AS name,
                   e.entity_type AS entity_type, e.support_count AS support_count
            ORDER BY e.name ASC
            LIMIT $limit
            """,
            {"q_norm": q_norm, "limit": limit},
            timeout=QUERY_TIMEOUT,
        )
        return [
            {
                "id": r.get("canonical_id"),
                "name": r.get("name"),
                "type": r.get("entity_type"),
                "support_count": r.get("support_count"),
            }
            for r in rows
        ]

    async def _resolve_id(self, id_or_name: str) -> str | None:
        """canonical_id 或 name_norm 定位（DISCOVERY.md §2.4 Cypher b）。"""
        if id_or_name.startswith("cent_"):
            rows = await self._query(
                "MATCH (e:SRT2CanonicalEntity {canonical_id: $id}) "
                "RETURN e.canonical_id AS canonical_id",
                {"id": id_or_name},
                timeout=QUERY_TIMEOUT,
            )
        else:
            rows = await self._query(
                "MATCH (e:SRT2CanonicalEntity {name_norm: $name_norm}) "
                "RETURN e.canonical_id AS canonical_id LIMIT 5",
                {"name_norm": id_or_name.lower()},
                timeout=QUERY_TIMEOUT,
            )
        return rows[0].get("canonical_id") if rows else None

    # ---------- /api/kg/entities/{id} ----------

    async def entity(self, id_or_name: str) -> dict[str, Any] | None:
        """按 canonical_id 或 name_norm 取实体全量属性；未命中返回 None。"""
        canonical_id = await self._resolve_id(id_or_name)
        if canonical_id is None:
            return None
        rows = await self._query(
            "MATCH (e:SRT2CanonicalEntity {canonical_id: $id}) "
            "RETURN e {.*} AS node",
            {"id": canonical_id},
            timeout=QUERY_TIMEOUT,
        )
        if not rows:
            return None
        props = rows[0].get("node") or {}
        return {
            "id": props.get("canonical_id"),
            "name": props.get("name"),
            "type": props.get("entity_type"),
            "properties": props,
        }

    # ---------- /api/kg/entities/{id}/neighborhood ----------

    async def neighborhood(
        self, id_or_name: str, depth: int = 1, limit: int = 200
    ) -> dict[str, Any]:
        """有界邻域扩展（DISCOVERY.md §2.4 Cypher (c) 思路）。

        - depth ∈ {1,2}，默认 1；limit ≤ 200。
        - 每跳按 paper_count/support_count 降序截断（防 hub 吃满预算）。
        - OPTIONAL MATCH 保证孤立节点不丢。
        - 无向匹配（-()-），entity 按 canonical_id 去重。
        - depth=1 绝不执行第二跳（两跳实测约 13s，不默认开启）。
        - 返回 API.md §6 的 {center, center_name, depth, nodes, edges}。
        """
        self._canonical(await self.select())  # 非 canonical 实例在此抛 Neo4jUnavailable
        depth = depth if depth in (1, 2) else 1
        limit = min(max(int(limit), 1), 200)
        canonical_id = await self._resolve_id(id_or_name)
        if canonical_id is None:
            return {"center": None, "center_name": None, "depth": depth, "nodes": [], "edges": []}

        params = {"seed": canonical_id, "max_nodes": limit, "max_edges": limit}
        if depth == 1:
            cypher = """
                MATCH (s:SRT2CanonicalEntity {canonical_id: $seed})
                OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]-(n1:SRT2CanonicalEntity)
                WITH s, r1, n1
                ORDER BY coalesce(r1.paper_count, 0) DESC, coalesce(r1.support_count, 0) DESC
                WITH s,
                     collect(DISTINCT n1)[..$max_nodes] AS level1,
                     [r IN collect(DISTINCT r1)[..$max_edges] | r {.*,
                         source_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
                                                   THEN s.canonical_id ELSE endNode(r).canonical_id END,
                         target_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
                                                   THEN endNode(r).canonical_id ELSE s.canonical_id END}] AS edges1
                RETURN s {.*} AS start,
                       [n IN level1 | n {.*}] AS nodes,
                       edges1 AS edges
            """
        else:
            cypher = """
                MATCH (s:SRT2CanonicalEntity {canonical_id: $seed})
                OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]-(n1:SRT2CanonicalEntity)
                WITH s, r1, n1
                ORDER BY coalesce(r1.paper_count, 0) DESC, coalesce(r1.support_count, 0) DESC
                WITH s,
                     collect(DISTINCT n1)[..$max_nodes] AS level1,
                     [r IN collect(DISTINCT r1)[..$max_edges] | r {.*,
                         source_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
                                                   THEN s.canonical_id ELSE endNode(r).canonical_id END,
                         target_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
                                                   THEN endNode(r).canonical_id ELSE s.canonical_id END}] AS edges1
                OPTIONAL MATCH (n1)-[r2:SRT2_CANONICAL_ASSERTS]-(n2:SRT2CanonicalEntity)
                WHERE n1 IN level1
                WITH s, level1, edges1, r2, n2
                ORDER BY coalesce(r2.paper_count, 0) DESC, coalesce(r2.support_count, 0) DESC
                WITH s, level1, edges1,
                     collect(DISTINCT n2)[..$max_nodes] AS level2,
                     [r IN collect(DISTINCT r2)[..$max_edges] | r {.*,
                         source_canonical_id: startNode(r).canonical_id,
                         target_canonical_id: endNode(r).canonical_id}] AS edges2
                RETURN s {.*} AS start,
                       [n IN (level1 + level2) | n {.*}] AS nodes,
                       (edges1 + edges2) AS edges
            """
        rows = await self._query(
            cypher, params, timeout=NEIGHBORHOOD_QUERY_TIMEOUT
        )
        if not rows:
            return {"center": None, "center_name": None, "depth": depth, "nodes": [], "edges": []}
        row = rows[0]
        start = row.get("start") or {}

        node_list: list[dict[str, Any]] = [
            {
                "id": start.get("canonical_id"),
                "label": start.get("name"),
                "type": start.get("entity_type"),
            }
        ]
        node_index: dict[str, int] = {}
        if start.get("canonical_id") is not None:
            node_index[start["canonical_id"]] = 0
        for raw in row.get("nodes") or []:
            nid = raw.get("canonical_id")
            if nid is None or nid in node_index:
                continue
            node_index[nid] = len(node_list)
            node_list.append(
                {
                    "id": nid,
                    "label": raw.get("name"),
                    "type": raw.get("entity_type"),
                }
            )

        edges: list[dict[str, Any]] = []
        for raw in row.get("edges") or []:
            src = raw.get("source_canonical_id")
            tgt = raw.get("target_canonical_id")
            if src is None or tgt is None:
                continue
            if src not in node_index or tgt not in node_index:
                continue
            # 自环边（实体自身关系，如 gene_structure）对可视化无意义，过滤
            if src == tgt:
                continue
            edges.append(
                {
                    "source": src,
                    "target": tgt,
                    "relation": raw.get("relation_category")
                    if raw.get("relation_category") is not None
                    else raw.get("relation_norm"),
                    "label": raw.get("relation_category")
                    if raw.get("relation_category") is not None
                    else raw.get("relation_norm"),
                    "assertion_id": raw.get("assertion_id"),
                    "paper_count": raw.get("paper_count"),
                    "support_count": raw.get("support_count"),
                    "title": raw.get("title"),
                    "doi": raw.get("doi"),
                    "pmid": raw.get("pmid"),
                    "chain_id": None,
                    "hop": None,
                }
            )

        return {
            "center": start.get("canonical_id"),
            "center_name": start.get("name"),
            "depth": depth,
            "nodes": node_list,
            "edges": edges,
        }

    # ---------- 调试/冒烟 ----------

    async def raw_query(self, cypher: str, **params: Any) -> list[dict[str, Any]]:
        """只读原始查询（冒烟测试用）。"""
        return await self._query(cypher, params, timeout=QUERY_TIMEOUT)


# 模块级单例：health 路由与 kg 路由共用同一连接池（app 关闭时 close）。
service = Neo4jService()


async def close_neo4j() -> None:
    """app shutdown 钩子。"""
    await service.close()


async def probe_neo4j() -> dict[str, Any]:
    """/api/health 用的探测（不抛错）。"""
    return await service.probe()