"""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