216 lines
9.1 KiB
Python
216 lines
9.1 KiB
Python
"""A Bedrock Converse-API answer generator (DeepSeek / Qwen / Nova / GLM …).
|
|
|
|
The counterpart to `bedrock_claude.py`, for every non-Anthropic model on
|
|
Bedrock. Those models are reached through the unified `converse` operation on
|
|
`bedrock-runtime` — the same client `embedding.py` already uses — so this adapter
|
|
adds no new SDK: `boto3` is imported lazily, and `rag/` still imports nothing.
|
|
|
|
One structural difference from the Anthropic path drives the shape of this file:
|
|
Converse has **no** server-side response schema (no `output_config.format`), so
|
|
the JSON envelope `answer.py` parses cannot be enforced by the API. It is asked
|
|
for in the prompt and then isolated here (`_extract_json`) before returning. If
|
|
the model still emits something unparseable, `answer.py` falls back to the
|
|
verbatim source text — losing the rewrite, never the answer.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from rag.ports import AnswerGenerationUnavailable, RerankUnavailable
|
|
|
|
BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
|
|
DEEPSEEK_V3_2 = "deepseek.v3.2"
|
|
|
|
# Sized for a rewrite of the retrieved evidence, not for open-ended generation:
|
|
# the section route can hand over a long section, and the answer restates it.
|
|
MAX_OUTPUT_TOKENS = 4096
|
|
|
|
# stopReasons that mean "a successful HTTP response carrying no usable answer".
|
|
# Treated as an outage so the caller degrades to the extractive text instead of
|
|
# reading content that was filtered away.
|
|
_EMPTY_STOP_REASONS = frozenset({"content_filtered", "guardrail_intervened"})
|
|
|
|
|
|
def _provider_error_types() -> tuple[type[BaseException], ...]:
|
|
"""botocore's error classes, or none when botocore is absent."""
|
|
try:
|
|
from botocore.exceptions import BotoCoreError, ClientError
|
|
except ImportError:
|
|
return ()
|
|
return (BotoCoreError, ClientError)
|
|
|
|
|
|
def _extract_json(text: str) -> str:
|
|
"""Isolate the JSON object from a Converse text block.
|
|
|
|
Converse cannot pin the output shape, so a model may wrap the object in a
|
|
```json fence or add a sentence around it. This returns the outermost
|
|
`{...}` span so `answer.py`'s `json.loads` sees the same clean envelope the
|
|
Anthropic adapter's schema-constrained path produces. If no object is found
|
|
the original text is returned, and the caller's parse fails closed.
|
|
"""
|
|
start = text.find("{")
|
|
end = text.rfind("}")
|
|
if start != -1 and end != -1 and end > start:
|
|
return text[start : end + 1]
|
|
return text
|
|
|
|
|
|
class BedrockConverseAnswerGenerator:
|
|
"""Rewrites retrieved evidence into prose via the Bedrock Converse API.
|
|
|
|
Model-agnostic: the model id is injected, so switching from DeepSeek to Qwen
|
|
or GLM is one config value (and one IAM resource ARN), no code change. What
|
|
the model returns is not trusted — `rag.grounding.verify` runs on every
|
|
answer this produces, so a fabricated figure yields a discarded generation,
|
|
not a wrong answer.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
region: str = "us-east-1",
|
|
client: Any | None = None,
|
|
model_id: str = DEEPSEEK_V3_2,
|
|
max_tokens: int = MAX_OUTPUT_TOKENS,
|
|
) -> None:
|
|
self._region = region
|
|
self._client = client
|
|
self._model_id = model_id
|
|
self._max_tokens = max_tokens
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
return self._model_id
|
|
|
|
def _runtime(self) -> Any:
|
|
if self._client is None:
|
|
import boto3
|
|
from botocore.config import Config
|
|
|
|
self._client = boto3.client(
|
|
BEDROCK_RUNTIME_SERVICE,
|
|
region_name=self._region,
|
|
config=Config(
|
|
connect_timeout=10,
|
|
read_timeout=60,
|
|
# "adaptive" was tried 2026-08-07 and reverted same day: its
|
|
# client-side rate limiter remembers "throttled" across
|
|
# requests and paces even unrelated, otherwise-healthy calls
|
|
# down after a burst — turned a single answerable turn's
|
|
# baseline ~9s into 1-5 MINUTES following this session's own
|
|
# heavy adversarial test traffic. "standard" retries each
|
|
# call independently, no shared state to get stuck in a bad
|
|
# regime. max_attempts alone (3->4) is kept.
|
|
retries={"max_attempts": 4, "mode": "standard"},
|
|
),
|
|
)
|
|
return self._client
|
|
|
|
def generate(self, system: str, user: str, schema: dict) -> str:
|
|
# The schema cannot be enforced by Converse, so it is stated in the
|
|
# message. Temperature 0: a formulary restatement is not a place for
|
|
# sampling variety.
|
|
directive = (
|
|
"Trả về DUY NHẤT một đối tượng JSON đúng schema sau, không kèm văn "
|
|
"bản nào khác, không dùng khối markdown ```:\n"
|
|
f"{json.dumps(schema, ensure_ascii=False)}"
|
|
)
|
|
try:
|
|
response = self._runtime().converse(
|
|
modelId=self._model_id,
|
|
system=[{"text": system}],
|
|
messages=[{"role": "user", "content": [{"text": f"{user}\n\n{directive}"}]}],
|
|
inferenceConfig={"maxTokens": self._max_tokens, "temperature": 0},
|
|
)
|
|
except _provider_error_types() as error:
|
|
raise AnswerGenerationUnavailable(
|
|
f"{self._model_id} could not be invoked: {type(error).__name__}"
|
|
) from error
|
|
|
|
if response.get("stopReason") in _EMPTY_STOP_REASONS:
|
|
raise AnswerGenerationUnavailable(
|
|
f"{self._model_id} produced no usable content "
|
|
f"(stopReason={response.get('stopReason')})"
|
|
)
|
|
|
|
blocks = response.get("output", {}).get("message", {}).get("content", [])
|
|
text = "".join(block.get("text", "") for block in blocks if isinstance(block, dict))
|
|
if not text.strip():
|
|
raise AnswerGenerationUnavailable(f"{self._model_id} returned no text content")
|
|
return _extract_json(text)
|
|
|
|
|
|
class BedrockCohereReranker:
|
|
"""Reorders candidate chunks by relevance with Cohere Rerank on Bedrock.
|
|
|
|
A cross-encoder rerank is the standard fix for the weak spot of pure vector
|
|
similarity: hit@1 0.544 was measured letting the embedding alone pick the
|
|
section, because the large pharmacology section sits close to every question.
|
|
Rerank scores each (query, chunk) pair jointly, so it recovers precision the
|
|
bi-encoder cannot. It is an optional improvement on the similarity fallback,
|
|
never on the deterministic section route — losing it reorders nothing, it
|
|
does not lose an answer.
|
|
"""
|
|
|
|
COHERE_RERANK_V3_5 = "cohere.rerank-v3-5:0"
|
|
|
|
def __init__(
|
|
self,
|
|
region: str = "us-east-1",
|
|
client: Any | None = None,
|
|
model_id: str = COHERE_RERANK_V3_5,
|
|
) -> None:
|
|
self._region = region
|
|
self._client = client
|
|
self._model_id = model_id
|
|
|
|
def _runtime(self) -> Any:
|
|
if self._client is None:
|
|
import boto3
|
|
from botocore.config import Config
|
|
|
|
self._client = boto3.client(
|
|
BEDROCK_RUNTIME_SERVICE,
|
|
region_name=self._region,
|
|
config=Config(
|
|
connect_timeout=10,
|
|
read_timeout=30,
|
|
# "adaptive" was tried 2026-08-07 and reverted same day: its
|
|
# client-side rate limiter remembers "throttled" across
|
|
# requests and paces even unrelated, otherwise-healthy calls
|
|
# down after a burst — turned a single answerable turn's
|
|
# baseline ~9s into 1-5 MINUTES following this session's own
|
|
# heavy adversarial test traffic. "standard" retries each
|
|
# call independently, no shared state to get stuck in a bad
|
|
# regime. max_attempts alone (3->4) is kept.
|
|
retries={"max_attempts": 4, "mode": "standard"},
|
|
),
|
|
)
|
|
return self._client
|
|
|
|
def rerank(self, query: str, documents: list[str], top_n: int | None = None) -> list[int]:
|
|
"""Return document indices, most relevant first. Never drops silently:
|
|
on any provider error it raises, and the caller keeps the input order."""
|
|
if not documents:
|
|
return []
|
|
n = top_n or len(documents)
|
|
try:
|
|
response = self._runtime().invoke_model(
|
|
modelId=self._model_id,
|
|
body=json.dumps(
|
|
{"query": query, "documents": documents, "top_n": n, "api_version": 2}
|
|
),
|
|
accept="*/*",
|
|
contentType="application/json",
|
|
)
|
|
except _provider_error_types() as error:
|
|
raise RerankUnavailable(
|
|
f"{self._model_id} could not be invoked: {type(error).__name__}"
|
|
) from error
|
|
body = json.loads(response["body"].read())
|
|
results = body.get("results")
|
|
if not results:
|
|
raise RerankUnavailable(f"{self._model_id} returned no results")
|
|
return [item["index"] for item in results]
|