354 lines
13 KiB
Python
354 lines
13 KiB
Python
"""Generate outside publication locks; recheck reviewed evidence before persistence/query."""
|
||
|
||
import hashlib
|
||
from dataclasses import asdict
|
||
from time import perf_counter
|
||
from uuid import UUID, uuid4
|
||
|
||
from fastapi import HTTPException
|
||
from sqlalchemy import Float, func, select
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.models.catalog import utc_now
|
||
from app.models.indexing import Chunk, IndexVersion
|
||
from app.models.retrieval import RetrievalRun
|
||
from app.models.vectors import (
|
||
ChunkVector,
|
||
EmbeddingSpace,
|
||
VectorPublication,
|
||
VectorState,
|
||
VectorVersion,
|
||
)
|
||
from app.schemas.vectors import (
|
||
ActiveVectors,
|
||
SpaceRead,
|
||
VectorHit,
|
||
VectorRead,
|
||
VectorRequest,
|
||
VectorRetrievalConfig,
|
||
VectorRetrievalResponse,
|
||
)
|
||
from app.services.embedding import EmbeddingError, EmbeddingProvider, validate_vectors
|
||
from app.services.indexing import locked_state, split_text
|
||
from app.services.retrieval import corpus_fingerprint, load_corpus
|
||
|
||
MAX_VECTOR_CHUNKS = 2000
|
||
|
||
|
||
def vector_state(db: Session) -> VectorState:
|
||
# Always acquire index_state first: the same ordering as catalog/text publication.
|
||
state = db.scalar(
|
||
select(VectorState)
|
||
.where(VectorState.id == 1)
|
||
.with_for_update()
|
||
.execution_options(populate_existing=True)
|
||
)
|
||
if state is None:
|
||
raise HTTPException(503, "向量状态未初始化，请检查迁移")
|
||
return state
|
||
|
||
|
||
def vector_or_404(db: Session, version_id: str) -> VectorVersion:
|
||
version = db.get(VectorVersion, version_id)
|
||
if version is None:
|
||
raise HTTPException(404, "向量版本不存在")
|
||
return version
|
||
|
||
|
||
def current_text(db: Session, state):
|
||
version = db.get(IndexVersion, state.active_version_id) if state.active_version_id else None
|
||
if version is None or version.catalog_revision != state.catalog_revision:
|
||
raise HTTPException(409, "尚无当前有效的文本发布版本，请构建并发布")
|
||
corpus = load_corpus(db, None)
|
||
if version.corpus_fingerprint != corpus_fingerprint(corpus):
|
||
raise HTTPException(409, "审核依据已变化，请重新构建并发布文本版本")
|
||
return version, corpus
|
||
|
||
|
||
def encode(provider: EmbeddingProvider, texts: list[str], *, query=False):
|
||
try:
|
||
spec = provider.spec
|
||
result = provider.encode(texts, query=query)
|
||
if provider.spec != spec:
|
||
raise EmbeddingError("编码期间模型身份改变")
|
||
validate_vectors(result, len(texts), spec)
|
||
return result
|
||
except HTTPException:
|
||
raise
|
||
except EmbeddingError as exc:
|
||
raise HTTPException(422, "编码校验失败，请核对文本 token 上限、向量维度和归一化") from exc
|
||
except Exception as exc:
|
||
# Optional dependency/cache failures must not expose server paths or library tracebacks.
|
||
raise HTTPException(503, "本地模型不可用，请检查可选运行依赖和固定版本缓存") from exc
|
||
|
||
|
||
def active_vectors(db: Session) -> ActiveVectors:
|
||
state = locked_state(db)
|
||
vector = vector_state(db)
|
||
version = db.get(VectorVersion, vector.active_version_id) if vector.active_version_id else None
|
||
current = False
|
||
if version and version.text_index_version_id == state.active_version_id:
|
||
try:
|
||
text_version, corpus = current_text(db, state)
|
||
checked_chunks(db, text_version, corpus)
|
||
complete(db, version)
|
||
current = version.chunk_count == text_version.chunk_count
|
||
except HTTPException as exc:
|
||
if exc.status_code != 409:
|
||
raise
|
||
return ActiveVectors(
|
||
generation=vector.generation,
|
||
text_generation=state.generation,
|
||
catalog_revision=state.catalog_revision,
|
||
is_current=current,
|
||
active_version=VectorRead.model_validate(version) if version else None,
|
||
)
|
||
|
||
|
||
def checked_chunks(db, text_version, corpus):
|
||
if text_version.chunk_count > MAX_VECTOR_CHUNKS:
|
||
raise HTTPException(503, "超过同步向量构建容量（2000 chunks）")
|
||
chunks = list(
|
||
db.scalars(
|
||
select(Chunk)
|
||
.where(Chunk.index_version_id == text_version.id)
|
||
.order_by(Chunk.knowledge_id, Chunk.ordinal)
|
||
.limit(MAX_VECTOR_CHUNKS + 1)
|
||
)
|
||
)
|
||
if len(chunks) != text_version.chunk_count:
|
||
raise HTTPException(409, "文本切分快照不完整，请重新构建")
|
||
expected = [
|
||
(
|
||
item["knowledge_id"],
|
||
ordinal,
|
||
start,
|
||
end,
|
||
body,
|
||
hashlib.sha256(body.encode()).hexdigest(),
|
||
{k: v for k, v in item.items() if k != "content"},
|
||
)
|
||
for item in corpus
|
||
for ordinal, (start, end, body) in enumerate(
|
||
split_text(
|
||
item["content"], text_version.config["chunk_size"], text_version.config["overlap"]
|
||
)
|
||
)
|
||
]
|
||
actual = [
|
||
(
|
||
c.knowledge_id,
|
||
c.ordinal,
|
||
c.start_offset,
|
||
c.end_offset,
|
||
c.text,
|
||
c.text_sha256,
|
||
c.provenance,
|
||
)
|
||
for c in chunks
|
||
]
|
||
if actual != expected:
|
||
raise HTTPException(409, "文本切分与审核语料不一致，请重新构建")
|
||
return [(c.id, c.text) for c in chunks]
|
||
|
||
|
||
def build_vectors(db: Session, text_version_id: UUID, actor_id: int, provider: EmbeddingProvider):
|
||
state = locked_state(db)
|
||
text_version, corpus = current_text(db, state)
|
||
if str(text_version_id) != text_version.id:
|
||
raise HTTPException(409, "只能对当前已发布的文本版本构建向量")
|
||
inputs = checked_chunks(db, text_version, corpus)
|
||
token = (state.generation, state.catalog_revision, text_version.id)
|
||
spec = provider.spec
|
||
db.commit() # No publication/catalog row locks or open read transaction during model work.
|
||
vectors = encode(provider, [body for _, body in inputs]) if inputs else []
|
||
if provider.spec != spec:
|
||
raise HTTPException(409, "模型身份已变化，请重试")
|
||
state = locked_state(db)
|
||
new_text, new_corpus = current_text(db, state)
|
||
if token != (state.generation, state.catalog_revision, state.active_version_id):
|
||
raise HTTPException(409, "编码期间文本或审核版本已变化，请重新构建")
|
||
if checked_chunks(db, new_text, new_corpus) != inputs:
|
||
raise HTTPException(409, "编码期间文本切分快照已变化，请重新构建")
|
||
space = db.get(EmbeddingSpace, spec.space_id)
|
||
if space is None:
|
||
db.add(EmbeddingSpace(id=spec.space_id, dimension=spec.dimension, spec=asdict(spec)))
|
||
db.flush()
|
||
elif space.spec != asdict(spec) or space.dimension != spec.dimension:
|
||
raise HTTPException(409, "已保存向量空间配置不一致")
|
||
version = VectorVersion(
|
||
id=str(uuid4()),
|
||
text_index_version_id=str(text_version_id),
|
||
space_id=spec.space_id,
|
||
dimension=spec.dimension,
|
||
chunk_count=len(inputs),
|
||
created_by=actor_id,
|
||
status="ready",
|
||
)
|
||
db.add(version)
|
||
db.flush()
|
||
db.add_all(
|
||
[
|
||
ChunkVector(
|
||
version_id=version.id,
|
||
chunk_id=chunk_id,
|
||
text_index_version_id=version.text_index_version_id,
|
||
space_id=spec.space_id,
|
||
dimension=spec.dimension,
|
||
embedding=embedding,
|
||
)
|
||
for (chunk_id, _), embedding in zip(inputs, vectors, strict=True)
|
||
]
|
||
)
|
||
result = VectorRead.model_validate(version)
|
||
db.commit()
|
||
return result
|
||
|
||
|
||
def complete(db, version):
|
||
count = db.scalar(
|
||
select(func.count()).select_from(ChunkVector).where(ChunkVector.version_id == version.id)
|
||
)
|
||
if count != version.chunk_count:
|
||
raise HTTPException(409, "向量快照不完整，请重新构建")
|
||
|
||
|
||
def publish_vectors(db: Session, version_id: UUID, expected_generation: int, actor_id: int):
|
||
state = locked_state(db)
|
||
vector = vector_state(db)
|
||
if vector.generation != expected_generation:
|
||
raise HTTPException(409, "向量发布代数已变化，请刷新后重试")
|
||
version = vector_or_404(db, str(version_id))
|
||
text_version, corpus = current_text(db, state)
|
||
if version.text_index_version_id != text_version.id:
|
||
raise HTTPException(409, "该向量版本不属于当前文本发布版本")
|
||
checked_chunks(db, text_version, corpus)
|
||
complete(db, version)
|
||
if version.chunk_count != text_version.chunk_count:
|
||
raise HTTPException(409, "向量与文本切分数量不一致")
|
||
if vector.active_version_id == version.id:
|
||
return active_vectors(db)
|
||
previous = vector.active_version_id
|
||
if previous:
|
||
db.get(VectorVersion, previous).status = "superseded"
|
||
vector.active_version_id = version.id
|
||
vector.generation += 1
|
||
version.status, version.published_at = "published", utc_now()
|
||
db.add(
|
||
VectorPublication(
|
||
generation=vector.generation,
|
||
previous_version_id=previous,
|
||
version_id=version.id,
|
||
actor_id=actor_id,
|
||
published_at=version.published_at,
|
||
)
|
||
)
|
||
result = ActiveVectors(
|
||
generation=vector.generation,
|
||
text_generation=state.generation,
|
||
catalog_revision=state.catalog_revision,
|
||
is_current=True,
|
||
active_version=VectorRead.model_validate(version),
|
||
)
|
||
db.commit()
|
||
return result
|
||
|
||
|
||
def retrieval_scope(db, provider):
|
||
state = locked_state(db)
|
||
vector = vector_state(db)
|
||
text_version, corpus = current_text(db, state)
|
||
version = db.get(VectorVersion, vector.active_version_id) if vector.active_version_id else None
|
||
if version is None or version.text_index_version_id != text_version.id:
|
||
raise HTTPException(409, "尚无匹配当前文本版本的向量发布，请先构建并发布")
|
||
checked_chunks(db, text_version, corpus)
|
||
space = db.get(EmbeddingSpace, version.space_id)
|
||
if (
|
||
space.id != provider.spec.space_id
|
||
or space.spec != asdict(provider.spec)
|
||
or space.dimension != provider.spec.dimension
|
||
or version.dimension != provider.spec.dimension
|
||
):
|
||
raise HTTPException(
|
||
409, "当前编码器与已发布向量空间不一致，请使用同一配置重建或切换运行配置"
|
||
)
|
||
complete(db, version)
|
||
if version.chunk_count != text_version.chunk_count:
|
||
raise HTTPException(409, "向量与文本切分数量不一致")
|
||
token = (state.generation, state.catalog_revision, vector.generation, version.id)
|
||
return state, vector, text_version, version, space, token
|
||
|
||
|
||
def retrieve_vectors(db: Session, payload: VectorRequest, provider: EmbeddingProvider):
|
||
started = perf_counter()
|
||
*_, token = retrieval_scope(db, provider)
|
||
db.commit()
|
||
query_vector = encode(provider, [payload.query], query=True)[0]
|
||
state, vector, text_version, version, space, current = retrieval_scope(db, provider)
|
||
if token != current:
|
||
raise HTTPException(409, "查询编码期间发布版本或审核依据已变化，请重试")
|
||
# Extension-qualified operator works even with an isolated business schema search_path.
|
||
distance = ChunkVector.embedding.op("OPERATOR(public.<=>)", return_type=Float)(query_vector)
|
||
statement = (
|
||
select(Chunk, distance.label("distance"))
|
||
.join(ChunkVector, ChunkVector.chunk_id == Chunk.id)
|
||
.where(
|
||
ChunkVector.version_id == version.id,
|
||
ChunkVector.space_id == space.id,
|
||
ChunkVector.text_index_version_id == text_version.id,
|
||
ChunkVector.dimension == space.dimension,
|
||
)
|
||
)
|
||
if payload.category is not None:
|
||
statement = statement.where(Chunk.provenance["category"].as_string() == payload.category)
|
||
total = db.scalar(select(func.count()).select_from(statement.subquery()))
|
||
rows = db.execute(
|
||
statement.order_by(distance, Chunk.knowledge_id, Chunk.ordinal, Chunk.id).limit(
|
||
payload.top_k
|
||
)
|
||
)
|
||
hits = [
|
||
VectorHit(
|
||
**chunk.provenance,
|
||
rank=rank,
|
||
chunk_id=chunk.id,
|
||
chunk_ordinal=chunk.ordinal,
|
||
snippet=chunk.text,
|
||
snippet_start=chunk.start_offset,
|
||
snippet_end=chunk.end_offset,
|
||
score=round(max(-1.0, min(1.0, 1.0 - float(dist))), 8),
|
||
matched_terms=[],
|
||
)
|
||
for rank, (chunk, dist) in enumerate(rows, 1)
|
||
]
|
||
result = VectorRetrievalResponse(
|
||
run_id=uuid4(),
|
||
query=payload.query,
|
||
status="ok" if hits else "no_evidence",
|
||
top_k=payload.top_k,
|
||
category=payload.category,
|
||
vector_version_id=version.id,
|
||
text_index_version_id=text_version.id,
|
||
generation=vector.generation,
|
||
text_generation=state.generation,
|
||
corpus_fingerprint=text_version.corpus_fingerprint,
|
||
corpus_size=total,
|
||
total_hits=total,
|
||
results=hits,
|
||
elapsed_ms=round((perf_counter() - started) * 1000, 3),
|
||
created_at=utc_now(),
|
||
config=VectorRetrievalConfig(space=SpaceRead.model_validate(space)),
|
||
)
|
||
db.add(
|
||
RetrievalRun(
|
||
id=str(result.run_id),
|
||
query=result.query,
|
||
mode="vector",
|
||
corpus_fingerprint=result.corpus_fingerprint,
|
||
response=result.model_dump(mode="json"),
|
||
created_at=result.created_at,
|
||
)
|
||
)
|
||
db.commit()
|
||
return result
|