148 lines
4.6 KiB
Python
148 lines
4.6 KiB
Python
"""The Converse adapter must translate provider failure into the domain error,
|
|
isolate the JSON envelope the API cannot enforce, and never leak an SDK type.
|
|
|
|
No network: a stub client stands in for `bedrock-runtime`. The grounding check
|
|
that actually guards correctness lives in `rag.grounding` and is tested there;
|
|
here we prove only the adapter's envelope and failure contract.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from adapters.bedrock_converse import (
|
|
BedrockConverseAnswerGenerator,
|
|
BedrockCohereReranker,
|
|
RerankUnavailable,
|
|
_extract_json,
|
|
)
|
|
from rag.ports import AnswerGenerationUnavailable
|
|
|
|
|
|
def _converse_reply(text: str, stop_reason: str = "end_turn") -> dict:
|
|
return {
|
|
"stopReason": stop_reason,
|
|
"output": {"message": {"content": [{"text": text}]}},
|
|
"usage": {"inputTokens": 10, "outputTokens": 5},
|
|
}
|
|
|
|
|
|
class _Client:
|
|
def __init__(self, reply: dict) -> None:
|
|
self._reply = reply
|
|
self.calls: list[dict] = []
|
|
|
|
def converse(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
return self._reply
|
|
|
|
|
|
def _payload(answer: str = "Liều 500 mg [1]", sufficient: bool = True) -> str:
|
|
return json.dumps({"answer": answer, "evidence_sufficient": sufficient}, ensure_ascii=False)
|
|
|
|
|
|
def test_returns_the_models_json_on_success():
|
|
client = _Client(_converse_reply(_payload()))
|
|
gen = BedrockConverseAnswerGenerator(client=client, model_id="deepseek.v3.2")
|
|
|
|
raw = gen.generate("system", "user", {"type": "object"})
|
|
|
|
parsed = json.loads(raw)
|
|
assert parsed["answer"] == "Liều 500 mg [1]"
|
|
assert parsed["evidence_sufficient"] is True
|
|
# The prompt carries the schema and a JSON-only directive to the model.
|
|
sent = client.calls[0]["messages"][0]["content"][0]["text"]
|
|
assert "JSON" in sent
|
|
|
|
|
|
def test_strips_a_markdown_fence_the_api_cannot_forbid():
|
|
fenced = f"```json\n{_payload()}\n```"
|
|
gen = BedrockConverseAnswerGenerator(client=_Client(_converse_reply(fenced)))
|
|
|
|
parsed = json.loads(gen.generate("s", "u", {}))
|
|
|
|
assert parsed["answer"] == "Liều 500 mg [1]"
|
|
|
|
|
|
def test_extracts_json_when_the_model_adds_prose_around_it():
|
|
noisy = f"Đây là câu trả lời: {_payload()} Hết."
|
|
assert json.loads(_extract_json(noisy))["evidence_sufficient"] is True
|
|
|
|
|
|
def test_a_filtered_stop_reason_is_an_outage_not_an_answer():
|
|
client = _Client(_converse_reply("", stop_reason="content_filtered"))
|
|
gen = BedrockConverseAnswerGenerator(client=client)
|
|
|
|
with pytest.raises(AnswerGenerationUnavailable):
|
|
gen.generate("s", "u", {})
|
|
|
|
|
|
def test_empty_text_raises_rather_than_returning_blank():
|
|
gen = BedrockConverseAnswerGenerator(client=_Client(_converse_reply(" ")))
|
|
|
|
with pytest.raises(AnswerGenerationUnavailable):
|
|
gen.generate("s", "u", {})
|
|
|
|
|
|
def test_real_botocore_client_error_becomes_the_domain_error():
|
|
botocore_exceptions = pytest.importorskip("botocore.exceptions")
|
|
|
|
class _Refusing:
|
|
def converse(self, **kwargs):
|
|
raise botocore_exceptions.ClientError(
|
|
{"Error": {"Code": "AccessDeniedException", "Message": "denied"}},
|
|
"Converse",
|
|
)
|
|
|
|
gen = BedrockConverseAnswerGenerator(client=_Refusing())
|
|
|
|
with pytest.raises(AnswerGenerationUnavailable):
|
|
gen.generate("s", "u", {})
|
|
|
|
|
|
# --- reranker -----------------------------------------------------------------
|
|
|
|
|
|
class _RerankBody:
|
|
def __init__(self, payload: dict) -> None:
|
|
self._data = json.dumps(payload).encode("utf-8")
|
|
|
|
def read(self) -> bytes:
|
|
return self._data
|
|
|
|
|
|
class _RerankClient:
|
|
def __init__(self, payload: dict) -> None:
|
|
self._payload = payload
|
|
|
|
def invoke_model(self, **kwargs):
|
|
return {"body": _RerankBody(self._payload)}
|
|
|
|
|
|
def test_rerank_returns_indices_most_relevant_first():
|
|
client = _RerankClient({"results": [{"index": 2}, {"index": 0}, {"index": 1}]})
|
|
reranker = BedrockCohereReranker(client=client)
|
|
|
|
order = reranker.rerank("chống chỉ định", ["a", "b", "c"])
|
|
|
|
assert order == [2, 0, 1]
|
|
|
|
|
|
def test_rerank_raises_so_the_caller_keeps_original_order():
|
|
botocore_exceptions = pytest.importorskip("botocore.exceptions")
|
|
|
|
class _Refusing:
|
|
def invoke_model(self, **kwargs):
|
|
raise botocore_exceptions.ClientError(
|
|
{"Error": {"Code": "AccessDeniedException", "Message": "denied"}},
|
|
"InvokeModel",
|
|
)
|
|
|
|
with pytest.raises(RerankUnavailable):
|
|
BedrockCohereReranker(client=_Refusing()).rerank("q", ["a", "b"])
|
|
|
|
|
|
def test_rerank_of_nothing_is_empty():
|
|
assert BedrockCohereReranker(client=_RerankClient({})).rerank("q", []) == []
|