script-build-and-framework-tmp/visualization-old/backend/app/retrieval_event_projection.py

from __future__ import annotations

import json
import re
from collections import Counter
from datetime import datetime, timezone
from typing import Any
from zoneinfo import ZoneInfo

from .business_detail_text import clean_report
from .runtime_event_index import (
    RuntimeEventIndex,
    event_input,
    event_order,
    event_output,
    event_summary,
    integer,
)
from .runtime_tool_catalog import (
    MAIN_DECISION_READ_TOOLS,
    RETRIEVAL_AGENT_LABELS,
    business_tool_label,
)

_AGENT_CONTEXT_TOOLS = {"think_and_plan", "get_script_snapshot"}
_DATABASE_TIMEZONE = ZoneInfo("Asia/Shanghai")


class RetrievalEventProjector:
    """Project implementer retrieval activity from an explicit event tree."""

    def project(
        self,
        index: RuntimeEventIndex,
        *,
        valid_rounds: set[int],
        valid_branches: set[int],
        captured_at: datetime,
        run_status: str | None = None,
    ) -> dict[str, Any]:
        ordered = index.ordered
        stages_by_branch: dict[int, dict[str, Any]] = {}
        unassigned: list[dict[str, Any]] = []

        implementers: list[dict[str, Any]] = []
        for event in ordered:
            if (
                str(event.get("event_type") or "") != "agent_invoke"
                or str(event.get("event_name") or "") != "script_implementer"
            ):
                continue
            round_index = integer(event.get("round_index"))
            branch_id = integer(event.get("branch_id"))
            if round_index in valid_rounds and branch_id in valid_branches:
                implementers.append(event)
            else:
                unassigned.append(event_summary(event, "missing-business-context"))

        retrieval_agents = [
            event
            for event in ordered
            if str(event.get("event_type") or "") == "agent_invoke"
            and str(event.get("event_name") or "") in RETRIEVAL_AGENT_LABELS
        ]
        attached_retrieval_ids: set[int] = set()

        for implementer in implementers:
            implementer_id = integer(implementer.get("id"))
            branch_id = integer(implementer.get("branch_id"))
            round_index = integer(implementer.get("round_index"))
            if implementer_id is None or branch_id is None or round_index is None:
                continue

            child_agents = [
                event
                for event in retrieval_agents
                if integer(event.get("parent_event_id")) == implementer_id
                and integer(event.get("branch_id")) == branch_id
                and integer(event.get("round_index")) == round_index
            ]
            attached_retrieval_ids.update(
                value
                for event in child_agents
                if (value := integer(event.get("id"))) is not None
            )
            direct_groups = _direct_tool_groups(index, implementer, run_status)
            agent_runs = [
                _retrieval_agent_run(event, index, run_status)
                for event in sorted(child_agents, key=event_order)
            ]
            operations: list[dict[str, Any]] = [*direct_groups, *agent_runs]
            waves, observed_mode = _assign_waves(operations, captured_at)
            status = _stage_status(operations, implementer, run_status)
            stage = {
                "id": f"round:{round_index}:branch:{branch_id}:retrieval",
                "roundIndex": round_index,
                "branchId": branch_id,
                "implementerEventId": implementer_id,
                "observedMode": observed_mode,
                "waves": waves,
                "directToolGroups": direct_groups,
                "agentRuns": agent_runs,
                "status": status,
                "startedAt": _iso_value(
                    min(
                        (_start(item) for item in operations if _start(item)),
                        default=None,
                    )
                ),
                "endedAt": _iso_value(
                    max(
                        (
                            _end(item, captured_at)
                            for item in operations
                            if _end(item, captured_at)
                        ),
                        default=None,
                    )
                ),
                "detailRef": f"round:{round_index}:branch:{branch_id}:retrieval",
            }
            previous = stages_by_branch.get(branch_id)
            if previous is None:
                stages_by_branch[branch_id] = stage
                continue

            # Multiple implementer invocations for one branch are rare but
            # valid. Preserve every operation instead of overwriting history.
            combined = [
                *previous.get("directToolGroups", []),
                *previous.get("agentRuns", []),
                *direct_groups,
                *agent_runs,
            ]
            waves, mode = _assign_waves(combined, captured_at)
            previous["directToolGroups"].extend(direct_groups)
            previous["agentRuns"].extend(agent_runs)
            previous["waves"] = waves
            previous["observedMode"] = mode
            previous["status"] = _combined_stage_status(previous["status"], status)

        for event in retrieval_agents:
            event_id = integer(event.get("id"))
            if event_id not in attached_retrieval_ids:
                unassigned.append(event_summary(event, "missing-parent-implementer"))

        return {
            "retrievalStagesByBranch": stages_by_branch,
            "unassigned": unassigned,
        }


def _direct_tool_groups(
    index: RuntimeEventIndex,
    implementer: dict[str, Any],
    run_status: str | None,
) -> list[dict[str, Any]]:
    implementer_id = integer(implementer.get("id"))
    if implementer_id is None:
        return []
    scoped = [
        event
        for event in index.events_in_scope(implementer_id)
        if integer(event.get("id")) != implementer_id
    ]
    groups: list[list[dict[str, Any]]] = []
    pending: list[dict[str, Any]] = []
    for event in sorted(scoped, key=event_order):
        is_direct = (
            str(event.get("event_type") or "") == "tool_call"
            and str(event.get("agent_role") or "") == "script_implementer"
            and str(event.get("event_name") or "") in MAIN_DECISION_READ_TOOLS
        )
        if is_direct:
            pending.append(event)
            continue
        if pending:
            groups.append(pending)
            pending = []
    if pending:
        groups.append(pending)
    return [_direct_group(group, run_status) for group in groups]


def _direct_group(
    events: list[dict[str, Any]], run_status: str | None
) -> dict[str, Any]:
    first_id = integer(events[0].get("id")) or 0
    calls = [_direct_call(event, run_status) for event in events]
    counts = Counter(call["status"] for call in calls)
    labels = Counter(call["businessLabel"] for call in calls)
    result = {
        "id": f"retrieval-direct:{first_id}",
        "kind": "direct-tools",
        "owner": "script_implementer",
        "callCount": len(calls),
        "successCount": counts["success"],
        "failureCount": counts["failure"],
        "runningCount": counts["running"],
        "sources": [
            {"businessLabel": label, "callCount": count}
            for label, count in labels.items()
        ],
        "calls": calls,
        "startedAt": _iso_value(
            min((_date_value(item.get("started_at")) for item in events), default=None)
        ),
        "endedAt": _iso_value(
            max((_date_value(item.get("ended_at")) for item in events), default=None)
        ),
        "detailRef": f"retrieval-direct:{first_id}",
    }
    if counts["interrupted"]:
        result["interruptedCount"] = counts["interrupted"]
    return result


def _direct_call(event: dict[str, Any], run_status: str | None) -> dict[str, Any]:
    event_id = integer(event.get("id")) or 0
    status = _operation_status(event, run_status)
    return {
        "id": f"event:{event_id}",
        "eventId": event_id,
        "businessLabel": business_tool_label(
            str(event.get("event_name") or "")
        ),
        "status": status,
        "resultSummary": (
            "读取失败"
            if status == "failure"
            else "构建中断"
            if status == "interrupted"
            else "正在读取"
            if status == "running"
            else "已读取"
        ),
        "detailRef": f"event:{event_id}",
    }


def _retrieval_agent_run(
    agent: dict[str, Any], index: RuntimeEventIndex, run_status: str | None
) -> dict[str, Any]:
    event_id = integer(agent.get("id")) or 0
    agent_type = str(agent.get("event_name") or "")
    child_tools = [
        event
        for event in index.ordered
        if str(event.get("event_type") or "") == "tool_call"
        and (
            integer(event.get("scope_event_id")) == event_id
            or integer(event.get("parent_event_id")) == event_id
        )
    ]
    attempts = [
        _query_attempt(event, run_status)
        for event in sorted(child_tools, key=event_order)
        if str(event.get("event_name") or "") not in _AGENT_CONTEXT_TOOLS
    ]
    counts = Counter(attempt["status"] for attempt in attempts)
    screening = clean_report(_agent_summary(agent))
    status = _agent_status(agent, run_status)
    result = {
        "id": f"retrieval-agent:{event_id}",
        "kind": "retrieval-agent",
        "eventId": event_id,
        "agentType": agent_type,
        "businessLabel": RETRIEVAL_AGENT_LABELS.get(agent_type, "取数 Agent"),
        "taskPreview": _preview(_agent_task(agent), 180),
        "status": status,
        "querySummary": {
            "total": len(attempts),
            "hit": counts["hit"],
            "empty": counts["empty"],
            "failure": counts["failure"],
            "unknown": counts["unknown"] + counts["running"],
        },
        "attempts": attempts,
        "screening": {
            "state": (
                "available"
                if screening
                else "running"
                if status == "running"
                else "missing"
            ),
            "preview": _preview(screening, 220),
        },
        "startedAt": _iso_value(_date_value(agent.get("started_at"))),
        "endedAt": _iso_value(_date_value(agent.get("ended_at"))),
        "durationMs": integer(agent.get("duration_ms")),
        "detailRef": f"retrieval-agent:{event_id}",
    }
    if counts["interrupted"]:
        result["querySummary"]["interrupted"] = counts["interrupted"]
    return result


def _query_attempt(event: dict[str, Any], run_status: str | None) -> dict[str, Any]:
    event_id = integer(event.get("id")) or 0
    status, count = _query_status(event, run_status)
    return {
        "id": f"event:{event_id}",
        "eventId": event_id,
        "sequence": integer(event.get("event_seq")) or 0,
        "queryLabel": _query_label(event_input(event)),
        "status": status,
        "resultCount": count,
        "detailRef": f"event:{event_id}",
    }


def _query_status(
    event: dict[str, Any], run_status: str | None = None
) -> tuple[str, int | None]:
    summary = event.get("retrievalOutcome")
    if isinstance(summary, dict) and summary.get("state"):
        state = str(summary.get("state"))
        if (
            str(run_status or "").lower()
            in {"success", "partial", "failed", "stopped", "completed", "cancelled", "cancel", "cancle"}
            and (str(event.get("status") or "").lower() == "running" or not event.get("ended_at"))
        ):
            state = "interrupted"
        return state, integer(summary.get("count"))
    raw_status = str(event.get("status") or "").lower()
    if (
        str(run_status or "").lower()
        in {"success", "partial", "failed", "stopped", "completed", "cancelled", "cancel", "cancle"}
        and (raw_status == "running" or not event.get("ended_at"))
    ):
        return "interrupted", None
    if raw_status in {"error", "failed", "failure"}:
        return "failure", None
    if raw_status == "running" or not event.get("ended_at"):
        return "running", None
    output = event_output(event)
    if isinstance(output, dict) and _output_reports_failure(output):
        return "failure", None
    count = _result_count(output)
    if count is not None:
        return ("hit" if count > 0 else "empty"), count
    text = str(event.get("output_preview") or "").strip()
    if text == "[]":
        return "empty", 0
    return "unknown", None


def _output_reports_failure(value: dict[str, Any]) -> bool:
    success = value.get("success")
    if success is False or (isinstance(success, str) and success.lower() == "false"):
        return True
    error = value.get("error")
    return error not in (None, "", False, [], {})


def _result_count(value: Any) -> int | None:
    parsed = value
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            match = re.search(r'"count"\s*:\s*(\d+)', value)
            return int(match.group(1)) if match else None
    if isinstance(parsed, list):
        return len(parsed)
    if not isinstance(parsed, dict):
        return None
    for key in ("count", "returned_count", "result_count"):
        number = integer(parsed.get(key))
        if number is not None:
            return number
    for key in ("results", "data", "items"):
        items = parsed.get(key)
        if isinstance(items, list):
            return len(items)
    return None


def _query_label(value: Any) -> str:
    if not isinstance(value, dict):
        return "查询内容详见详情"
    parts: list[str] = []
    keyword = value.get("keyword") or value.get("query") or value.get("topic_name")
    if isinstance(keyword, str) and keyword.strip():
        parts.append(keyword.strip())
    if not parts and isinstance(value.get("account_name"), str):
        parts.append(f"账号 {value['account_name'].strip()}")
    for key in ("names", "category_names"):
        names = value.get(key)
        if isinstance(names, list) and names:
            parts.append("、".join(str(item) for item in names[:4]))
    return_field = value.get("return_field")
    if isinstance(return_field, str) and return_field.strip():
        parts.append(f"返回{return_field.strip()}")
    platform = value.get("platform_channel")
    if isinstance(platform, str) and platform.strip():
        parts.append(platform.strip())
    return " · ".join(parts) or "查询内容详见详情"


def _assign_waves(
    operations: list[dict[str, Any]], captured_at: datetime
) -> tuple[list[dict[str, Any]], str]:
    if not operations:
        return [], "unknown"
    if any(_start(item) is None for item in operations):
        for item in operations:
            item["waveIndex"] = None
        return [], "unknown"
    ordered = sorted(operations, key=lambda item: (_start(item), item.get("id")))
    waves: list[dict[str, Any]] = []
    current: list[dict[str, Any]] = []
    current_end: datetime | None = None
    for operation in ordered:
        start = _start(operation)
        end = _end(operation, captured_at) or start
        if current and current_end is not None and start is not None and start >= current_end:
            waves.append(_wave(len(waves), current, current_end))
            current = []
            current_end = None
        current.append(operation)
        current_end = max(filter(None, (current_end, end)), default=None)
    if current:
        waves.append(_wave(len(waves), current, current_end))
    by_id = {item.get("id"): item for item in operations}
    for wave in waves:
        for operation_id in wave["operationIds"]:
            by_id[operation_id]["waveIndex"] = wave["index"]
    if len(operations) == 1:
        mode = "single"
    elif len(waves) == 1:
        mode = "parallel"
    elif any(len(wave["operationIds"]) > 1 for wave in waves):
        mode = "mixed"
    else:
        mode = "sequential"
    return waves, mode


def _wave(
    index: int, operations: list[dict[str, Any]], end: datetime | None
) -> dict[str, Any]:
    return {
        "index": index,
        "operationIds": [item.get("id") for item in operations],
        "startedAt": _iso_value(
            min((_start(item) for item in operations), default=None)
        ),
        "endedAt": _iso_value(end),
    }


def _stage_status(
    operations: list[dict[str, Any]],
    implementer: dict[str, Any],
    run_status: str | None = None,
) -> str:
    implementer_status = _operation_status(implementer, run_status)
    if not operations:
        return "partial" if implementer_status == "interrupted" else implementer_status if implementer_status == "running" else "missing"
    statuses = {_operation_status(item, run_status) for item in operations}
    if "running" in statuses or implementer_status == "running":
        return "running"
    if "failure" in statuses or "interrupted" in statuses or implementer_status == "interrupted":
        return "partial"
    return "completed"


def _combined_stage_status(left: str, right: str) -> str:
    if "running" in {left, right}:
        return "running"
    if "partial" in {left, right}:
        return "partial"
    if "completed" in {left, right}:
        return "completed"
    return "missing"


def _operation_status(
    value: dict[str, Any], run_status: str | None = None
) -> str:
    if int(value.get("failureCount") or 0) > 0:
        return "failure"
    if int(value.get("runningCount") or 0) > 0:
        return "running"
    status = str(value.get("status") or "").lower()
    if status == "interrupted":
        return "interrupted"
    if (
        str(run_status or "").lower()
        in {"success", "partial", "failed", "stopped", "completed", "cancelled", "cancel", "cancle"}
        and (status == "running" or not value.get("ended_at") and not value.get("endedAt"))
    ):
        return "interrupted"
    if status in {"error", "failed", "failure"}:
        return "failure"
    if status == "running" or not value.get("ended_at") and not value.get("endedAt"):
        return "running"
    return "success"


def _agent_status(value: dict[str, Any], run_status: str | None = None) -> str:
    status = _operation_status(value, run_status)
    return (
        "failed"
        if status == "failure"
        else "interrupted"
        if status == "interrupted"
        else "running"
        if status == "running"
        else "completed"
    )


def _agent_task(event: dict[str, Any]) -> str | None:
    value = event_input(event)
    if isinstance(value, dict) and isinstance(value.get("task"), str):
        return value["task"]
    return None


def _agent_summary(event: dict[str, Any]) -> str | None:
    value = event_output(event)
    if isinstance(value, dict) and isinstance(value.get("summary"), str):
        return value["summary"]
    return None


def _start(item: dict[str, Any]) -> datetime | None:
    return _date_value(item.get("startedAt") or item.get("started_at"))


def _end(item: dict[str, Any], captured_at: datetime) -> datetime | None:
    value = _date_value(item.get("endedAt") or item.get("ended_at"))
    if value is None and _operation_status(item) == "running":
        return captured_at
    return value


def _date_value(value: Any) -> datetime | None:
    if isinstance(value, datetime):
        parsed = value
        if parsed.tzinfo is None:
            parsed = parsed.replace(tzinfo=_DATABASE_TIMEZONE)
        return parsed.astimezone(timezone.utc)
    if isinstance(value, str) and value:
        try:
            parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
            if parsed.tzinfo is None:
                parsed = parsed.replace(tzinfo=_DATABASE_TIMEZONE)
            return parsed.astimezone(timezone.utc)
        except ValueError:
            return None
    return None


def _iso_value(value: datetime | None) -> str | None:
    return value.isoformat() if value is not None else None


def _preview(value: Any, limit: int) -> str | None:
    text = str(value or "").strip()
    if not text:
        return None
    text = re.sub(r"\s+", " ", text)
    return text if len(text) <= limit else text[: limit - 1].rstrip() + "…"