fix(retrieval): 按索引策略重建并融合跨空间检索
This commit is contained in:
@@ -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;
|
||||||
|
""",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
@@ -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", "向量索引写入失败,原索引已保留,请检查数据库和磁盘状态。")
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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())
|
||||||
|
|||||||
Reference in New Issue
Block a user