148 lines
5.4 KiB
Python
148 lines
5.4 KiB
Python
"""Bounded chunk snapshots and serialized, optimistic publication."""
|
|
|
|
import hashlib
|
|
from uuid import UUID, uuid4, uuid5
|
|
|
|
from fastapi import Depends, HTTPException, Request
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.db.session import get_db
|
|
from app.models.catalog import utc_now
|
|
from app.models.indexing import Chunk, IndexPublication, IndexState, IndexVersion
|
|
from app.schemas.indexing import ActiveIndex, BuildIndex, IndexRead
|
|
from app.services.retrieval import corpus_fingerprint, load_corpus
|
|
|
|
MAX_CHUNKS = 20_000
|
|
|
|
|
|
def locked_state(db: Session) -> IndexState:
|
|
state = db.scalar(
|
|
select(IndexState)
|
|
.where(IndexState.id == 1)
|
|
.with_for_update()
|
|
.execution_options(populate_existing=True)
|
|
)
|
|
if state is None:
|
|
raise HTTPException(503, "索引状态尚未初始化，请检查数据库迁移")
|
|
return state
|
|
|
|
|
|
def catalog_write_guard(request: Request, db: Session = Depends(get_db)):
|
|
if request.method in {"POST", "PATCH", "PUT", "DELETE"}:
|
|
# Always before catalog/parent row locks; failures roll back this revision as well.
|
|
state = locked_state(db)
|
|
state.catalog_revision += 1
|
|
|
|
|
|
def version_or_404(db: Session, version_id: str) -> IndexVersion:
|
|
version = db.get(IndexVersion, version_id)
|
|
if version is None:
|
|
raise HTTPException(404, "索引版本不存在")
|
|
return version
|
|
|
|
|
|
def split_text(content: str, size: int, overlap: int):
|
|
"""Offsets are original Python Unicode character offsets, end-exclusive."""
|
|
for start in range(0, len(content), size - overlap):
|
|
end = min(len(content), start + size)
|
|
yield start, end, content[start:end]
|
|
if end == len(content):
|
|
break
|
|
|
|
|
|
def build_index(db: Session, payload: BuildIndex, actor_id: int) -> IndexRead:
|
|
state = locked_state(db)
|
|
corpus = load_corpus(db, None)
|
|
version_id = uuid4()
|
|
chunks = []
|
|
for item in corpus:
|
|
provenance = {key: value for key, value in item.items() if key != "content"}
|
|
for ordinal, (start, end, content) in enumerate(
|
|
split_text(item["content"], payload.chunk_size, payload.overlap)
|
|
):
|
|
if len(chunks) >= MAX_CHUNKS:
|
|
raise HTTPException(503, "切分数量超过本阶段容量，请增加切分长度或减少重叠")
|
|
chunks.append(
|
|
Chunk(
|
|
id=str(uuid5(version_id, f"{item['knowledge_id']}:{ordinal}")),
|
|
index_version_id=str(version_id),
|
|
knowledge_id=item["knowledge_id"],
|
|
ordinal=ordinal,
|
|
start_offset=start,
|
|
end_offset=end,
|
|
text=content,
|
|
text_sha256=hashlib.sha256(content.encode()).hexdigest(),
|
|
provenance=provenance,
|
|
)
|
|
)
|
|
version = IndexVersion(
|
|
id=str(version_id),
|
|
status="ready",
|
|
config={"splitter": "unicode-window-v1", **payload.model_dump()},
|
|
corpus_fingerprint=corpus_fingerprint(corpus),
|
|
catalog_revision=state.catalog_revision,
|
|
item_count=len(corpus),
|
|
chunk_count=len(chunks),
|
|
created_by=actor_id,
|
|
)
|
|
db.add(version)
|
|
db.flush()
|
|
db.add_all(chunks)
|
|
result = IndexRead.model_validate(version)
|
|
db.commit()
|
|
return result
|
|
|
|
|
|
def active_index(db: Session, state: IndexState | None = None) -> ActiveIndex:
|
|
state = state or locked_state(db)
|
|
version = db.get(IndexVersion, state.active_version_id) if state.active_version_id else None
|
|
current = bool(version and version.catalog_revision == state.catalog_revision)
|
|
if current:
|
|
current = version.corpus_fingerprint == corpus_fingerprint(load_corpus(db, None))
|
|
return ActiveIndex(
|
|
generation=state.generation,
|
|
catalog_revision=state.catalog_revision,
|
|
is_current=current,
|
|
active_version=IndexRead.model_validate(version) if version else None,
|
|
)
|
|
|
|
|
|
def publish_index(db: Session, version_id: UUID, expected_generation: int, actor_id: int):
|
|
state = locked_state(db)
|
|
if state.generation != expected_generation:
|
|
raise HTTPException(409, "已发布版本发生变化，请刷新后重试")
|
|
version = version_or_404(db, str(version_id))
|
|
if (
|
|
version.catalog_revision != state.catalog_revision
|
|
or version.corpus_fingerprint != corpus_fingerprint(load_corpus(db, None))
|
|
):
|
|
raise HTTPException(409, "资料或审核依据已变化，请重新构建后发布")
|
|
if state.active_version_id == version.id:
|
|
return active_index(db, state)
|
|
previous = state.active_version_id
|
|
if previous:
|
|
db.get(IndexVersion, previous).status = "superseded"
|
|
state.active_version_id = version.id
|
|
state.generation += 1
|
|
version.status = "published"
|
|
version.published_at = utc_now()
|
|
db.add(
|
|
IndexPublication(
|
|
generation=state.generation,
|
|
previous_version_id=previous,
|
|
version_id=version.id,
|
|
actor_id=actor_id,
|
|
published_at=version.published_at,
|
|
)
|
|
)
|
|
# Freeze the response before commit expires ORM objects and releases the lock.
|
|
result = ActiveIndex(
|
|
generation=expected_generation + 1,
|
|
catalog_revision=version.catalog_revision,
|
|
is_current=True,
|
|
active_version=IndexRead.model_validate(version),
|
|
)
|
|
db.commit()
|
|
return result
|