border-red-post-nalati/backend/app/services/exports.py

"""Read stored snapshots in one SQL statement; finish validation before any download."""

import hashlib
from typing import Annotated
from uuid import uuid4

from fastapi import HTTPException
from pydantic import Field, TypeAdapter, ValidationError
from sqlalchemy import String, literal, select, union_all
from sqlalchemy.orm import Session

from app.models.catalog import utc_now
from app.models.rag import RagContext
from app.models.retrieval import RetrievalRun
from app.schemas.exports import ExportHistory
from app.schemas.hybrid import ChunkKeywordResponse, HybridResponse
from app.schemas.rag import ContextRead
from app.schemas.retrieval import RetrievalResponse
from app.schemas.vectors import VectorRetrievalResponse
from app.services.completions import canonical_json, request_body

MAX_CONTEXTS = 100
MAX_BYTES = 8 * 1024 * 1024
SCHEMA_VERSION = "nalati-retrieval-export-v1"
RUN = TypeAdapter(
    Annotated[
        RetrievalResponse | ChunkKeywordResponse | VectorRetrievalResponse | HybridResponse,
        Field(discriminator="mode"),
    ]
)


def invalid_history():
    return HTTPException(409, "历史记录格式或关联不一致，未生成导出文件")


def same_fields(raw, parsed):
    """Reject unknown/missing fields instead of silently dropping snapshot metadata."""
    if isinstance(parsed, dict):
        return (
            isinstance(raw, dict)
            and raw.keys() == parsed.keys()
            and all(same_fields(raw[key], value) for key, value in parsed.items())
        )
    if isinstance(parsed, list):
        return (
            isinstance(raw, list)
            and len(raw) == len(parsed)
            and all(same_fields(left, right) for left, right in zip(raw, parsed))
        )
    return True


def history_rows(db: Session, ids: list[str], include_contexts: bool):
    runs = select(
        literal(0).label("record_order"),
        RetrievalRun.id.label("id"),
        RetrievalRun.id.label("run_id"),
        RetrievalRun.created_at.label("created_at"),
        RetrievalRun.query.label("query"),
        RetrievalRun.mode.label("mode"),
        RetrievalRun.corpus_fingerprint.label("corpus_fingerprint"),
        RetrievalRun.response.label("payload"),
    ).where(RetrievalRun.id.in_(ids))
    if include_contexts:
        contexts = select(
            literal(1).label("record_order"),
            RagContext.id.label("id"),
            RagContext.retrieval_run_id.label("run_id"),
            RagContext.created_at.label("created_at"),
            literal(None, String).label("query"),
            literal(None, String).label("mode"),
            literal(None, String).label("corpus_fingerprint"),
            RagContext.response.label("payload"),
        ).where(RagContext.retrieval_run_id.in_(ids))
        statement = union_all(runs, contexts)
    else:
        statement = runs
    # One PostgreSQL statement snapshot, including optional preparations. No offset
    # pagination, ORM auth rows, current corpus lookups, or model dependencies.
    return list(
        db.execute(
            statement.order_by("record_order", "created_at", "id").limit(
                len(ids) + MAX_CONTEXTS + 1
            )
        ).mappings()
    )


def checked_run(row):
    raw = row["payload"]
    try:
        value = RUN.validate_json(canonical_json(raw), strict=True)
        if (
            not same_fields(raw, value.model_dump(mode="json"))
            or str(value.run_id) != row["id"]
            or value.mode != row["mode"]
            or value.query != row["query"]
            or value.corpus_fingerprint != row["corpus_fingerprint"]
            or value.created_at != row["created_at"]
            or (value.status == "ok") != bool(value.results)
        ):
            raise ValueError("run mismatch")
    except (ValidationError, ValueError, TypeError) as exc:
        raise invalid_history() from exc
    return value


def checked_context(row, runs):
    raw = row["payload"]
    try:
        value = ContextRead.model_validate_json(canonical_json(raw), strict=True)
        run = runs[row["run_id"]]
        if (
            not same_fields(raw, value.model_dump(mode="json"))
            or str(value.id) != row["id"]
            or str(value.retrieval_run_id) != row["run_id"]
            or value.created_at != row["created_at"]
            or value.query != run.query
            or value.retrieval_mode != run.mode
            or value.corpus_fingerprint != run.corpus_fingerprint
            or value.text_index_version_id != run.text_index_version_id
            or value.text_generation != run.text_generation
            or value.catalog_revision != getattr(run, "catalog_revision", value.catalog_revision)
            or value.vector_version_id != getattr(run, "vector_version_id", None)
            or value.vector_generation != getattr(run, "generation", None)
        ):
            raise ValueError("context mismatch")
        if value.status == "prepared":
            expected = request_body(
                value.query,
                [evidence.model_dump(mode="json") for evidence in value.evidence],
                value.completion,
            )
            if (
                not value.evidence
                or canonical_json(value.request_body) != canonical_json(expected)
                or value.request_sha256
                != hashlib.sha256(canonical_json(expected).encode()).hexdigest()
            ):
                raise ValueError("request mismatch")
        elif value.request_body is not None or value.request_sha256 is not None or value.evidence:
            raise ValueError("unexpected request")
    except (ValidationError, ValueError, TypeError, KeyError, AttributeError) as exc:
        raise invalid_history() from exc


def export_history(db: Session, payload: ExportHistory):
    ids = sorted(str(value) for value in payload.run_ids)
    rows = history_rows(db, ids, payload.include_contexts)
    context_count = sum(row["record_order"] == 1 for row in rows)
    if context_count > MAX_CONTEXTS:
        raise HTTPException(413, "关联上下文超过100条，请缩小检索记录选择范围")
    run_rows = [row for row in rows if row["record_order"] == 0]
    if {row["id"] for row in run_rows} != set(ids):
        raise HTTPException(404, "部分检索记录不存在，未生成导出文件")
    runs = {row["id"]: checked_run(row) for row in run_rows}
    lines = []
    size = 0
    for row in rows:
        if row["record_order"]:
            checked_context(row, runs)
            record = {"record_type": "rag_context", "id": row["id"], "payload": row["payload"]}
        else:
            record = {
                "record_type": "retrieval_run",
                "id": row["id"],
                "ranking_unit": "knowledge_item" if row["mode"] == "keyword" else "chunk",
                "payload": row["payload"],
            }
        try:
            line = (canonical_json(record) + "\n").encode("utf-8")
        except (TypeError, ValueError, UnicodeError) as exc:
            raise invalid_history() from exc
        size += len(line)
        if size > MAX_BYTES:
            raise HTTPException(413, "导出超过8 MiB，请缩小检索记录选择范围")
        lines.append(line)
    records = b"".join(lines)
    digest = hashlib.sha256(records).hexdigest()
    export_id = str(uuid4())
    manifest = {
        "record_type": "manifest",
        "schema_version": SCHEMA_VERSION,
        "export_id": export_id,
        "exported_at": utc_now().isoformat(),
        "selection": {"run_ids": ids, "include_contexts": payload.include_contexts},
        "counts": {"retrieval_runs": len(run_rows), "rag_contexts": context_count},
        "records_sha256": digest,
        "snapshot_semantics": "stored_history_not_current_corpus",
        "generation_measurements": {
            "status": "not_called",
            "usage": None,
            "cost": None,
            "currency": None,
        },
    }
    content = (canonical_json(manifest) + "\n").encode("utf-8") + records
    if len(content) > MAX_BYTES:
        raise HTTPException(413, "导出超过8 MiB，请缩小检索记录选择范围")
    return content, export_id, digest