206 lines
8.2 KiB
Python
206 lines
8.2 KiB
Python
"""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, ResponsesOptions
|
|
from app.schemas.retrieval import RetrievalResponse
|
|
from app.schemas.vectors import VectorRetrievalResponse
|
|
from app.services.completions import canonical_json, request_body
|
|
from app.services.responses import request_body as responses_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":
|
|
build_body = (
|
|
responses_body if isinstance(value.completion, ResponsesOptions) else request_body
|
|
)
|
|
expected = build_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
|