duocthu/apps/ai-service/adapters/bedrock_converse.py

"""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` fails closed with a
diagnosable abstention instead of presenting a raw dump as a generated 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
        # Converse's own response usage, captured for the observability layer
        # (rag/instrumentation.py) to attach to its span. Not part of the
        # `generate()` return contract -- every existing caller is unaffected.
        self.last_usage: dict[str, int] | None = None

    @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=5,
                    # The request budget is checked between calls and cannot
                    # interrupt boto3 while a call is in flight. A 60s read
                    # timeout with four SDK attempts allowed one turn to run
                    # for minutes. Allow one bounded retry: production smoke
                    # tests showed an isolated Qwen read timeout immediately
                    # followed by a healthy 3s response for the same request.
                    read_timeout=20,
                    retries={"total_max_attempts": 2, "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

        self.last_usage = dict(response.get("usage") or {}) or None

        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]