325 lines
10 KiB
Python
325 lines
10 KiB
Python
"""Provider adapters, exercised with no AWS account and no network.
|
|
|
|
Every Bedrock call goes through a recording stub, so what is under test is the
|
|
part that can actually be wrong offline: the request body we send, and our
|
|
reading of the response bodies AWS documents. The one thing these tests cannot
|
|
establish is whether AWS accepts that body — that needs the live probe, and
|
|
the coordination handoff says so explicitly.
|
|
"""
|
|
import io
|
|
import json
|
|
import math
|
|
|
|
import pytest
|
|
|
|
from ingestion.embed import (
|
|
BGE_M3,
|
|
COHERE_V4,
|
|
TITAN_V2,
|
|
INPUT_DOCUMENT,
|
|
INPUT_QUERY,
|
|
Boto3BedrockInvoker,
|
|
EmbeddingVector,
|
|
build_provider,
|
|
provider_names,
|
|
text_digest,
|
|
)
|
|
from ingestion.embed import probe
|
|
from ingestion.embed.bedrock_cohere import CohereEmbedV4
|
|
from ingestion.embed.bedrock_titan import TitanTextEmbeddingsV2
|
|
from ingestion.embed.local_bge_m3 import BgeM3Local
|
|
|
|
|
|
class RecordingInvoker:
|
|
"""Stands in for Bedrock; remembers every request it was handed."""
|
|
|
|
def __init__(self, responses):
|
|
self._responses = list(responses)
|
|
self.calls = []
|
|
|
|
def invoke_json(self, model_id, payload, accept="application/json"):
|
|
self.calls.append(
|
|
{"model_id": model_id, "payload": payload, "accept": accept}
|
|
)
|
|
return self._responses.pop(0)
|
|
|
|
|
|
def _titan_response(dimensions=1024, token_count=12):
|
|
return {
|
|
"embedding": [0.01] * dimensions,
|
|
"inputTextTokenCount": token_count,
|
|
"embeddingsByType": {"float": [0.01] * dimensions},
|
|
}
|
|
|
|
|
|
def _cohere_by_type_response(rows, dimensions=1024):
|
|
return {
|
|
"id": "stub-id",
|
|
"response_type": "embeddings_by_type",
|
|
"embeddings": {"float": [[0.02] * dimensions for _ in range(rows)]},
|
|
"texts": ["stub"] * rows,
|
|
}
|
|
|
|
|
|
def _cohere_floats_response(rows, dimensions=1024):
|
|
return {
|
|
"id": "stub-id",
|
|
"response_type": "embeddings_floats",
|
|
"embeddings": [[0.02] * dimensions for _ in range(rows)],
|
|
}
|
|
|
|
|
|
def test_titan_request_body_matches_the_documented_v2_shape():
|
|
invoker = RecordingInvoker([_titan_response()])
|
|
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024, normalize=True)
|
|
|
|
provider.embed_documents(["paracetamol"])
|
|
|
|
payload = invoker.calls[0]["payload"]
|
|
assert invoker.calls[0]["model_id"] == "amazon.titan-embed-text-v2:0"
|
|
assert payload == {
|
|
"inputText": "paracetamol",
|
|
"dimensions": 1024,
|
|
"normalize": True,
|
|
}
|
|
|
|
|
|
def test_titan_records_provenance_and_reported_token_count():
|
|
invoker = RecordingInvoker([_titan_response(token_count=7)])
|
|
provider = TitanTextEmbeddingsV2(invoker)
|
|
|
|
vector = provider.embed_documents(["paracetamol"]).vectors[0]
|
|
|
|
assert vector.model_id == "amazon.titan-embed-text-v2:0"
|
|
assert vector.provider == TITAN_V2
|
|
assert vector.dimensions == 1024
|
|
assert vector.input_kind == INPUT_DOCUMENT
|
|
assert vector.normalized is True
|
|
assert vector.input_token_count == 7
|
|
assert vector.text_sha256 == text_digest("paracetamol")
|
|
|
|
|
|
def test_titan_sends_one_request_per_text():
|
|
invoker = RecordingInvoker([_titan_response(), _titan_response()])
|
|
provider = TitanTextEmbeddingsV2(invoker)
|
|
|
|
batch = provider.embed_documents(["a", "b"])
|
|
|
|
assert batch.request_count == 2
|
|
assert len(batch.vectors) == 2
|
|
|
|
|
|
def test_titan_rejects_a_dimension_the_model_does_not_offer():
|
|
with pytest.raises(ValueError, match="supports"):
|
|
TitanTextEmbeddingsV2(RecordingInvoker([]), dimensions=768)
|
|
|
|
|
|
def test_cohere_uses_search_document_for_corpus_and_search_query_for_queries():
|
|
invoker = RecordingInvoker(
|
|
[_cohere_by_type_response(1), _cohere_by_type_response(1)]
|
|
)
|
|
provider = CohereEmbedV4(invoker)
|
|
|
|
provider.embed_documents(["metformin"])
|
|
provider.embed_queries(["liều metformin"])
|
|
|
|
assert invoker.calls[0]["payload"]["input_type"] == "search_document"
|
|
assert invoker.calls[1]["payload"]["input_type"] == "search_query"
|
|
|
|
|
|
def test_cohere_request_body_pins_dimension_float_type_and_no_truncation():
|
|
invoker = RecordingInvoker([_cohere_by_type_response(2)])
|
|
provider = CohereEmbedV4(invoker, dimensions=1024)
|
|
|
|
provider.embed_documents(["a", "b"])
|
|
|
|
payload = invoker.calls[0]["payload"]
|
|
assert invoker.calls[0]["model_id"] == "cohere.embed-v4:0"
|
|
assert payload["texts"] == ["a", "b"]
|
|
assert payload["embedding_types"] == ["float"]
|
|
# Left unset the model would return 1536, which no 1024-wide collection
|
|
# can accept.
|
|
assert payload["output_dimension"] == 1024
|
|
# An over-length input must fail, not arrive silently shortened.
|
|
assert payload["truncate"] == "NONE"
|
|
assert invoker.calls[0]["accept"] == "*/*"
|
|
|
|
|
|
def test_cohere_reads_the_embeddings_by_type_response():
|
|
invoker = RecordingInvoker([_cohere_by_type_response(2)])
|
|
provider = CohereEmbedV4(invoker)
|
|
|
|
batch = provider.embed_documents(["a", "b"])
|
|
|
|
assert len(batch.vectors) == 2
|
|
assert all(len(v.values) == 1024 for v in batch.vectors)
|
|
assert batch.request_count == 1
|
|
|
|
|
|
def test_cohere_also_reads_the_plain_embeddings_floats_response():
|
|
invoker = RecordingInvoker([_cohere_floats_response(2)])
|
|
provider = CohereEmbedV4(invoker)
|
|
|
|
batch = provider.embed_documents(["a", "b"])
|
|
|
|
assert len(batch.vectors) == 2
|
|
assert all(len(v.values) == 1024 for v in batch.vectors)
|
|
|
|
|
|
def test_cohere_leaves_normalization_unknown_because_the_docs_do_not_say():
|
|
invoker = RecordingInvoker([_cohere_by_type_response(1)])
|
|
|
|
vector = CohereEmbedV4(invoker).embed_documents(["a"]).vectors[0]
|
|
|
|
assert vector.normalized is None
|
|
|
|
|
|
def test_cohere_splits_at_the_documented_96_text_ceiling():
|
|
invoker = RecordingInvoker(
|
|
[_cohere_by_type_response(96), _cohere_by_type_response(4)]
|
|
)
|
|
provider = CohereEmbedV4(invoker)
|
|
|
|
batch = provider.embed_documents([f"t{i}" for i in range(100)])
|
|
|
|
assert batch.request_count == 2
|
|
assert len(invoker.calls[0]["payload"]["texts"]) == 96
|
|
assert len(invoker.calls[1]["payload"]["texts"]) == 4
|
|
assert len(batch.vectors) == 100
|
|
|
|
|
|
def test_cohere_rejects_a_batch_size_above_the_documented_ceiling():
|
|
with pytest.raises(ValueError, match="batch_size"):
|
|
CohereEmbedV4(RecordingInvoker([]), batch_size=97)
|
|
|
|
|
|
def test_a_wrong_width_vector_fails_instead_of_entering_the_corpus():
|
|
invoker = RecordingInvoker([_titan_response(dimensions=512)])
|
|
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024)
|
|
|
|
with pytest.raises(ValueError, match="512 dimensions"):
|
|
provider.embed_documents(["a"])
|
|
|
|
|
|
def test_a_response_missing_its_vectors_fails_loudly():
|
|
invoker = RecordingInvoker([{"id": "stub", "response_type": "x"}])
|
|
|
|
with pytest.raises(ValueError, match="no 'embeddings' field"):
|
|
CohereEmbedV4(invoker).embed_documents(["a"])
|
|
|
|
|
|
def test_a_count_mismatch_between_texts_and_vectors_fails():
|
|
invoker = RecordingInvoker([_cohere_by_type_response(1)])
|
|
|
|
with pytest.raises(ValueError, match="1 vectors for 2 texts"):
|
|
CohereEmbedV4(invoker).embed_documents(["a", "b"])
|
|
|
|
|
|
def test_an_unknown_input_kind_is_refused_before_any_request_is_made():
|
|
invoker = RecordingInvoker([])
|
|
|
|
with pytest.raises(ValueError, match="input_kind"):
|
|
CohereEmbedV4(invoker).embed(["a"], "search_document")
|
|
assert invoker.calls == []
|
|
|
|
|
|
def test_empty_text_is_refused_before_any_request_is_made():
|
|
invoker = RecordingInvoker([])
|
|
|
|
with pytest.raises(ValueError, match="empty"):
|
|
CohereEmbedV4(invoker).embed_documents(["a", " "])
|
|
assert invoker.calls == []
|
|
|
|
|
|
def test_bge_m3_runs_through_an_injected_encoder_with_no_weights_loaded():
|
|
seen = []
|
|
|
|
def encoder(texts):
|
|
seen.append(list(texts))
|
|
unit = 1.0 / math.sqrt(1024)
|
|
return [[unit] * 1024 for _ in texts]
|
|
|
|
provider = BgeM3Local(encoder=encoder, batch_size=2)
|
|
batch = provider.embed_queries(["a", "b", "c"])
|
|
|
|
assert seen == [["a", "b"], ["c"]]
|
|
assert batch.request_count == 2
|
|
assert len(batch.vectors) == 3
|
|
assert batch.vectors[0].input_kind == INPUT_QUERY
|
|
assert batch.vectors[0].model_id == "BAAI/bge-m3"
|
|
# Injected encoder: we did not set normalize_embeddings, so we do not claim it.
|
|
assert batch.vectors[0].normalized is None
|
|
|
|
|
|
def test_registry_builds_every_provider_without_touching_an_sdk():
|
|
assert set(provider_names()) == {TITAN_V2, COHERE_V4, BGE_M3}
|
|
|
|
titan = build_provider(TITAN_V2, invoker=RecordingInvoker([]))
|
|
cohere = build_provider(COHERE_V4, invoker=RecordingInvoker([]))
|
|
local = build_provider(BGE_M3)
|
|
|
|
assert (titan.dimensions, cohere.dimensions, local.dimensions) == (
|
|
1024,
|
|
1024,
|
|
1024,
|
|
)
|
|
assert titan.max_batch_size == 1
|
|
assert cohere.max_batch_size == 96
|
|
|
|
|
|
def test_registry_rejects_an_unknown_provider_name():
|
|
with pytest.raises(ValueError, match="unknown embedding provider"):
|
|
build_provider("text-embedding-3-small")
|
|
|
|
|
|
class FakeBotoClient:
|
|
"""The shape boto3's bedrock-runtime client returns: a streaming body."""
|
|
|
|
def __init__(self, response_body):
|
|
self._response_body = response_body
|
|
self.kwargs = None
|
|
|
|
def invoke_model(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
return {"body": io.BytesIO(json.dumps(self._response_body).encode())}
|
|
|
|
|
|
def test_boto3_invoker_serialises_the_request_and_reads_the_streamed_body():
|
|
client = FakeBotoClient({"embedding": [0.5]})
|
|
invoker = Boto3BedrockInvoker(region="us-east-1", client=client)
|
|
|
|
body = invoker.invoke_json("some.model", {"inputText": "à"}, accept="*/*")
|
|
|
|
assert body == {"embedding": [0.5]}
|
|
assert client.kwargs["modelId"] == "some.model"
|
|
assert client.kwargs["contentType"] == "application/json"
|
|
assert client.kwargs["accept"] == "*/*"
|
|
# Vietnamese must survive the round trip as characters, not \\u escapes
|
|
# the model would then embed literally.
|
|
assert json.loads(client.kwargs["body"]) == {"inputText": "à"}
|
|
|
|
|
|
def test_probe_measures_the_l2_norm_rather_than_trusting_the_docs():
|
|
unit = 1.0 / math.sqrt(4)
|
|
assert probe._l2_norm([unit] * 4) == pytest.approx(1.0)
|
|
assert probe._l2_norm([3.0, 4.0]) == pytest.approx(5.0)
|
|
|
|
|
|
def test_probe_reports_a_vector_without_raising(capsys):
|
|
vector = EmbeddingVector(
|
|
values=[0.5, 0.5, 0.5, 0.5],
|
|
text_sha256=text_digest("x"),
|
|
provider=TITAN_V2,
|
|
model_id="amazon.titan-embed-text-v2:0",
|
|
dimensions=4,
|
|
input_kind=INPUT_DOCUMENT,
|
|
normalized=True,
|
|
input_token_count=3,
|
|
)
|
|
|
|
probe._report(vector, latency_ms=12.5, requests=1)
|
|
|
|
out = capsys.readouterr().out
|
|
assert "amazon.titan-embed-text-v2:0" in out
|
|
assert "measured L2 norm: 1.000000" in out
|