admin
d703ab64e3
CI / docs-check (push) Canceled after 0s
CI / backend-test (push) Canceled after 0s
CI / service-test (push) Canceled after 0s
CI / frontend-test (push) Canceled after 0s
CI / rust-core (push) Canceled after 0s
CI / docs-check (pull_request) Canceled after 0s
CI / backend-test (pull_request) Canceled after 0s
CI / service-test (pull_request) Canceled after 0s
CI / frontend-test (pull_request) Canceled after 0s
CI / rust-core (pull_request) Canceled after 0s
297 lines
13 KiB
Python
297 lines
13 KiB
Python
"""可选的 API 嵌入，与稳定的 hash/sqlite-vec 索引相互隔离。
|
|
|
|
运行时的 model_id 是权威空间标识，涵盖提供商 URL、端点、模型与维度；维度相同并不表示兼容。
|
|
持久化向量用于按需构建各空间和维度的 sqlite-vec 索引。原生精确 KNN 避免每次搜索都由 Python
|
|
解码 JSON 并计算点积。覆盖率检查与排序使用同一事务。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import math
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
from typing import Protocol
|
|
|
|
from app.database.db import connect_knowledge as connect, transaction
|
|
from app.errors import ApiError
|
|
from app.operation_logs import log_event
|
|
from app.retrieval.vectorstore import VectorHit
|
|
from app.retrieval.provenance import record_embedding
|
|
from app.retrieval.hybrid import rrf_fuse
|
|
from app.retrieval import space_index
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EmbeddingResult(Protocol):
|
|
vectors: list[list[float]]
|
|
source: str
|
|
model_id: str
|
|
dimensions: int
|
|
fallback_reason: str | None
|
|
|
|
|
|
class EmbeddingRuntime(Protocol):
|
|
async def embed(self, texts: list[str], *, local_only=False) -> EmbeddingResult: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteEmbeddings:
|
|
space_id: str
|
|
dimensions: int
|
|
vectors: list[list[float]]
|
|
source: str = "api"
|
|
|
|
|
|
def get_model_routing() -> EmbeddingRuntime | None:
|
|
"""惰性集成钩子；测试可以注入运行时而无需任何网络 I/O。"""
|
|
from app.container import container
|
|
|
|
return getattr(container, "model_routing", None)
|
|
|
|
|
|
def _unit_vector(vector: list[float], dimensions: int) -> list[float]:
|
|
if len(vector) != dimensions:
|
|
raise ValueError("embedding dimension mismatch")
|
|
if any(isinstance(value, bool) or not isinstance(value, (int, float)) for value in vector):
|
|
raise ValueError("embedding must be numeric")
|
|
if not all(math.isfinite(value) for value in vector):
|
|
raise ValueError("embedding must be finite")
|
|
scale = max(abs(value) for value in vector)
|
|
if scale == 0:
|
|
raise ValueError("embedding must be nonzero")
|
|
# 缩放首先避免有限但极端的 API 值的上溢/下溢。
|
|
scaled = [value / scale for value in vector]
|
|
norm = math.sqrt(math.fsum(value * value for value in scaled))
|
|
return [value / norm for value in scaled]
|
|
|
|
|
|
async def embed_remote(texts: list[str], *, accept_local=False, strict=False, local_only=False) -> RemoteEmbeddings | None:
|
|
"""返回经过验证的 API 向量，或 None 以使用调用者的本地基线。不要使用运行时的本地结果：调用者可能已经注入了自己的嵌入/存储对。异常特意排除取消。"""
|
|
if not texts:
|
|
return None
|
|
try:
|
|
runtime = get_model_routing()
|
|
if runtime is None:
|
|
if strict:
|
|
raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 服务未就绪，请检查模型路由和本地运行环境。")
|
|
return None
|
|
result = await runtime.embed(texts, local_only=True) if local_only else await runtime.embed(texts)
|
|
if result.source != "api" and not accept_local:
|
|
record_embedding(fallback_reason=result.fallback_reason)
|
|
return None
|
|
if not isinstance(result.model_id, str) or not result.model_id or result.model_id == "hash-v1":
|
|
raise ValueError("API embedding needs a distinct space ID")
|
|
if type(result.dimensions) is not int or result.dimensions <= 0:
|
|
raise ValueError("invalid embedding dimensions")
|
|
if len(result.vectors) != len(texts):
|
|
raise ValueError("embedding count mismatch")
|
|
return RemoteEmbeddings(
|
|
space_id=result.model_id,
|
|
dimensions=result.dimensions,
|
|
vectors=[_unit_vector(vector, result.dimensions) for vector in result.vectors],
|
|
source=result.source,
|
|
)
|
|
except Exception as exc:
|
|
log_event('vectors', 'embedding.failed', level='ERROR' if strict else 'WARNING', error=exc,
|
|
count=len(texts), fallback='none' if strict else 'local_index')
|
|
# 避免记录包含凭据或笔记文本的提供程序异常。
|
|
record_embedding(fallback_reason="REMOTE_EMBEDDING_UNAVAILABLE")
|
|
logger.warning("Remote embedding unavailable (%s); using local index", type(exc).__name__)
|
|
if strict:
|
|
if isinstance(exc, ApiError):
|
|
raise
|
|
raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 调用失败或返回无效，请检查模型路由、API 和本地模型运行状态。") from exc
|
|
return None
|
|
|
|
|
|
def _ensure_table(conn: sqlite3.Connection) -> None:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS routed_block_vectors (
|
|
space_id TEXT NOT NULL,
|
|
block_id TEXT NOT NULL REFERENCES blocks(block_id) ON DELETE CASCADE,
|
|
dimensions INTEGER NOT NULL CHECK (dimensions > 0),
|
|
vector TEXT NOT NULL,
|
|
PRIMARY KEY (space_id, dimensions, block_id)
|
|
)
|
|
""")
|
|
primary = [row[1] for row in sorted(conn.execute('PRAGMA table_info(routed_block_vectors)'), key=lambda row: row[5]) if row[5]]
|
|
if primary == ['space_id', 'block_id']:
|
|
conn.execute('CREATE TABLE routed_block_vectors_upgrade (space_id TEXT NOT NULL, block_id TEXT NOT NULL REFERENCES blocks(block_id) ON DELETE CASCADE, dimensions INTEGER NOT NULL CHECK(dimensions>0), vector TEXT NOT NULL, PRIMARY KEY(space_id,dimensions,block_id))')
|
|
conn.execute('INSERT INTO routed_block_vectors_upgrade SELECT * FROM routed_block_vectors')
|
|
conn.execute('DROP TABLE routed_block_vectors')
|
|
conn.execute('ALTER TABLE routed_block_vectors_upgrade RENAME TO routed_block_vectors')
|
|
conn.execute("""
|
|
CREATE INDEX IF NOT EXISTS routed_block_vectors_block_id
|
|
ON routed_block_vectors(block_id)
|
|
""")
|
|
|
|
|
|
def store_remote(
|
|
conn: sqlite3.Connection, block_ids: list[str], batch: RemoteEmbeddings | None,
|
|
) -> None:
|
|
"""在调用方的元数据事务内尽力写入辅助索引。
|
|
|
|
savepoint 可阻止只写入部分远程批次，并将存储故障与笔记保存隔离；替换或删除内容块时，
|
|
所有旧空间都会自动级联清理。
|
|
"""
|
|
if batch is None:
|
|
return
|
|
try:
|
|
conn.execute("SAVEPOINT routed_vectors_write")
|
|
try:
|
|
if len(block_ids) != len(batch.vectors):
|
|
raise ValueError("block/vector count mismatch")
|
|
_ensure_table(conn)
|
|
conn.executemany(
|
|
"""INSERT INTO routed_block_vectors (space_id, block_id, dimensions, vector)
|
|
VALUES (?, ?, ?, ?)
|
|
ON CONFLICT (space_id, dimensions, block_id) DO UPDATE SET
|
|
dimensions = excluded.dimensions, vector = excluded.vector""",
|
|
[
|
|
(batch.space_id, block_id, batch.dimensions, json.dumps(vector, allow_nan=False))
|
|
for block_id, vector in zip(block_ids, batch.vectors)
|
|
],
|
|
)
|
|
space_index.upsert(conn, block_ids, batch)
|
|
except BaseException:
|
|
conn.execute("ROLLBACK TO routed_vectors_write")
|
|
raise
|
|
finally:
|
|
conn.execute("RELEASE routed_vectors_write")
|
|
except Exception as exc:
|
|
logger.warning("Remote vector storage unavailable (%s); local index retained", type(exc).__name__)
|
|
|
|
|
|
async def search_remote(query: str, *, top_k: int, accept_local=False, strict=False) -> list[VectorHit] | None:
|
|
"""None 表示回退，包括任何丢失/无效的当前块向量。将覆盖率和向量一起读取，以便并发笔记更新无法生成明显完整的子集。切勿用本地命中来填补缺失的远程命中。"""
|
|
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
|
|
|
|
if not await _prepare_for_search([batch], strict):
|
|
return None
|
|
return await asyncio.to_thread(_search_space, batch, top_k, strict)
|
|
|
|
|
|
async def _prepare_indexes(batches):
|
|
from app.services.coordination import vault_mutation_lock
|
|
def prepare(check_only=False):
|
|
conn = connect()
|
|
try:
|
|
if check_only:
|
|
return space_index.is_ready(conn, batches)
|
|
space_index.prepare(conn, batches)
|
|
finally:
|
|
conn.close()
|
|
if await asyncio.to_thread(prepare, True):
|
|
return
|
|
# 与保存共享协作门：当迁移在另一个线程中拥有 SQLite 写锁时，永远不会阻塞 SQLite 写锁上的事件循环。
|
|
async with vault_mutation_lock():
|
|
work = asyncio.create_task(asyncio.to_thread(prepare))
|
|
cancelled = False
|
|
while not work.done():
|
|
try:
|
|
await asyncio.shield(work)
|
|
except asyncio.CancelledError:
|
|
cancelled = True
|
|
work.result()
|
|
if cancelled:
|
|
raise asyncio.CancelledError
|
|
|
|
|
|
async def _prepare_for_search(batches, strict):
|
|
try:
|
|
await _prepare_indexes(batches)
|
|
return True
|
|
except Exception as exc:
|
|
record_embedding(fallback_reason='REMOTE_INDEX_UNAVAILABLE')
|
|
if strict:
|
|
raise ApiError(409, 'SEMANTIC_INDEX_UNAVAILABLE', '向量索引准备失败，请检查索引状态。') from exc
|
|
return False
|
|
|
|
|
|
def _search_space(batch, top_k, strict):
|
|
record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
|
|
try:
|
|
conn = connect()
|
|
try:
|
|
with transaction(conn):
|
|
exists = conn.execute(
|
|
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'routed_block_vectors'"
|
|
).fetchone()
|
|
if exists is None:
|
|
record_embedding(fallback_reason="REMOTE_INDEX_MISSING")
|
|
if not conn.execute("SELECT 1 FROM blocks LIMIT 1").fetchone():
|
|
return []
|
|
if strict:
|
|
raise ValueError("semantic index missing")
|
|
return None
|
|
result = space_index.search(conn, batch, top_k)
|
|
record_embedding(source=batch.source, model_id=batch.space_id,
|
|
dimensions=batch.dimensions, fallback_reason=None)
|
|
return result
|
|
finally:
|
|
conn.close()
|
|
except Exception as exc:
|
|
record_embedding(fallback_reason="REMOTE_INDEX_UNAVAILABLE")
|
|
logger.debug("Remote vector search unavailable (%s); using local index", type(exc).__name__)
|
|
if strict:
|
|
raise ApiError(409, "SEMANTIC_INDEX_UNAVAILABLE",
|
|
"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):
|
|
"""按策略嵌入；独立对每个空间进行排名并融合排名，而不是向量。"""
|
|
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
|
|
if not await _prepare_for_search(list(batches.values()), strict):
|
|
return None
|
|
return await asyncio.to_thread(_search_partitions, batches, policies, top_k, strict)
|
|
|
|
|
|
def _search_partitions(batches, policies, top_k, strict):
|
|
conn = connect()
|
|
try:
|
|
with transaction(conn):
|
|
# 在打开单个读取快照之前，查询向量已准备就绪。
|
|
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():
|
|
ranked.append(space_index.search(conn, batch, top_k, policy))
|
|
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()
|