202 lines
8.0 KiB
Python
202 lines
8.0 KiB
Python
"""Optional API embeddings, isolated from the stable hash/sqlite-vec index.
|
|
|
|
The runtime's model_id is the authoritative space ID (including provider URL,
|
|
endpoint, model and dimensions); equal dimensions alone never imply compatibility.
|
|
This phase uses a lazy, rebuildable SQLite side table instead of a schema migration.
|
|
Search scans only current blocks in one database snapshot and requires complete
|
|
coverage. Cosine ranking costs O(blocks * dimensions) with an O(top_k) heap; this
|
|
small-vault implementation should become a per-space ANN index at larger scale.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import heapq
|
|
import json
|
|
import logging
|
|
import math
|
|
import sqlite3
|
|
from dataclasses import dataclass
|
|
from typing import Protocol
|
|
|
|
from app.database.db import connect, transaction
|
|
from app.retrieval.vectorstore import VectorHit
|
|
from app.retrieval.provenance import record_embedding
|
|
|
|
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]) -> EmbeddingResult: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteEmbeddings:
|
|
space_id: str
|
|
dimensions: int
|
|
vectors: list[list[float]]
|
|
|
|
|
|
def get_model_routing() -> EmbeddingRuntime | None:
|
|
"""Lazy integration hook; tests can inject a runtime without any network 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")
|
|
# Scaling first avoids overflow/underflow for finite but extreme API values.
|
|
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]) -> RemoteEmbeddings | None:
|
|
"""Return validated API vectors, or None to use the caller's local baseline.
|
|
|
|
Do not use the runtime's local result: the caller may have injected its own
|
|
embedding/store pair. Exception deliberately excludes cancellation.
|
|
"""
|
|
if not texts:
|
|
return None
|
|
try:
|
|
runtime = get_model_routing()
|
|
if runtime is None:
|
|
return None
|
|
result = await runtime.embed(texts)
|
|
if result.source != "api":
|
|
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],
|
|
)
|
|
except Exception as exc:
|
|
# Avoid logging provider exceptions containing credentials or note text.
|
|
record_embedding(fallback_reason="REMOTE_EMBEDDING_UNAVAILABLE")
|
|
logger.warning("Remote embedding unavailable (%s); using local index", type(exc).__name__)
|
|
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, block_id)
|
|
)
|
|
""")
|
|
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:
|
|
"""Best-effort side-index write inside the caller's metadata transaction.
|
|
|
|
A savepoint prevents partial remote batches and isolates storage failures from
|
|
note saving. Replacing/deleting blocks cascades all old spaces automatically.
|
|
"""
|
|
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, 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)
|
|
],
|
|
)
|
|
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) -> list[VectorHit] | None:
|
|
"""None means fallback, including any missing/invalid current-block vector.
|
|
|
|
Read coverage and vectors together so concurrent note updates cannot produce
|
|
an apparently complete subset. Never fill missing remote hits with local hits.
|
|
"""
|
|
batch = await embed_remote([query])
|
|
if batch is None:
|
|
return None
|
|
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")
|
|
return None
|
|
rows = conn.execute(
|
|
"""SELECT b.block_id, r.vector
|
|
FROM blocks AS b
|
|
LEFT JOIN routed_block_vectors AS r
|
|
ON r.block_id = b.block_id AND r.space_id = ? AND r.dimensions = ?
|
|
ORDER BY b.block_id""",
|
|
(batch.space_id, batch.dimensions),
|
|
)
|
|
|
|
def hits():
|
|
for row in rows:
|
|
if row["vector"] is None:
|
|
raise ValueError("remote space has incomplete block 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)))
|
|
|
|
result = heapq.nlargest(top_k, hits(), key=lambda hit: hit.score)
|
|
record_embedding(source="api", 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__)
|
|
return None
|