145 lines
5.2 KiB
Python
145 lines
5.2 KiB
Python
"""The embedding boundary: what a provider must do, and what it must record.
|
|
|
|
Two things drive this design.
|
|
|
|
**Asymmetric models make the input kind part of the contract.** Cohere embeds
|
|
a corpus record and a search query into deliberately different subspaces —
|
|
the same string sent as `search_document` and as `search_query` does not come
|
|
back as the same vector. Getting that backwards raises no error; recall just
|
|
quietly drops. So `input_kind` is a required argument of `embed()`, not an
|
|
optional keyword a caller can forget, and the value used is recorded on every
|
|
vector so a mismatch is detectable after the fact.
|
|
|
|
**Provenance applies to vectors too.** CLAUDE.md requires an extracted unit to
|
|
stay traceable to its source; a vector is no different. `model_id`,
|
|
`dimensions`, `input_kind` and the sha256 of the exact text embedded are what
|
|
let a collection be checked for the one mistake that is invisible from the
|
|
outside — vectors from two different models mixed into one Qdrant collection,
|
|
where every query still returns *something*.
|
|
|
|
`normalized` is deliberately three-valued. Titan is asked to normalize and
|
|
says so; a local encoder is told to; Cohere's Bedrock documentation does not
|
|
state whether its float vectors are unit-length, so the field stays `None`
|
|
rather than guessing. An unmeasured claim does not get written down as a fact.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Optional, Sequence
|
|
|
|
INPUT_DOCUMENT = "document"
|
|
INPUT_QUERY = "query"
|
|
|
|
INPUT_KINDS = (INPUT_DOCUMENT, INPUT_QUERY)
|
|
|
|
|
|
def text_digest(text: str) -> str:
|
|
"""sha256 of the exact string sent to the provider.
|
|
|
|
Shared by every adapter so a cached vector can be matched to its text by
|
|
the same rule that produced it.
|
|
"""
|
|
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EmbeddingVector:
|
|
values: List[float]
|
|
text_sha256: str
|
|
provider: str
|
|
model_id: str
|
|
dimensions: int
|
|
input_kind: str
|
|
# None means the provider does not document it — not "no".
|
|
normalized: Optional[bool] = None
|
|
# Only some providers report it (Titan does, Cohere's documented text
|
|
# response does not).
|
|
input_token_count: Optional[int] = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EmbeddingBatch:
|
|
"""Vectors plus what a benchmark needs to compare providers fairly."""
|
|
|
|
vectors: List[EmbeddingVector] = field(default_factory=list)
|
|
request_count: int = 0
|
|
latency_ms: float = 0.0
|
|
|
|
|
|
class EmbeddingProvider(ABC):
|
|
"""One embedding model, reachable without the caller knowing its SDK.
|
|
|
|
Subclasses implement `_embed_batch` for a single request; `embed` owns
|
|
input validation, splitting into provider-sized requests, and timing, so
|
|
that logic exists once rather than per adapter.
|
|
"""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def name(self) -> str:
|
|
"""Short registry key, e.g. `titan-v2`."""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def model_id(self) -> str:
|
|
"""Provider-side identifier, e.g. `amazon.titan-embed-text-v2:0`."""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def dimensions(self) -> int:
|
|
"""Vector length this instance is configured to produce."""
|
|
|
|
@property
|
|
def max_batch_size(self) -> int:
|
|
"""Texts accepted per request. Default is the safest possible value."""
|
|
return 1
|
|
|
|
@abstractmethod
|
|
def _embed_batch(
|
|
self, texts: Sequence[str], input_kind: str
|
|
) -> List[EmbeddingVector]:
|
|
"""Embed at most `max_batch_size` texts in one provider request."""
|
|
|
|
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
|
|
if input_kind not in INPUT_KINDS:
|
|
raise ValueError(
|
|
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
|
|
)
|
|
if any(not t.strip() for t in texts):
|
|
raise ValueError("refusing to embed an empty or whitespace-only text")
|
|
|
|
vectors: List[EmbeddingVector] = []
|
|
requests = 0
|
|
started = time.perf_counter()
|
|
for start in range(0, len(texts), self.max_batch_size):
|
|
window = texts[start : start + self.max_batch_size]
|
|
vectors.extend(self._embed_batch(window, input_kind))
|
|
requests += 1
|
|
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
|
|
|
if len(vectors) != len(texts):
|
|
raise ValueError(
|
|
f"{self.name} returned {len(vectors)} vectors for "
|
|
f"{len(texts)} texts"
|
|
)
|
|
return EmbeddingBatch(
|
|
vectors=vectors, request_count=requests, latency_ms=elapsed_ms
|
|
)
|
|
|
|
def embed_documents(self, texts: Sequence[str]) -> EmbeddingBatch:
|
|
return self.embed(texts, INPUT_DOCUMENT)
|
|
|
|
def embed_queries(self, texts: Sequence[str]) -> EmbeddingBatch:
|
|
return self.embed(texts, INPUT_QUERY)
|
|
|
|
def _check_dimensions(self, values: Sequence[float]) -> None:
|
|
"""A wrong-length vector is a corpus-wide defect; fail on the first."""
|
|
if len(values) != self.dimensions:
|
|
raise ValueError(
|
|
f"{self.model_id} returned {len(values)} dimensions, "
|
|
f"expected {self.dimensions}"
|
|
)
|