DataOps-platform/app/core/knowledge/retrieval/device.py

from __future__ import annotations

import re
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import Protocol

from sqlalchemy import text

from app.core.knowledge.access import KnowledgeAccessContext
from app.core.knowledge.retrieval.contracts import KnowledgeEvidence

MAX_DEVICE_QUERY_LENGTH = 300
MAX_DEVICE_RESULTS = 100

_EVENT_LABELS = {
    "fault": "故障",
    "alarm": "告警",
    "maintenance": "维护",
    "downtime": "停机",
}
_EVENT_ORDER = {
    event_type: position for position, event_type in enumerate(_EVENT_LABELS)
}
_LABELED_IDENTIFIER = re.compile(
    r"(?:源\s*ID|平台\s*UID)\s*(?:为|是|[:：])?\s*"
    r"([0-9A-Za-z][0-9A-Za-z._:-]{1,})",
    re.IGNORECASE,
)
_QUESTION_ENTITY_PATTERNS = tuple(
    re.compile(pattern)
    for pattern in (
        r"^(?:请问|查询|查找|搜索)?\s*(.+?)\s*"
        r"(?:的)?(?:责任人|负责人)(?:是)?谁[?？。]?$",
        r"^(?:请问|查询|查找|搜索)?\s*(.+?)\s*"
        r"(?:最近)?\s*有哪些\s*(?:故障|告警|维护|停机)"
        r"(?:记录|事件)?[?？。]?$",
        r"^(?:请问|查询|查找|搜索)?\s*(.+?)\s*"
        r"负责哪些设备[?？。]?$",
        r"^(?:请问|查询|查找|搜索)?\s*(.+?)\s*"
        r"有哪些设备[?？。]?$",
        r"^(?:请问|查询|查找|搜索)?\s*(.+?)\s*"
        r"(?:位于哪里|在哪里)[?？。]?$",
    )
)


@dataclass(frozen=True)
class DeviceSearchRow:
    asset_uid: str
    asset_type: str
    name: str
    current_version: int
    location: str | None
    organization: str | None
    responsible_person: str | None
    source_codes: tuple[str, ...]
    related_events: tuple[tuple[str, str, str | None], ...]
    business_domain_uid: str | None
    updated_at: datetime | None
    rank: float


class DeviceKnowledgeRepository(Protocol):
    def search(
        self,
        *,
        query: str,
        global_access: bool,
        business_domain_uids: tuple[str, ...],
        limit: int,
    ) -> Sequence[DeviceSearchRow]: ...


def normalize_device_query(query: str) -> str:
    normalized = query.strip()
    if not normalized:
        raise ValueError("设备检索词不能为空")
    if len(normalized) > MAX_DEVICE_QUERY_LENGTH:
        raise ValueError(
            f"设备检索词长度不能超过 {MAX_DEVICE_QUERY_LENGTH} 个字符"
        )
    return normalized


def device_query_terms(query: str) -> tuple[str, ...]:
    normalized = normalize_device_query(query)
    candidates: list[str] = []

    identifier = _LABELED_IDENTIFIER.search(normalized)
    if identifier:
        candidates.append(identifier.group(1))

    for pattern in _QUESTION_ENTITY_PATTERNS:
        match = pattern.match(normalized)
        if match:
            candidate = match.group(1).strip(" \t\r\n，,。；;：:")
            if len(candidate) >= 2:
                candidates.append(candidate)
            break

    candidates.append(normalized)
    return tuple(dict.fromkeys(candidates))


def _clean_values(values: Sequence[str]) -> tuple[str, ...]:
    return tuple(sorted({value.strip() for value in values if value.strip()}))


def _clean_events(
    values: Sequence[tuple[str, str, str | None]],
) -> tuple[tuple[str, str, str | None], ...]:
    normalized = {
        (
            event_type.strip().lower(),
            title.strip(),
            source_code.strip() or None if source_code else None,
        )
        for event_type, title, source_code in values
        if event_type.strip().lower() in _EVENT_LABELS and title.strip()
    }
    return tuple(
        sorted(
            normalized,
            key=lambda item: (
                _EVENT_ORDER[item[0]],
                item[1],
                item[2] or "",
            ),
        )
    )


def build_device_evidence(row: DeviceSearchRow) -> KnowledgeEvidence:
    source_codes = _clean_values(row.source_codes)
    events = _clean_events(row.related_events)
    lines = [
        f"设备名称：{row.name}",
        f"平台 UID：{row.asset_uid}",
    ]
    if source_codes:
        lines.append(f"源 ID：{'、'.join(source_codes)}")
    if row.location:
        lines.append(f"位置：{row.location}")
    if row.organization:
        lines.append(f"所属组织：{row.organization}")
    if row.responsible_person:
        lines.append(f"责任人：{row.responsible_person}")
    for event_type, title, source_code in events:
        detail = f"{title}（{source_code}）" if source_code else title
        lines.append(f"{_EVENT_LABELS[event_type]}：{detail}")

    point_key = f"DeviceAsset/{row.asset_uid}/summary"
    return KnowledgeEvidence(
        chunk_id=f"device:{row.asset_uid}:v{row.current_version}",
        content="\n".join(lines),
        score=float(row.rank),
        retriever="device",
        object_uid=row.asset_uid,
        object_type="DeviceAsset",
        object_version=row.current_version,
        business_domain_uid=row.business_domain_uid,
        point_keys=(point_key,),
        point_revisions=(row.current_version,),
        generation=row.current_version,
        source_updated_at=(
            row.updated_at.isoformat() if row.updated_at is not None else None
        ),
    )


class SqlDeviceKnowledgeRetriever:
    def __init__(self, repository: DeviceKnowledgeRepository):
        self._repository = repository

    def retrieve(
        self,
        query: str,
        context: KnowledgeAccessContext,
        limit: int,
    ) -> tuple[KnowledgeEvidence, ...]:
        normalized_query = normalize_device_query(query)
        bounded_limit = max(1, min(int(limit), MAX_DEVICE_RESULTS))
        rows = self._repository.search(
            query=normalized_query,
            global_access=context.global_access,
            business_domain_uids=tuple(sorted(context.business_domain_uids)),
            limit=bounded_limit,
        )
        return tuple(build_device_evidence(row) for row in rows)


_AUTHORIZED_DEVICE_CTES = """
    WITH authorized_sources AS (
        SELECT source.uid,
               CASE
                   WHEN :global_access THEN NULL::uuid
                   ELSE MIN(scope.domain_uid)::uuid
               END AS business_domain_uid
        FROM public.ingestion_sources source
        LEFT JOIN LATERAL (
            SELECT value AS domain_uid
            FROM jsonb_array_elements_text(
                CASE
                    WHEN jsonb_typeof(
                        source.permission_scope -> 'business_domains'
                    ) = 'array'
                    THEN source.permission_scope -> 'business_domains'
                    ELSE '[]'::jsonb
                END
            ) values(value)
        ) scope ON TRUE
        WHERE source.status = 'active'
          AND (
              :global_access
              OR scope.domain_uid = ANY(CAST(:domains AS text[]))
          )
        GROUP BY source.uid
    ),
    authorized_mappings AS (
        SELECT mapping.asset_uid,
               MIN(
                   authorized.business_domain_uid::text
               )::uuid AS business_domain_uid,
               array_agg(
                   DISTINCT mapping.source_code
                   ORDER BY mapping.source_code
               ) AS source_codes,
               MAX(
                   COALESCE(
                       mapping.source_updated_at,
                       mapping.last_seen_at
                   )
               ) AS source_updated_at
        FROM public.device_asset_source_mappings mapping
        JOIN authorized_sources authorized
          ON authorized.uid = mapping.source_uid
        GROUP BY mapping.asset_uid
    ),
    authorized_events AS (
        SELECT event.asset_uid,
               jsonb_agg(
                   DISTINCT jsonb_build_object(
                       'event_type', event.event_type,
                       'title', event.title,
                       'source_code', event.source_code
                   )
               ) AS related_events,
               MAX(event.occurred_at) AS event_updated_at
        FROM public.device_operational_events event
        JOIN authorized_sources authorized
          ON authorized.uid = event.source_uid
        GROUP BY event.asset_uid
    )
"""

_DEVICE_SELECT = """
    SELECT asset.uid AS asset_uid,
           asset.asset_type,
           asset.name,
           asset.current_version,
           asset.location,
           asset.organization,
           asset.responsible_person,
           mapping.source_codes,
           COALESCE(events.related_events, '[]'::jsonb) AS related_events,
           mapping.business_domain_uid,
           GREATEST(
               asset.updated_at,
               mapping.source_updated_at,
               events.event_updated_at
           ) AS updated_at,
           {rank} AS rank
    FROM public.device_assets asset
    JOIN authorized_mappings mapping
      ON mapping.asset_uid = asset.uid
    LEFT JOIN authorized_events events
      ON events.asset_uid = asset.uid
    WHERE asset.status = 'active'
      {predicate}
"""


def _device_search_row(row) -> DeviceSearchRow:
    events = tuple(
        (
            str(event["event_type"]),
            str(event["title"]),
            (
                str(event["source_code"])
                if event.get("source_code") is not None
                else None
            ),
        )
        for event in (row["related_events"] or ())
    )
    return DeviceSearchRow(
        asset_uid=str(row["asset_uid"]),
        asset_type=str(row["asset_type"]),
        name=str(row["name"]),
        current_version=int(row["current_version"]),
        location=row["location"],
        organization=row["organization"],
        responsible_person=row["responsible_person"],
        source_codes=tuple(str(value) for value in row["source_codes"]),
        related_events=_clean_events(events),
        business_domain_uid=(
            str(row["business_domain_uid"])
            if row["business_domain_uid"] is not None
            else None
        ),
        updated_at=row["updated_at"],
        rank=float(row["rank"]),
    )


class SqlDeviceKnowledgeRepository:
    def __init__(self, session):
        self._session = session

    @staticmethod
    def _access_parameters(
        *,
        global_access: bool,
        business_domain_uids: tuple[str, ...],
    ) -> dict[str, object]:
        return {
            "global_access": bool(global_access),
            "domains": list(business_domain_uids),
        }

    def search(
        self,
        *,
        query: str,
        global_access: bool,
        business_domain_uids: tuple[str, ...],
        limit: int,
    ) -> tuple[DeviceSearchRow, ...]:
        statement = text(
            _AUTHORIZED_DEVICE_CTES
            + _DEVICE_SELECT.format(
                rank="""
                    CASE
                        WHEN EXISTS (
                            SELECT 1
                            FROM unnest(CAST(:terms AS text[])) term(value)
                            WHERE lower(asset.uid::text) = lower(term.value)
                        ) THEN 1.0
                        WHEN EXISTS (
                            SELECT 1
                            FROM unnest(CAST(:terms AS text[])) term(value)
                            WHERE lower(asset.name) = lower(term.value)
                        ) THEN 0.98
                        WHEN EXISTS (
                            SELECT 1
                            FROM unnest(CAST(:terms AS text[])) term(value)
                            JOIN unnest(mapping.source_codes) code
                              ON lower(code) = lower(term.value)
                        ) THEN 0.96
                        WHEN EXISTS (
                            SELECT 1
                            FROM unnest(CAST(:terms AS text[])) term(value)
                            WHERE left(
                                lower(asset.name),
                                char_length(term.value)
                            ) = lower(term.value)
                        ) THEN 0.90
                        ELSE 0.80
                    END
                """,
                predicate="""
                    AND EXISTS (
                        SELECT 1
                        FROM unnest(CAST(:terms AS text[])) term(value)
                        WHERE
                            strpos(
                                lower(asset.uid::text),
                                lower(term.value)
                            ) > 0
                            OR strpos(
                                lower(asset.name),
                                lower(term.value)
                            ) > 0
                            OR strpos(
                                lower(COALESCE(asset.location, '')),
                                lower(term.value)
                            ) > 0
                            OR strpos(
                                lower(COALESCE(asset.organization, '')),
                                lower(term.value)
                            ) > 0
                            OR strpos(
                                lower(COALESCE(asset.responsible_person, '')),
                                lower(term.value)
                            ) > 0
                            OR EXISTS (
                                SELECT 1
                                FROM unnest(mapping.source_codes) code
                                WHERE strpos(
                                    lower(code),
                                    lower(term.value)
                                ) > 0
                            )
                            OR EXISTS (
                                SELECT 1
                                FROM jsonb_array_elements(
                                    COALESCE(
                                        events.related_events,
                                        '[]'::jsonb
                                    )
                                ) event
                                WHERE strpos(
                                    lower(COALESCE(event ->> 'title', '')),
                                    lower(term.value)
                                ) > 0
                                   OR strpos(
                                       lower(
                                           COALESCE(
                                               event ->> 'source_code',
                                               ''
                                           )
                                       ),
                                       lower(term.value)
                                   ) > 0
                            )
                    )
                    ORDER BY rank DESC, asset.updated_at DESC, asset.uid
                    LIMIT :limit
                """,
            )
        )
        parameters = self._access_parameters(
            global_access=global_access,
            business_domain_uids=business_domain_uids,
        )
        parameters.update(
            {
                "terms": list(device_query_terms(query)),
                "limit": max(1, min(int(limit), MAX_DEVICE_RESULTS)),
            }
        )
        rows = self._session.execute(statement, parameters).mappings()
        return tuple(_device_search_row(row) for row in rows)

    def get_detail(
        self,
        asset_uid: str,
        *,
        global_access: bool,
        business_domain_uids: tuple[str, ...],
    ) -> DeviceSearchRow | None:
        statement = text(
            _AUTHORIZED_DEVICE_CTES
            + _DEVICE_SELECT.format(
                rank="1.0",
                predicate="""
                    AND asset.uid = CAST(:asset_uid AS uuid)
                    LIMIT 1
                """,
            )
        )
        parameters = self._access_parameters(
            global_access=global_access,
            business_domain_uids=business_domain_uids,
        )
        parameters["asset_uid"] = str(asset_uid)
        row = self._session.execute(statement, parameters).mappings().first()
        return _device_search_row(row) if row is not None else None