fix(retrieval): 按索引策略重建并融合跨空间检索

This commit is contained in:
2026-09-04 19:48:41 +08:00
parent d1cffb2f40
commit 6661f4f8c6
5 changed files with 156 additions and 5 deletions
+4
View File
@@ -127,6 +127,10 @@ MIGRATIONS: list[str] = [
query TEXT NOT NULL UNIQUE query TEXT NOT NULL UNIQUE
); );
""", """,
# v6: persist each block's embedding policy for partitioned retrieval.
"""
ALTER TABLE blocks ADD COLUMN embedding_local_only INTEGER NOT NULL DEFAULT 0;
""",
] ]
+62
View File
@@ -22,6 +22,7 @@ from app.database.db import connect, transaction
from app.errors import ApiError from app.errors import ApiError
from app.retrieval.vectorstore import VectorHit from app.retrieval.vectorstore import VectorHit
from app.retrieval.provenance import record_embedding from app.retrieval.provenance import record_embedding
from app.retrieval.hybrid import rrf_fuse
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -167,9 +168,18 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa
Read coverage and vectors together so concurrent note updates cannot produce Read coverage and vectors together so concurrent note updates cannot produce
an apparently complete subset. Never fill missing remote hits with local hits. an apparently complete subset. Never fill missing remote hits with local hits.
""" """
if accept_local:
conn = connect()
try:
policies = {bool(row[0]) for row in conn.execute("SELECT DISTINCT embedding_local_only FROM blocks")}
finally:
conn.close()
if True in policies:
return await _search_partitioned(query, policies, top_k=top_k, strict=strict)
batch = await embed_remote([query], accept_local=accept_local, strict=strict) batch = await embed_remote([query], accept_local=accept_local, strict=strict)
if batch is None: if batch is None:
return None return None
record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions}) record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
try: try:
conn = connect() conn = connect()
@@ -221,3 +231,55 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa
"Embedding 已可用,但当前模型的向量索引缺失、不完整或已失效。请在「设置 → 索引与模型」中重建全部索引。", "Embedding 已可用,但当前模型的向量索引缺失、不完整或已失效。请在「设置 → 索引与模型」中重建全部索引。",
{"model_id": batch.space_id, "dimensions": batch.dimensions, "source": batch.source}) from exc {"model_id": batch.space_id, "dimensions": batch.dimensions, "source": batch.source}) from exc
return None return None
async def _search_partitioned(query: str, policies: set[bool], *, top_k: int, strict: bool):
"""Embed per policy; rank each space independently and fuse ranks, not vectors."""
batches = {}
for policy in sorted(policies):
batch = await embed_remote([query], accept_local=True, strict=strict, local_only=policy)
if batch is None:
return None
batches[policy] = batch
conn = connect()
try:
with transaction(conn):
# Query vectors are ready before opening the single read snapshot.
current = {bool(row[0]) for row in conn.execute("SELECT DISTINCT embedding_local_only FROM blocks")}
if current != policies:
raise ValueError("embedding policies changed while querying")
ranked = []
for policy, batch in batches.items():
rows = conn.execute(
"SELECT b.block_id,r.vector FROM blocks b LEFT JOIN routed_block_vectors r "
"ON r.block_id=b.block_id AND r.space_id=? AND r.dimensions=? "
"WHERE b.embedding_local_only=? ORDER BY b.block_id",
(batch.space_id, batch.dimensions, int(policy)),
)
def hits():
for row in rows:
if row['vector'] is None:
raise ValueError("incomplete policy coverage")
vector = _unit_vector(json.loads(row['vector']), batch.dimensions)
score = math.fsum(a * b for a, b in zip(batch.vectors[0], vector))
yield VectorHit(id=row['block_id'], score=max(0.0, min(1.0, score)))
try:
ranked.append(heapq.nlargest(top_k, hits(), key=lambda hit: hit.score))
finally:
rows.close()
spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions,
"local_only": policy} for policy, b in batches.items()]
record_embedding(source="mixed" if len({b.source for b in batches.values()}) > 1 else batch.source,
spaces=spaces, fallback_reason=None)
if len(ranked) == 1:
return ranked[0]
fused = rrf_fuse([[hit.id for hit in group] for group in ranked])
return [VectorHit(id=key, score=score) for key, score in
sorted(fused.items(), key=lambda item: (-item[1], item[0]))[:top_k]]
except Exception as exc:
record_embedding(source="unavailable", fallback_reason="REMOTE_INDEX_UNAVAILABLE")
if strict:
raise ApiError(409, "SEMANTIC_INDEX_UNAVAILABLE", "部分索引分区缺失或已失效,请重建全部索引。") from exc
return None
finally:
conn.close()
+6 -5
View File
@@ -85,7 +85,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
)) ))
try: try:
prepared_notes = [] prepared_notes = []
semantic_space = None semantic_spaces = {}
for rel, folder, markdown, created, updated in docs: for rel, folder, markdown, created, updated in docs:
parsed = parse_note( parsed = parse_note(
markdown=markdown, file_path=rel, folder=folder, tags=None, markdown=markdown, file_path=rel, folder=folder, tags=None,
@@ -97,9 +97,10 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
if batch is None: if batch is None:
raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 未生成向量,重建已停止,原索引已保留。") raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 未生成向量,重建已停止,原索引已保留。")
space = (batch.space_id, batch.dimensions) space = (batch.space_id, batch.dimensions)
if semantic_space is not None and semantic_space != space: policy = parsed.embedding_local_only
if policy in semantic_spaces and semantic_spaces[policy] != space:
raise ApiError(409, "EMBEDDING_SPACE_CHANGED", "重建期间 Embedding 模型发生切换,原索引已保留,请待模型服务稳定后重试。") raise ApiError(409, "EMBEDDING_SPACE_CHANGED", "重建期间 Embedding 模型发生切换,原索引已保留,请待模型服务稳定后重试。")
semantic_space = space semantic_spaces[policy] = space
prepared_notes.append((parsed, prepared)) prepared_notes.append((parsed, prepared))
# All network/model awaits precede the transaction. The concrete SQLite # All network/model awaits precede the transaction. The concrete SQLite
# methods below complete synchronously despite their async interfaces. # methods below complete synchronously despite their async interfaces.
@@ -114,12 +115,12 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
await vector_store.clear(conn=conn) await vector_store.clear(conn=conn)
for parsed, prepared in prepared_notes: for parsed, prepared in prepared_notes:
await index_note(parsed, prepared=prepared, conn=conn) await index_note(parsed, prepared=prepared, conn=conn)
if semantic_space is not None: for policy, space in semantic_spaces.items():
exists = conn.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='routed_block_vectors'").fetchone() exists = conn.execute("SELECT 1 FROM sqlite_master WHERE type='table' AND name='routed_block_vectors'").fetchone()
missing = not exists or conn.execute( missing = not exists or conn.execute(
"SELECT 1 FROM blocks b LEFT JOIN routed_block_vectors r " "SELECT 1 FROM blocks b LEFT JOIN routed_block_vectors r "
"ON r.block_id=b.block_id AND r.space_id=? AND r.dimensions=? " "ON r.block_id=b.block_id AND r.space_id=? AND r.dimensions=? "
"WHERE r.block_id IS NULL LIMIT 1", semantic_space, "WHERE b.embedding_local_only=? AND r.block_id IS NULL LIMIT 1", (*space, int(policy)),
).fetchone() ).fetchone()
if missing: if missing:
raise ApiError(500, "SEMANTIC_INDEX_WRITE_FAILED", "向量索引写入失败,原索引已保留,请检查数据库和磁盘状态。") raise ApiError(500, "SEMANTIC_INDEX_WRITE_FAILED", "向量索引写入失败,原索引已保留,请检查数据库和磁盘状态。")
+2
View File
@@ -118,6 +118,8 @@ async def index_note(
blocks=parsed.blocks, blocks=parsed.blocks,
) )
old_ids = set(old_block_ids) old_ids = set(old_block_ids)
conn.execute("UPDATE blocks SET embedding_local_only=? WHERE note_id=?",
(int(parsed.embedding_local_only), parsed.note_id))
new_ids = {block.block_id for block in parsed.blocks} new_ids = {block.block_id for block in parsed.blocks}
stale_ids = [bid for bid in old_ids if bid not in new_ids] stale_ids = [bid for bid in old_ids if bid not in new_ids]
if stale_ids: if stale_ids:
+82
View File
@@ -537,3 +537,85 @@ def test_real_embedding_rebuild_failure_preserves_index(runtime, production_engi
def test_empty_vault_vector_search_returns_empty(runtime, production_engine): def test_empty_vault_vector_search_returns_empty(runtime, production_engine):
assert asyncio.run(production_engine.search(request())).items == [] assert asyncio.run(production_engine.search(request())).items == []
@pytest.fixture
def policy_runtime(monkeypatch):
class PolicyRuntime:
fallback = False
calls = []
async def embed(self, texts, *, local_only=False):
self.calls.append((list(texts), local_only))
local = local_only or self.fallback
dim = 3 if local else 2
return SimpleNamespace(source='local' if local else 'api', model_id='local-space' if local else 'api-space',
dimensions=dim, vectors=[[1.0] + [0.0] * (dim - 1) for _ in texts],
fallback_reason='PROVIDER_TIMEOUT' if self.fallback and not local_only else None)
runtime = PolicyRuntime()
monkeypatch.setattr(routed_vectors, 'get_model_routing', lambda: runtime)
return runtime
async def seed_policies():
normal = await note_service.create_note(title='Normal', markdown='apple public', folder=None, tags=[])
private = await note_service.create_note(title='Private', markdown='---\nembedding_local_only: true\n---\napple private', folder=None, tags=[])
return normal, private
@pytest.mark.parametrize('fallback', [False, True])
def test_mixed_policy_rebuild_and_retrieval(policy_runtime, production_engine, fallback):
policy_runtime.fallback = fallback
async def scenario():
notes = await seed_policies()
await index_service.rebuild(IndexRebuildRequest())
for mode in (SearchMode.vector, SearchMode.hybrid):
result = await production_engine.search(SearchRequest(query='apple', mode=mode))
assert {item.note_id for item in result.items} == {note.note_id for note in notes}
for texts, local_only in policy_runtime.calls:
if any('private' in text for text in texts):
assert local_only
if not fallback:
assert {r[0] for r in rows('SELECT DISTINCT space_id FROM routed_block_vectors')} == {'api-space', 'local-space'}
asyncio.run(scenario())
def test_local_only_vault_never_requests_api_for_search(policy_runtime, production_engine):
async def scenario():
await note_service.create_note(title='Private', markdown='---\nembedding_local_only: true\n---\napple private', folder=None, tags=[])
await index_service.rebuild(IndexRebuildRequest())
assert (await production_engine.search(request())).items
assert all(local_only for _, local_only in policy_runtime.calls)
asyncio.run(scenario())
def test_partition_storage_failure_rolls_back_all_partitions(policy_runtime, production_engine, monkeypatch):
from app.errors import ApiError
async def scenario():
await seed_policies()
before = [tuple(row) for row in rows('SELECT * FROM routed_block_vectors ORDER BY block_id')]
original = routed_vectors.store_remote
def fail_local(conn, ids, batch):
if batch.source != 'local':
original(conn, ids, batch)
monkeypatch.setattr(routed_vectors, 'store_remote', fail_local)
with pytest.raises(ApiError) as error:
await index_service.rebuild(IndexRebuildRequest())
assert error.value.code == 'SEMANTIC_INDEX_WRITE_FAILED'
assert [tuple(row) for row in rows('SELECT * FROM routed_block_vectors ORDER BY block_id')] == before
asyncio.run(scenario())
def test_missing_partition_does_not_silently_return_partial_hits(policy_runtime, production_engine):
from app.errors import ApiError
async def scenario():
await seed_policies()
conn = connect()
try:
conn.execute("DELETE FROM routed_block_vectors WHERE space_id='local-space'")
finally:
conn.close()
with pytest.raises(ApiError) as error:
await production_engine.search(request())
assert error.value.code == 'SEMANTIC_INDEX_UNAVAILABLE'
assert (await production_engine.search(request(SearchMode.hybrid))).items
asyncio.run(scenario())