194 lines
7.4 KiB
Python
194 lines
7.4 KiB
Python
"""Prepare immutable, current evidence snapshots without invoking a generation model."""
|
|
|
|
import hashlib
|
|
import math
|
|
from uuid import uuid4
|
|
|
|
from fastapi import HTTPException
|
|
from pydantic import TypeAdapter, ValidationError
|
|
from sqlalchemy import select
|
|
|
|
from app.models.catalog import utc_now
|
|
from app.models.indexing import Chunk
|
|
from app.models.rag import RagContext
|
|
from app.models.retrieval import RetrievalRun
|
|
from app.models.vectors import EmbeddingSpace, VectorVersion
|
|
from app.schemas.hybrid import ChunkKeywordResponse, HybridResponse
|
|
from app.schemas.rag import ContextRead, Evidence, ExcludedEvidence, PrepareContext
|
|
from app.schemas.vectors import SpaceRead, VectorRetrievalResponse
|
|
from app.services.completions import canonical_json
|
|
from app.services.embedding import EmbeddingSpec
|
|
from app.services.indexing import locked_state
|
|
from app.services.responses import request_body
|
|
from app.services.vectors import checked_chunks, complete, current_text, vector_state
|
|
|
|
RUN = TypeAdapter(ChunkKeywordResponse | VectorRetrievalResponse | HybridResponse)
|
|
|
|
|
|
def unusable():
|
|
return HTTPException(409, "检索记录与当前审核/发布快照不一致，请重新检索")
|
|
|
|
|
|
def current_run(db, run_id):
|
|
row = db.get(RetrievalRun, str(run_id))
|
|
if row is None:
|
|
raise HTTPException(404, "检索记录不存在")
|
|
if row.mode == "keyword":
|
|
raise HTTPException(422, "上下文需要 chunk 检索记录，请使用 chunks、vector 或 hybrid 接口")
|
|
try:
|
|
run = RUN.validate_python(row.response)
|
|
except ValidationError as exc:
|
|
raise unusable() from exc
|
|
if (
|
|
row.mode != run.mode
|
|
or row.query != run.query
|
|
or row.corpus_fingerprint != run.corpus_fingerprint
|
|
or (run.status == "ok") != bool(run.results)
|
|
):
|
|
raise unusable()
|
|
state = locked_state(db)
|
|
text_version, corpus = current_text(db, state)
|
|
checked_chunks(db, text_version, corpus)
|
|
if (
|
|
str(run.text_index_version_id) != text_version.id
|
|
or run.text_generation != state.generation
|
|
or run.corpus_fingerprint != text_version.corpus_fingerprint
|
|
or (hasattr(run, "catalog_revision") and run.catalog_revision != state.catalog_revision)
|
|
):
|
|
raise unusable()
|
|
if run.mode in {"vector", "hybrid"}:
|
|
active = vector_state(db)
|
|
version = (
|
|
db.get(VectorVersion, active.active_version_id) if active.active_version_id else None
|
|
)
|
|
if (
|
|
version is None
|
|
or version.id != str(run.vector_version_id)
|
|
or active.generation != run.generation
|
|
or version.text_index_version_id != text_version.id
|
|
or version.chunk_count != text_version.chunk_count
|
|
):
|
|
raise unusable()
|
|
complete(db, version)
|
|
space = db.get(EmbeddingSpace, version.space_id)
|
|
try:
|
|
spec = EmbeddingSpec(**space.spec)
|
|
if (
|
|
space.id != spec.space_id
|
|
or space.dimension != spec.dimension
|
|
or version.dimension != space.dimension
|
|
or SpaceRead.model_validate(space) != run.config.space
|
|
):
|
|
raise ValueError("space mismatch")
|
|
except (TypeError, ValueError, AttributeError, ValidationError) as exc:
|
|
raise unusable() from exc
|
|
chunks = list(db.scalars(select(Chunk).where(Chunk.index_version_id == text_version.id)))
|
|
by_id = {chunk.id: chunk for chunk in chunks}
|
|
seen = set()
|
|
for rank, hit in enumerate(run.results, 1):
|
|
chunk = by_id.get(str(hit.chunk_id))
|
|
if (
|
|
chunk is None
|
|
or chunk.id in seen
|
|
or hit.rank != rank
|
|
or not math.isfinite(hit.score)
|
|
or hit.snippet != chunk.text
|
|
or hit.snippet_start != chunk.start_offset
|
|
or hit.snippet_end != chunk.end_offset
|
|
or hit.chunk_ordinal != chunk.ordinal
|
|
or {key: getattr(hit, key) for key in chunk.provenance} != chunk.provenance
|
|
or (run.category is not None and chunk.provenance["category"] != run.category)
|
|
):
|
|
raise unusable()
|
|
seen.add(chunk.id)
|
|
return run, state, by_id
|
|
|
|
|
|
def input_characters(body):
|
|
return len(body["instructions"]) + sum(
|
|
len(part["text"]) for message in body["input"] for part in message["content"]
|
|
)
|
|
|
|
|
|
def prepare_context(db, payload: PrepareContext, actor_id: int):
|
|
run, state, by_id = current_run(db, payload.retrieval_run_id)
|
|
selected, excluded = [], []
|
|
empty_body = request_body(run.query, [], payload.completion)
|
|
if input_characters(empty_body) > payload.max_input_characters:
|
|
raise HTTPException(422, "问题和固定提示词已超过字符预算，请增加预算")
|
|
for hit in run.results:
|
|
chunk = by_id[str(hit.chunk_id)]
|
|
reason = None
|
|
if len(selected) >= payload.max_evidence:
|
|
reason = "evidence_limit"
|
|
elif any(
|
|
e.knowledge_id == chunk.knowledge_id
|
|
and max(e.start_offset, chunk.start_offset) < min(e.end_offset, chunk.end_offset)
|
|
for e in selected
|
|
):
|
|
reason = "overlap"
|
|
evidence = Evidence(
|
|
**chunk.provenance,
|
|
citation_id=f"E{len(selected) + 1}",
|
|
chunk_id=chunk.id,
|
|
chunk_ordinal=chunk.ordinal,
|
|
retrieval_rank=hit.rank,
|
|
text=chunk.text,
|
|
text_sha256=chunk.text_sha256,
|
|
start_offset=chunk.start_offset,
|
|
end_offset=chunk.end_offset,
|
|
)
|
|
if reason is None:
|
|
body = request_body(
|
|
run.query,
|
|
[e.model_dump(mode="json") for e in [*selected, evidence]],
|
|
payload.completion,
|
|
)
|
|
if input_characters(body) > payload.max_input_characters:
|
|
reason = "budget"
|
|
if reason:
|
|
excluded.append(ExcludedEvidence(chunk_id=chunk.id, reason=reason))
|
|
else:
|
|
selected.append(evidence)
|
|
body = (
|
|
request_body(run.query, [e.model_dump(mode="json") for e in selected], payload.completion)
|
|
if selected
|
|
else None
|
|
)
|
|
response = ContextRead(
|
|
id=uuid4(),
|
|
status="prepared" if selected else "insufficient_evidence",
|
|
reason="candidates_available"
|
|
if selected
|
|
else ("context_budget" if run.results else "no_candidates"),
|
|
retrieval_run_id=run.run_id,
|
|
retrieval_mode=run.mode,
|
|
query=run.query,
|
|
text_index_version_id=run.text_index_version_id,
|
|
text_generation=state.generation,
|
|
catalog_revision=state.catalog_revision,
|
|
vector_version_id=getattr(run, "vector_version_id", None),
|
|
vector_generation=getattr(run, "generation", None),
|
|
corpus_fingerprint=run.corpus_fingerprint,
|
|
completion=payload.completion,
|
|
max_input_characters=payload.max_input_characters,
|
|
input_characters=input_characters(body) if body else 0,
|
|
max_evidence=payload.max_evidence,
|
|
evidence=selected,
|
|
excluded=excluded,
|
|
request_body=body,
|
|
request_sha256=hashlib.sha256(canonical_json(body).encode()).hexdigest() if body else None,
|
|
created_at=utc_now(),
|
|
)
|
|
db.add(
|
|
RagContext(
|
|
id=str(response.id),
|
|
retrieval_run_id=str(run.run_id),
|
|
created_by=actor_id,
|
|
response=response.model_dump(mode="json"),
|
|
created_at=response.created_at,
|
|
)
|
|
)
|
|
db.commit()
|
|
return response
|