128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["requests", "pyyaml", "markdown-it-py", "sqlite-vec==0.1.6"]
|
|
# ///
|
|
"""Embedding client and the identity guard over the vector space.
|
|
|
|
Two invariants live here:
|
|
|
|
* Vectors are stored L2-normalized, so cosine distance equals the dot product.
|
|
* The query prefix must be bit-identical at index and at query time. That is the real
|
|
reason it is persisted in `meta` rather than only read from the config — mixing
|
|
vectors produced under two different contracts degrades results silently, which is
|
|
the most expensive kind of bug.
|
|
|
|
An unreachable Ollama is not an error here: the caller stores chunks with
|
|
`embedded_at IS NULL` and the query side says out loud that it is FTS-only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
|
|
import requests
|
|
from wiki_chunker import CHUNKER_VERSION
|
|
from wiki_config import EmbeddingConfig
|
|
from wiki_db import EMBEDDING_DIMS
|
|
|
|
SCHEMA_VERSION = "1"
|
|
SQLITE_VEC_VERSION = "0.1.6"
|
|
REQUEST_TIMEOUT_SECONDS = 60
|
|
|
|
# Only these keys make two vectors comparable; the rest of `meta` is informational.
|
|
GUARDED_META_KEYS = (
|
|
"embedding_model",
|
|
"embedding_dims",
|
|
"normalized",
|
|
"query_prefix",
|
|
"chunker_version",
|
|
)
|
|
|
|
|
|
class EmbeddingUnavailable(RuntimeError):
|
|
"""Ollama could not be reached or refused the request."""
|
|
|
|
|
|
class IndexIdentityMismatch(RuntimeError):
|
|
"""The index was built under a different embedding contract — a reindex is needed."""
|
|
|
|
|
|
def expected_meta(config: EmbeddingConfig) -> dict[str, str]:
|
|
return {
|
|
"embedding_model": config.model,
|
|
"embedding_dims": str(config.dims),
|
|
"normalized": "l2",
|
|
"query_prefix": config.query_prefix,
|
|
"chunker_version": CHUNKER_VERSION,
|
|
"schema_version": SCHEMA_VERSION,
|
|
"sqlite_vec_version": SQLITE_VEC_VERSION,
|
|
}
|
|
|
|
|
|
def meta_mismatches(stored: dict[str, str], config: EmbeddingConfig) -> list[str]:
|
|
"""Guarded meta keys that disagree with the config. Empty on a fresh (unwritten) index."""
|
|
if not stored:
|
|
return []
|
|
expected = expected_meta(config)
|
|
return [key for key in GUARDED_META_KEYS if stored.get(key) != expected[key]]
|
|
|
|
|
|
def require_matching_index(stored: dict[str, str], config: EmbeddingConfig) -> None:
|
|
"""Refuse to query an index built under a different contract."""
|
|
if config.dims != EMBEDDING_DIMS:
|
|
raise IndexIdentityMismatch(f"reindex needed: config dims {config.dims} != schema dims {EMBEDDING_DIMS}")
|
|
mismatches = meta_mismatches(stored, config)
|
|
if mismatches:
|
|
detail = ", ".join(
|
|
f"{key}: index={stored.get(key)!r} config={expected_meta(config)[key]!r}" for key in mismatches
|
|
)
|
|
raise IndexIdentityMismatch(f"reindex needed: {detail}")
|
|
|
|
|
|
def l2_normalize(vector: Sequence[float]) -> list[float]:
|
|
norm = math.sqrt(sum(value * value for value in vector))
|
|
if norm == 0.0:
|
|
return list(vector)
|
|
return [value / norm for value in vector]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OllamaEmbedder:
|
|
config: EmbeddingConfig
|
|
|
|
def embed_documents(self, texts: Sequence[str]) -> list[list[float]]:
|
|
"""Documents are embedded without the instruct prefix (the model is asymmetric)."""
|
|
return self._embed(list(texts))
|
|
|
|
def embed_query(self, text: str) -> list[float]:
|
|
return self._embed([self.config.query_prefix + text])[0]
|
|
|
|
def probe(self) -> None:
|
|
"""Raise EmbeddingUnavailable unless the endpoint answers an embed request."""
|
|
self._embed(["ping"])
|
|
|
|
def _embed(self, inputs: list[str]) -> list[list[float]]:
|
|
if not inputs:
|
|
return []
|
|
payload = {
|
|
"model": self.config.model,
|
|
"input": inputs,
|
|
"keep_alive": self.config.keep_alive,
|
|
}
|
|
try:
|
|
response = requests.post(f"{self.config.endpoint}/api/embed", json=payload, timeout=REQUEST_TIMEOUT_SECONDS)
|
|
response.raise_for_status()
|
|
embeddings = response.json()["embeddings"]
|
|
except (requests.RequestException, KeyError, ValueError) as exc:
|
|
raise EmbeddingUnavailable(f"{self.config.endpoint}: {exc}") from exc
|
|
|
|
if len(embeddings) != len(inputs):
|
|
raise EmbeddingUnavailable(f"asked for {len(inputs)} vectors, got {len(embeddings)}")
|
|
for vector in embeddings:
|
|
if len(vector) != self.config.dims:
|
|
raise EmbeddingUnavailable(f"model returned {len(vector)} dims, config says {self.config.dims}")
|
|
return [l2_normalize(vector) for vector in embeddings]
|