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

from __future__ import annotations

import json
from datetime import date, datetime
from collections.abc import Mapping
from threading import Lock
from typing import Any

from sqlalchemy import MetaData, Table

from .runtime_bridge import load_runtime_modules, new_session, runtime_dir
from .runtime_tool_catalog import RETRIEVAL_AGENT_LABELS
from .retrieval_detail_projection import (
    QUERY_TOOL_NAMES,
    summarize_retrieval_output,
)
from .sanitizer import sanitize


_DETAIL_BODY_MAX = 240_000


class BuildNotFound(LookupError):
    pass


class ActivityNotFound(LookupError):
    pass


def _json_value(value: Any) -> Any:
    if isinstance(value, (datetime, date)):
        return value.isoformat()
    return sanitize(value)


def row_dict(row: Any) -> dict[str, Any]:
    if row is None:
        return {}
    if isinstance(row, Mapping):
        return {str(key): _json_value(value) for key, value in row.items()}
    mapping = getattr(row, "_mapping", None)
    if mapping is not None:
        return {str(key): _json_value(value) for key, value in mapping.items()}
    return {
        column.name: _json_value(getattr(row, column.name))
        for column in row.__table__.columns
    }


class ScriptBuildRepository:
    """The only database adapter used by V8.

    The adapter is deliberately query-only.  Sessions come from a dedicated
    MySQL read-only pool and this class never flushes, commits or mutates ORM
    objects.
    """

    def __init__(self, session_factory=new_session):
        self._session_factory = session_factory
        self._metadata = MetaData()
        self._reflected_tables: dict[str, Table] = {}
        self._table_lock = Lock()

    def list_builds(
        self, *, limit: int = 30, status: str | None = None
    ) -> list[dict[str, Any]]:
        _, models = load_runtime_modules()
        session = self._session_factory()
        try:
            query = session.query(models.ScriptBuildRecord).filter(
                models.ScriptBuildRecord.is_deleted.is_(False)
            )
            if status:
                query = query.filter(models.ScriptBuildRecord.status == status)
            rows = (
                query.order_by(models.ScriptBuildRecord.id.desc())
                .limit(max(1, min(limit, 100)))
                .all()
            )
            return [self._build_summary(row) for row in rows]
        finally:
            session.close()

    def load_bundle(self, script_build_id: int) -> dict[str, Any]:
        _, models = load_runtime_modules()
        session = self._session_factory()
        try:
            record = (
                session.query(models.ScriptBuildRecord)
                .filter(models.ScriptBuildRecord.id == script_build_id)
                .first()
            )
            if record is None:
                raise BuildNotFound(f"未找到脚本构建 #{script_build_id}")

            rounds = (
                session.query(models.ScriptBuildRound)
                .filter(models.ScriptBuildRound.script_build_id == script_build_id)
                .order_by(models.ScriptBuildRound.round_index.asc())
                .all()
            )
            branches = (
                session.query(models.ScriptBuildBranch)
                .filter(models.ScriptBuildBranch.script_build_id == script_build_id)
                .order_by(
                    models.ScriptBuildBranch.round_index.asc(),
                    models.ScriptBuildBranch.branch_id.asc(),
                )
                .all()
            )
            data_decisions = (
                session.query(models.ScriptBuildDataDecision)
                .filter(
                    models.ScriptBuildDataDecision.script_build_id == script_build_id,
                    models.ScriptBuildDataDecision.branch_id > 0,
                )
                .order_by(
                    models.ScriptBuildDataDecision.created_at.asc(),
                    models.ScriptBuildDataDecision.id.asc(),
                )
                .all()
            )
            ignored_legacy_decisions = (
                session.query(models.ScriptBuildDataDecision)
                .filter(
                    models.ScriptBuildDataDecision.script_build_id == script_build_id,
                    models.ScriptBuildDataDecision.branch_id <= 0,
                )
                .count()
            )
            multipath_decisions = (
                session.query(models.ScriptBuildMultipathDecision)
                .filter(
                    models.ScriptBuildMultipathDecision.script_build_id
                    == script_build_id
                )
                .order_by(
                    models.ScriptBuildMultipathDecision.created_at.asc(),
                    models.ScriptBuildMultipathDecision.id.asc(),
                )
                .all()
            )
            domain_info = (
                session.query(models.ScriptBuildDomainInfo)
                .filter(models.ScriptBuildDomainInfo.script_build_id == script_build_id)
                .order_by(
                    models.ScriptBuildDomainInfo.created_at.asc(),
                    models.ScriptBuildDomainInfo.id.asc(),
                )
                .all()
            )
            event_table = self._table(session, "script_build_event")
            events = (
                session.query(event_table)
                .filter(event_table.c.script_build_id == script_build_id)
                .order_by(
                    event_table.c.event_seq.asc(),
                    event_table.c.id.asc(),
                )
                .all()
            )
            event_payloads = [self._event_summary(row) for row in events]
            self._attach_light_event_bodies(session, event_payloads)

            paragraphs = (
                session.query(models.ScriptBuildParagraph)
                .filter(models.ScriptBuildParagraph.script_build_id == script_build_id)
                .order_by(
                    models.ScriptBuildParagraph.branch_id.asc(),
                    models.ScriptBuildParagraph.id.asc(),
                )
                .all()
            )
            elements = (
                session.query(models.ScriptBuildElement)
                .filter(models.ScriptBuildElement.script_build_id == script_build_id)
                .order_by(
                    models.ScriptBuildElement.branch_id.asc(),
                    models.ScriptBuildElement.id.asc(),
                )
                .all()
            )
            links = (
                session.query(models.ScriptBuildParagraphElement)
                .filter(
                    models.ScriptBuildParagraphElement.script_build_id
                    == script_build_id
                )
                .order_by(
                    models.ScriptBuildParagraphElement.branch_id.asc(),
                    models.ScriptBuildParagraphElement.id.asc(),
                )
                .all()
            )

            branch_payloads: list[dict[str, Any]] = []
            for branch in branches:
                payload = row_dict(branch)
                payload["candidate_snapshot"] = self._candidate_snapshot(
                    branch, paragraphs=paragraphs, elements=elements, links=links
                )
                branch_payloads.append(payload)

            return {
                "record": row_dict(record),
                "rounds": [row_dict(row) for row in rounds],
                "branches": branch_payloads,
                "dataDecisions": [row_dict(row) for row in data_decisions],
                "multipathDecisions": [
                    row_dict(row) for row in multipath_decisions
                ],
                "domainInfo": [row_dict(row) for row in domain_info],
                "events": event_payloads,
                "ignoredLegacyDataDecisionCount": int(
                    ignored_legacy_decisions or 0
                ),
                "currentArtifact": self._snapshot_for_branch(
                    0, paragraphs, elements, links
                ),
            }
        finally:
            session.close()

    def load_current_artifact(self, script_build_id: int) -> dict[str, Any]:
        """Read only the active Base rows needed by the full script viewer."""
        _, models = load_runtime_modules()
        session = self._session_factory()
        try:
            exists = (
                session.query(models.ScriptBuildRecord.id)
                .filter(models.ScriptBuildRecord.id == script_build_id)
                .first()
            )
            if exists is None:
                raise BuildNotFound(f"未找到脚本构建 #{script_build_id}")

            paragraphs = (
                session.query(models.ScriptBuildParagraph)
                .filter(
                    models.ScriptBuildParagraph.script_build_id == script_build_id,
                    models.ScriptBuildParagraph.branch_id == 0,
                    models.ScriptBuildParagraph.is_active.is_not(False),
                )
                .order_by(
                    models.ScriptBuildParagraph.paragraph_index.asc(),
                    models.ScriptBuildParagraph.id.asc(),
                )
                .all()
            )
            elements = (
                session.query(models.ScriptBuildElement)
                .filter(
                    models.ScriptBuildElement.script_build_id == script_build_id,
                    models.ScriptBuildElement.branch_id == 0,
                    models.ScriptBuildElement.is_active.is_not(False),
                )
                .order_by(models.ScriptBuildElement.id.asc())
                .all()
            )
            paragraph_ids = [row.id for row in paragraphs]
            element_ids = [row.id for row in elements]
            links = []
            if paragraph_ids and element_ids:
                links = (
                    session.query(models.ScriptBuildParagraphElement)
                    .filter(
                        models.ScriptBuildParagraphElement.script_build_id
                        == script_build_id,
                        models.ScriptBuildParagraphElement.branch_id == 0,
                        models.ScriptBuildParagraphElement.paragraph_id.in_(paragraph_ids),
                        models.ScriptBuildParagraphElement.element_id.in_(element_ids),
                    )
                    .order_by(models.ScriptBuildParagraphElement.id.asc())
                    .all()
                )
            return self._snapshot_for_branch(0, paragraphs, elements, links)
        finally:
            session.close()

    def load_event_detail(
        self, script_build_id: int, event_id: int
    ) -> dict[str, Any]:
        values = self.load_event_details(script_build_id, [event_id])
        event = values.get(int(event_id))
        if event is None:
            raise ActivityNotFound(f"未找到运行事件 {event_id}")
        return event

    def load_event_details(
        self, script_build_id: int, event_ids: list[int]
    ) -> dict[int, dict[str, Any]]:
        """Batch-load full Event + Event Body records for one Inspector.

        A creative card can reference dozens of events.  Loading the rows and
        bodies in two queries keeps the three-column endpoint independent of
        that cardinality.
        """
        wanted: list[int] = []
        for value in event_ids:
            try:
                event_id = int(value)
            except (TypeError, ValueError):
                continue
            if event_id > 0 and event_id not in wanted:
                wanted.append(event_id)
        wanted.sort()
        if not wanted:
            return {}
        session = self._session_factory()
        try:
            event_table = self._table(session, "script_build_event")
            events = (
                session.query(event_table)
                .filter(
                    event_table.c.script_build_id == script_build_id,
                    event_table.c.id.in_(wanted),
                )
                .order_by(event_table.c.event_seq.asc(), event_table.c.id.asc())
                .all()
            )
            body_table = self._table(session, "script_build_event_body")
            bodies = (
                session.query(body_table)
                .filter(
                    body_table.c.script_build_id == script_build_id,
                    body_table.c.event_id.in_(wanted),
                )
                .all()
            )
            body_by_event = {
                int(parsed["eventId"]): parsed
                for row in bodies
                for parsed in [self._event_body(row)]
                if parsed.get("eventId") is not None
            }
            result: dict[int, dict[str, Any]] = {}
            for row in events:
                payload = self._event_summary(row)
                event_id = int(payload["id"])
                parsed = body_by_event.get(event_id)
                payload["input"] = parsed.get("input") if parsed else None
                payload["output"] = parsed.get("output") if parsed else None
                result[event_id] = payload
            return result
        finally:
            session.close()

    def load_prompt_context(
        self, script_build_id: int, prompt_ref: str
    ) -> dict[str, Any]:
        """Read the exact delegated task plus the *current* prompt config.

        Prompt history is intentionally not guessed from event time: runtime
        events do not store a prompt snapshot or version reference.
        """

        ref = str(prompt_ref or "").strip()
        event: dict[str, Any] | None = None
        if ref.startswith("prompt:event:") or ref.startswith("event:"):
            try:
                event_id = int(ref.rsplit(":", 1)[1])
            except ValueError as exc:
                raise ActivityNotFound(f"无效的规则引用 {prompt_ref}") from exc
            event = self.load_event_detail(script_build_id, event_id)
            event_name = str(event.get("event_name") or "")
        elif ref in {"prompt:main", "main"}:
            event_name = "main"
        else:
            raise ActivityNotFound(f"未找到规则引用 {prompt_ref}")

        config = _prompt_config(event_name)
        task = None
        if event:
            side = event.get("input")
            content = side.get("content") if isinstance(side, dict) else None
            if isinstance(content, dict):
                task = content.get("task")
            elif isinstance(content, str):
                task = content

        session = self._session_factory()
        try:
            _, models = load_runtime_modules()
            exists = (
                session.query(models.ScriptBuildRecord.id)
                .filter(models.ScriptBuildRecord.id == script_build_id)
                .first()
            )
            if exists is None:
                raise BuildNotFound(f"未找到脚本构建 #{script_build_id}")
            prompt_row = None
            for biz_type in config["bizTypes"]:
                prompt_row = (
                    session.query(models.Prompt)
                    .filter(models.Prompt.biz_type == biz_type)
                    .order_by(models.Prompt.id.asc())
                    .first()
                )
                if prompt_row is not None:
                    break
            if prompt_row is not None:
                content = prompt_row.prompt_content
                source = "current-db"
                version = prompt_row.current_version
            else:
                path = runtime_dir() / "prompts" / "script" / config["file"]
                content = path.read_text(encoding="utf-8") if path.exists() else None
                source = "current-file"
                version = None
            return {
                "promptRef": ref,
                "actor": config["actor"],
                "task": task,
                "promptContent": content,
                "promptSource": source,
                "promptVersion": version,
                "modeNotice": config.get("modeNotice"),
            }
        finally:
            session.close()

    def _table(self, session: Any, name: str) -> Table:
        table = self._reflected_tables.get(name)
        if table is not None:
            return table
        # SQLAlchemy registers a Table in MetaData before reflection has
        # populated all columns. Parallel Inspector requests must not observe
        # that half-built object.
        with self._table_lock:
            table = self._reflected_tables.get(name)
            if table is None:
                table = Table(name, self._metadata, autoload_with=session.get_bind())
                self._reflected_tables[name] = table
            return table

    @staticmethod
    def _build_summary(row: Any) -> dict[str, Any]:
        return {
            "id": str(row.id),
            "status": row.status,
            "title": f"脚本构建 #{row.id}",
            "scriptDirection": row.script_direction,
            "createdAt": _json_value(row.start_time),
            "updatedAt": _json_value(row.end_time),
        }

    @staticmethod
    def _event_summary(row: Any) -> dict[str, Any]:
        payload = row_dict(row)
        payload["detailRef"] = f"event:{payload.get('id')}"
        return payload

    @staticmethod
    def _event_body(row: Any) -> dict[str, Any]:
        mapping = getattr(row, "_mapping", None)
        data = dict(mapping) if mapping is not None else {
            column.name: getattr(row, column.name)
            for column in row.__table__.columns
        }
        return {
            "id": data.get("id"),
            "eventId": data.get("event_id"),
            "input": ScriptBuildRepository._body_side(
                data.get("input_content_type"), data.get("input_content")
            ),
            "output": ScriptBuildRepository._body_side(
                data.get("output_content_type"), data.get("output_content")
            ),
        }

    @staticmethod
    def _body_side(content_type: Any, raw: Any) -> dict[str, Any] | None:
        if raw is None:
            return None
        raw_text = str(raw)
        content: Any = raw
        parsed = False
        if content_type == "json" or raw_text.lstrip().startswith(("{", "[")):
            try:
                content = json.loads(raw_text)
                parsed = True
            except (TypeError, json.JSONDecodeError):
                content = raw
        truncated = len(raw_text) > _DETAIL_BODY_MAX
        safe_content = (
            sanitize(raw_text, max_text=_DETAIL_BODY_MAX)
            if truncated
            # A valid JSON body below the declared detail limit must not have
            # an individual string silently clipped at a second, lower limit.
            else sanitize(content, max_text=_DETAIL_BODY_MAX)
        )
        return {
            "contentType": content_type,
            "content": safe_content,
            "truncated": truncated,
            **(
                {"omittedCharacters": len(raw_text) - _DETAIL_BODY_MAX}
                if truncated
                else {}
            ),
        }

    def _attach_light_event_bodies(
        self, session: Any, events: list[dict[str, Any]]
    ) -> None:
        """Attach only the small event fields needed by the business projector.

        Tool outputs can be very large.  The execution view intentionally reads
        only tool inputs and the event's bounded output preview; full output is
        loaded by ``load_event_detail`` after the user opens one query.
        """
        agent_ids = [
            int(item["id"])
            for item in events
            if item.get("id") is not None and self._needs_light_agent_body(item)
        ]
        tool_ids = [
            int(item["id"])
            for item in events
            if item.get("id") is not None and item.get("event_type") == "tool_call"
        ]
        query_tool_ids = [
            int(item["id"])
            for item in events
            if item.get("id") is not None
            and item.get("event_type") == "tool_call"
            and str(item.get("event_name") or "") in QUERY_TOOL_NAMES
        ]
        if not agent_ids and not tool_ids:
            return
        body_table = self._table(session, "script_build_event_body")
        build_ids = {
            int(item["script_build_id"])
            for item in events
            if item.get("script_build_id") is not None
        }
        body_by_event: dict[int, dict[str, Any]] = {}
        if agent_ids:
            rows = (
                session.query(
                    body_table.c.event_id,
                    body_table.c.input_content_type,
                    body_table.c.input_content,
                    body_table.c.output_content_type,
                    body_table.c.output_content,
                )
                .filter(
                    body_table.c.event_id.in_(agent_ids),
                    body_table.c.script_build_id.in_(build_ids),
                )
                .all()
            )
            for row in rows:
                data = row_dict(row)
                body_by_event[int(data["event_id"])] = {
                    "inputData": self._light_content(
                        data.get("input_content_type"), data.get("input_content")
                    ),
                    "agentOutputData": self._light_content(
                        data.get("output_content_type"), data.get("output_content")
                    ),
                }
        if tool_ids:
            rows = (
                session.query(
                    body_table.c.event_id,
                    body_table.c.input_content_type,
                    body_table.c.input_content,
                )
                .filter(
                    body_table.c.event_id.in_(tool_ids),
                    body_table.c.script_build_id.in_(build_ids),
                )
                .all()
            )
            for row in rows:
                data = row_dict(row)
                body_by_event.setdefault(int(data["event_id"]), {})["inputData"] = self._light_content(
                    data.get("input_content_type"), data.get("input_content")
                )
        if query_tool_ids:
            rows = (
                session.query(
                    body_table.c.event_id,
                    body_table.c.output_content,
                )
                .filter(
                    body_table.c.event_id.in_(query_tool_ids),
                    body_table.c.script_build_id.in_(build_ids),
                )
                .all()
            )
            event_by_id = {
                int(item["id"]): item
                for item in events
                if item.get("id") is not None
            }
            for row in rows:
                mapping = getattr(row, "_mapping", {})
                event_id = int(mapping.get("event_id"))
                event = event_by_id.get(event_id) or {}
                body_by_event.setdefault(event_id, {})["retrievalOutcome"] = (
                    summarize_retrieval_output(
                        str(event.get("event_name") or ""),
                        raw_status=event.get("status"),
                        ended_at=event.get("ended_at"),
                        output=mapping.get("output_content"),
                    )
                )
        for event in events:
            event.update(body_by_event.get(int(event.get("id") or 0), {}))

    @staticmethod
    def _needs_light_agent_body(event: dict[str, Any]) -> bool:
        if event.get("event_type") != "agent_invoke":
            return False
        name = str(event.get("event_name") or "")
        return name in RETRIEVAL_AGENT_LABELS or name in {
            "script_implementer",
            "script_multipath_evaluator",
            "script_evaluator",
        }

    @staticmethod
    def _light_content(content_type: Any, raw: Any) -> Any:
        if raw is None:
            return None
        if content_type == "json":
            try:
                return sanitize(json.loads(str(raw)), max_text=20_000)
            except (TypeError, json.JSONDecodeError):
                return sanitize(raw, max_text=20_000)
        return sanitize(raw, max_text=20_000)

    def _candidate_snapshot(
        self,
        branch: Any,
        *,
        paragraphs: list[Any],
        elements: list[Any],
        links: list[Any],
    ) -> dict[str, Any]:
        if branch.status in {"merged", "discarded"}:
            if branch.content_snapshot:
                return {
                    "snapshot": sanitize(branch.content_snapshot),
                    "sourceOrigin": "database",
                    "snapshotKind": "saved-at-disposition",
                    "historicalAccuracy": "exact",
                    "currentProjectionAccuracy": "not-applicable",
                    "note": "分支处置前保存的候选快照",
                }
            return {
                "snapshot": None,
                "sourceOrigin": "missing",
                "snapshotKind": "missing-disposition-snapshot",
                "historicalAccuracy": "missing",
                "currentProjectionAccuracy": "not-applicable",
                "note": "终态候选未保存处置前快照，不能用当前主脚本反推当时候选。",
            }
        return {
            "snapshot": self._snapshot_for_branch(
                int(branch.branch_id), paragraphs, elements, links
            ),
            "sourceOrigin": "database",
            "snapshotKind": "current-overlay",
            "historicalAccuracy": "unknown",
            "currentProjectionAccuracy": "exact",
            "note": (
                "基于当前主脚本计算的 overlay；不是暂存时的历史快照"
                if branch.status == "parked"
                else "基于当前主脚本计算的分支 overlay"
            ),
        }

    @staticmethod
    def _snapshot_for_branch(
        branch_id: int,
        paragraphs: list[Any],
        elements: list[Any],
        links: list[Any],
    ) -> dict[str, Any]:
        base_paras = [
            row
            for row in paragraphs
            if int(row.branch_id or 0) == 0 and row.is_active is not False
        ]
        base_elems = [
            row
            for row in elements
            if int(row.branch_id or 0) == 0 and row.is_active is not False
        ]
        active_base_para_ids = {row.id for row in base_paras}
        active_base_elem_ids = {row.id for row in base_elems}
        base_links = [
            row
            for row in links
            if int(row.branch_id or 0) == 0
            and row.paragraph_id in active_base_para_ids
            and row.element_id in active_base_elem_ids
        ]
        if branch_id == 0:
            para_rows, elem_rows, link_rows = base_paras, base_elems, base_links
            origin = "database"
        else:
            branch_paras = [
                row
                for row in paragraphs
                if int(row.branch_id or 0) == branch_id
                and row.is_active is not False
            ]
            branch_elems = [
                row
                for row in elements
                if int(row.branch_id or 0) == branch_id
                and row.is_active is not False
            ]
            overridden_para_ids = {
                row.base_ref_id
                for row in branch_paras
                if row.base_ref_id is not None
            }
            overridden_elem_ids = {
                row.base_ref_id
                for row in branch_elems
                if row.base_ref_id is not None
            }
            para_rows = [
                row for row in base_paras if row.id not in overridden_para_ids
            ] + branch_paras
            elem_rows = [
                row for row in base_elems if row.id not in overridden_elem_ids
            ] + branch_elems
            visible_para_ids = {row.id for row in para_rows} | overridden_para_ids
            visible_elem_ids = {row.id for row in elem_rows} | overridden_elem_ids
            branch_links = [
                row
                for row in links
                if int(row.branch_id or 0) == branch_id
                and row.paragraph_id in visible_para_ids
                and row.element_id in visible_elem_ids
            ]
            seen = {(row.paragraph_id, row.element_id) for row in base_links}
            link_rows = base_links + [
                row
                for row in branch_links
                if (row.paragraph_id, row.element_id) not in seen
            ]
            origin = "database-overlay"
        return {
            "paragraphs": [_artifact_row(row) for row in para_rows],
            "elements": [_artifact_row(row) for row in elem_rows],
            "paragraphElements": [row_dict(row) for row in link_rows],
            "sourceOrigin": origin,
        }


def _artifact_row(row: Any) -> dict[str, Any]:
    payload = row_dict(row)
    row_id = payload.get("id")
    base_ref_id = payload.get("base_ref_id")
    payload["rowId"] = row_id
    payload["baseRefId"] = base_ref_id
    payload["effectiveId"] = base_ref_id if base_ref_id is not None else row_id
    return payload


def _prompt_config(event_name: str) -> dict[str, Any]:
    configs = {
        "main": {
            "bizTypes": ("script_build_system",),
            "file": "script_build_system_v2.md",
            "actor": {"role": "main", "label": "主 Agent"},
        },
        "script_implementer": {
            "bizTypes": ("script_implementer",),
            "file": "script_implementer.md",
            "actor": {"role": "implementer", "label": "实现 Agent"},
        },
        "script_multipath_evaluator": {
            "bizTypes": ("script_build_eval_agent", "script_build_evaluator"),
            "file": "script_build_eval_agent.md",
            "actor": {"role": "multipath-evaluator", "label": "多方案评审 Agent"},
            "modeNotice": "当前调用使用多方案对比评审模式。",
        },
        "script_evaluator": {
            "bizTypes": ("script_build_eval_agent", "script_build_evaluator"),
            "file": "script_build_eval_agent.md",
            "actor": {"role": "overall-evaluator", "label": "整体评审 Agent"},
            "modeNotice": "当前调用使用主脚本整体评审模式。",
        },
    }
    config = configs.get(str(event_name))
    if config is None:
        raise ActivityNotFound(f"该活动没有可展示的 Agent 规则：{event_name}")
    return config