PlantAI-rag-new/agent/qa/src/rag/normalized_release.py

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any, Iterable


# 库名锁已 env 化：镜像文件 agent/legacy-rag/src/rag/normalized_release.py 必须逐字同步(md5 校验)。
EXPECTED_RELEASE = os.environ.get("NORMALIZED_RELEASE_ID") or "plant_srt_pc2025_v3_20260817_normalized"
MAX_SEEDS = 32
MAX_RESULTS = 200


def _text(value: Any) -> str:
    return "" if value is None else str(value).strip()


def _dedupe(values: Iterable[Any], *, normalize: bool = False) -> list[str]:
    output: list[str] = []
    seen: set[str] = set()
    for value in values:
        text = " ".join(_text(value).split())
        if normalize:
            text = text.casefold()
        if not text or text in seen:
            continue
        seen.add(text)
        output.append(text)
        if len(output) >= MAX_SEEDS:
            break
    return output


@dataclass(frozen=True)
class NormalizedReleaseSettings:
    release_id: str
    qdrant_collection: str

    @classmethod
    def from_env(cls) -> "NormalizedReleaseSettings":
        release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip() or EXPECTED_RELEASE
        collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip() or EXPECTED_RELEASE
        return cls(release_id=release_id, qdrant_collection=collection)


class NormalizedGraphAdapter:
    """Read-only, bounded access to the normalized Phase 2B Neo4j schema."""

    _ONE_HOP = """
        MATCH (seed:Entity)
        WHERE seed.canonical_id IN $canonical_ids
           OR seed.name_norm IN $name_norms
        MATCH (seed)-[r:ASSERTS]-(other:Entity)
        WITH DISTINCT startNode(r) AS source,
             endNode(r) AS target,
             r
        OPTIONAL MATCH (source)-[support:SUPPORTED_BY]->(paper:Paper)
        WHERE support.assertion_id = r.assertion_id
        RETURN r.assertion_id AS assertion_id,
               source.name AS source,
               source.entity_type AS source_type,
               source.canonical_id AS source_canonical_id,
               r.relation_norm AS relation,
               r.relation_category AS relation_category,
               target.name AS target,
               target.entity_type AS target_type,
               target.canonical_id AS target_canonical_id,
               r.layer AS layer,
               paper.paper_id AS paper_id,
               paper.pmid AS pmid,
               paper.pmcid AS pmcid,
               paper.doi AS doi,
               paper.title AS title
        LIMIT $limit
    """

    _IDENTITY = """
        MATCH (entity:Entity)
        WITH count(entity) AS entity_count
        MATCH (paper:Paper)
        RETURN entity_count,
               count(paper) AS paper_count
    """

    def __init__(self, driver: Any) -> None:
        self.driver = driver

    def identity(self) -> dict[str, Any]:
        with self.driver.session() as session:
            record = session.run(self._IDENTITY).single()
        if record is None:
            raise RuntimeError("normalized graph identity query returned no row")
        result = dict(record)
        result["release_id"] = EXPECTED_RELEASE
        result["graph_schema"] = "normalized"
        return result

    def resolve_entities(
        self,
        *,
        canonical_ids: Iterable[Any],
        names: Iterable[Any],
        limit: int = MAX_SEEDS,
    ) -> list[dict[str, Any]]:
        ids = _dedupe(canonical_ids)
        name_norms = _dedupe(names, normalize=True)
        if not ids and not name_norms:
            return []
        bounded_limit = max(1, min(int(limit or 1), MAX_RESULTS))
        statement = """
            MATCH (entity:Entity)
            WHERE entity.canonical_id IN $canonical_ids
               OR entity.name_norm IN $name_norms
            RETURN entity.canonical_id AS canonical_id,
                   entity.name AS name,
                   entity.name_norm AS name_norm,
                   entity.entity_type AS entity_type,
                   entity.layer AS layer
            LIMIT $limit
        """
        with self.driver.session() as session:
            records = session.run(
                statement,
                canonical_ids=ids,
                name_norms=name_norms,
                limit=bounded_limit,
            )
            return [dict(record) for record in records]

    def one_hop(
        self,
        *,
        canonical_ids: Iterable[Any],
        names: Iterable[Any],
        limit: int,
    ) -> list[dict[str, Any]]:
        ids = _dedupe(canonical_ids)
        name_norms = _dedupe(names, normalize=True)
        if not ids and not name_norms:
            return []
        bounded_limit = max(1, min(int(limit or 1), MAX_RESULTS))
        with self.driver.session() as session:
            records = session.run(
                self._ONE_HOP,
                canonical_ids=ids,
                name_norms=name_norms,
                limit=bounded_limit,
            )
            output = [dict(record) for record in records]
        for row in output:
            row.setdefault("relation_id", _text(row.get("assertion_id")))
            row.setdefault("evidence", "")
            row.setdefault("source_system", "V15_V16_canonical")
            row.setdefault("kb_version", EXPECTED_RELEASE)
            row.setdefault("output_layer", _text(row.get("layer")))
        return output

    def entity_type_counts(self) -> list[dict[str, Any]]:
        statement = """
            MATCH (entity:Entity)
            RETURN entity.entity_type AS entity_type,
                   count(entity) AS count
            ORDER BY count DESC
        """
        with self.driver.session() as session:
            return [dict(record) for record in session.run(statement)]

    def relationship_counts(self) -> list[dict[str, Any]]:
        statement = """
            MATCH ()-[relationship]->()
            WHERE type(relationship) IN ['ASSERTS', 'SUPPORTED_BY']
            RETURN type(relationship) AS rel_type,
                   count(relationship) AS count
            ORDER BY count DESC
        """
        with self.driver.session() as session:
            return [dict(record) for record in session.run(statement)]

    def list_entities(
        self,
        *,
        entity_types: Iterable[Any],
        limit: int,
        offset: int,
    ) -> list[dict[str, Any]]:
        types = _dedupe(entity_types)
        if not types:
            return []
        bounded_limit = max(1, min(int(limit or 1), 1000))
        bounded_offset = max(0, min(int(offset or 0), 1_000_000))
        statement = """
            MATCH (entity:Entity)
            WHERE entity.entity_type IN $entity_types
            RETURN entity.canonical_id AS canonical_id,
                   entity.name AS name,
                   entity.name_norm AS name_norm,
                   entity.entity_type AS entity_type,
                   entity.layer AS layer
            ORDER BY entity.name_norm, entity.canonical_id
            SKIP $offset
            LIMIT $limit
        """
        with self.driver.session() as session:
            records = session.run(
                statement,
                entity_types=types,
                offset=bounded_offset,
                limit=bounded_limit,
            )
            return [dict(record) for record in records]

    def search_entities(
        self,
        query: str,
        *,
        entity_types: Iterable[Any],
        limit: int,
    ) -> list[dict[str, Any]]:
        query_norm = " ".join(_text(query).casefold().split())
        if not query_norm:
            return []
        types = _dedupe(entity_types)
        bounded_limit = max(1, min(int(limit or 1), 100))
        statement = """
            MATCH (entity:Entity)
            WHERE ($all_types OR entity.entity_type IN $entity_types)
              AND (
                    entity.name_norm = $query_norm
                 OR entity.name_norm STARTS WITH $query_norm
              )
            RETURN entity.canonical_id AS canonical_id,
                   entity.name AS name,
                   entity.name_norm AS name_norm,
                   entity.entity_type AS entity_type,
                   entity.layer AS layer
            ORDER BY CASE WHEN entity.name_norm = $query_norm THEN 0 ELSE 1 END,
                     entity.name_norm,
                     entity.canonical_id
            LIMIT $limit
        """
        with self.driver.session() as session:
            records = session.run(
                statement,
                all_types=not types,
                entity_types=types,
                query_norm=query_norm,
                limit=bounded_limit,
            )
            return [dict(record) for record in records]

    def entity_detail(
        self,
        name: str,
        *,
        entity_types: Iterable[Any],
        neighbor_limit: int,
    ) -> dict[str, Any] | None:
        matches = self.search_entities(
            name,
            entity_types=entity_types,
            limit=10,
        )
        query_norm = " ".join(_text(name).casefold().split())
        entity = next(
            (
                row
                for row in matches
                if _text(row.get("name_norm")).casefold() == query_norm
            ),
            None,
        )
        if entity is None:
            return None
        relations = self.one_hop(
            canonical_ids=[entity.get("canonical_id")],
            names=[],
            limit=neighbor_limit,
        )
        return {"entity": entity, "relations": relations}