Wire the guarded conversational RAG answer layer end-to-end

2026-08-05 14:33:13 +07:00
parent 834d9e51b0
commit ef08b4929e
127 changed files with 37921 additions and 169 deletions
@@ -0,0 +1,324 @@
+"""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