Fix live multi-turn: pass the resolved drug, stop did-you-mean garbage
This commit is contained in:
@@ -0,0 +1,199 @@
|
||||
"""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,
|
||||
retries={"max_attempts": 3, "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,
|
||||
retries={"max_attempts": 3, "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]
|
||||
Reference in New Issue
Block a user