From 6661f4f8c693e115d7a91ab24bfcb53614bac9c1 Mon Sep 17 00:00:00 2001 From: KiriAky 107 Date: Fri, 4 Sep 2026 19:48:41 +0800 Subject: [PATCH] =?UTF-8?q?fix(retrieval):=20=E6=8C=89=E7=B4=A2=E5=BC=95?= =?UTF-8?q?=E7=AD=96=E7=95=A5=E9=87=8D=E5=BB=BA=E5=B9=B6=E8=9E=8D=E5=90=88?= =?UTF-8?q?=E8=B7=A8=E7=A9=BA=E9=97=B4=E6=A3=80=E7=B4=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/database/migrations.py | 4 ++ backend/app/retrieval/routed_vectors.py | 62 +++++++++++++++++++ backend/app/services/index_service.py | 11 ++-- backend/app/services/note_service.py | 2 + backend/tests/test_routed_retrieval.py | 82 +++++++++++++++++++++++++ 5 files changed, 156 insertions(+), 5 deletions(-) diff --git a/backend/app/database/migrations.py b/backend/app/database/migrations.py index c41f8e0..5fd1bd8 100644 --- a/backend/app/database/migrations.py +++ b/backend/app/database/migrations.py @@ -127,6 +127,10 @@ MIGRATIONS: list[str] = [ 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; + """, ] diff --git a/backend/app/retrieval/routed_vectors.py b/backend/app/retrieval/routed_vectors.py index 666526c..f259d3a 100644 --- a/backend/app/retrieval/routed_vectors.py +++ b/backend/app/retrieval/routed_vectors.py @@ -22,6 +22,7 @@ from app.database.db import connect, transaction from app.errors import ApiError from app.retrieval.vectorstore import VectorHit from app.retrieval.provenance import record_embedding +from app.retrieval.hybrid import rrf_fuse 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 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) if batch is None: return None + record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions}) try: conn = connect() @@ -221,3 +231,55 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa "Embedding 已可用，但当前模型的向量索引缺失、不完整或已失效。请在「设置 → 索引与模型」中重建全部索引。", {"model_id": batch.space_id, "dimensions": batch.dimensions, "source": batch.source}) from exc 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() diff --git a/backend/app/services/index_service.py b/backend/app/services/index_service.py index 7291cd2..4c72c92 100644 --- a/backend/app/services/index_service.py +++ b/backend/app/services/index_service.py @@ -85,7 +85,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob: )) try: prepared_notes = [] - semantic_space = None + semantic_spaces = {} for rel, folder, markdown, created, updated in docs: parsed = parse_note( markdown=markdown, file_path=rel, folder=folder, tags=None, @@ -97,9 +97,10 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob: if batch is None: raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 未生成向量，重建已停止，原索引已保留。") 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 模型发生切换，原索引已保留，请待模型服务稳定后重试。") - semantic_space = space + semantic_spaces[policy] = space prepared_notes.append((parsed, prepared)) # All network/model awaits precede the transaction. The concrete SQLite # methods below complete synchronously despite their async interfaces. @@ -114,12 +115,12 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob: await vector_store.clear(conn=conn) for parsed, prepared in prepared_notes: 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() missing = not exists or conn.execute( "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=? " - "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() if missing: raise ApiError(500, "SEMANTIC_INDEX_WRITE_FAILED", "向量索引写入失败，原索引已保留，请检查数据库和磁盘状态。") diff --git a/backend/app/services/note_service.py b/backend/app/services/note_service.py index 3c7da5c..a4f4eeb 100644 --- a/backend/app/services/note_service.py +++ b/backend/app/services/note_service.py @@ -118,6 +118,8 @@ async def index_note( blocks=parsed.blocks, ) 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} stale_ids = [bid for bid in old_ids if bid not in new_ids] if stale_ids: diff --git a/backend/tests/test_routed_retrieval.py b/backend/tests/test_routed_retrieval.py index 06e66c6..d830c7e 100644 --- a/backend/tests/test_routed_retrieval.py +++ b/backend/tests/test_routed_retrieval.py @@ -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): 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())