RAG/backend/app/adapters/bailian/_base.py

"""Shared HTTP and validation machinery for Alibaba Cloud Model Studio."""

from __future__ import annotations

import asyncio
import re
import secrets
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import UTC, datetime
from email.utils import parsedate_to_datetime
from time import perf_counter
from typing import Any, Self
from urllib.parse import urlsplit

import httpx

from app.ports.model_providers import (
    ModelProviderError,
    ProviderErrorKind,
    ProviderUsage,
    TokenCounter,
)

_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:/-]{0,127}$")
_BAILIAN_BEIJING_HOST = re.compile(
    r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.cn-beijing\.maas\.aliyuncs\.com$"
)


def conservative_utf8_token_count(text: str) -> int:
    """Return a dependency-free conservative upper bound for byte-BPE tokens.

    Production may inject the provider tokenizer through ``token_counter``.  A
    UTF-8 byte count is deliberately conservative and, unlike a word count,
    does not undercount Chinese text or byte-fallback tokens.
    """

    return len(text.encode("utf-8"))


def count_tokens(
    counter: TokenCounter,
    text: str,
    *,
    operation: str,
) -> int:
    try:
        count = counter(text)
    except Exception:
        raise sanitized_error(
            operation=operation,
            kind=ProviderErrorKind.INVALID_REQUEST,
            provider_code="token_count_failed",
        ) from None
    if isinstance(count, bool) or not isinstance(count, int) or count < 0:
        raise sanitized_error(
            operation=operation,
            kind=ProviderErrorKind.INVALID_REQUEST,
            provider_code="invalid_token_count",
        )
    return count


def sanitized_error(
    *,
    operation: str,
    kind: ProviderErrorKind,
    status_code: int | None = None,
    provider_code: str | None = None,
    request_id: str | None = None,
    retryable: bool = False,
) -> ModelProviderError:
    return ModelProviderError(
        operation=operation,
        kind=kind,
        status_code=status_code,
        provider_code=provider_code,
        request_id=request_id,
        retryable=retryable,
    )


def invalid_request(operation: str, code: str) -> ModelProviderError:
    return sanitized_error(
        operation=operation,
        kind=ProviderErrorKind.INVALID_REQUEST,
        provider_code=code,
    )


def invalid_response(operation: str, code: str) -> ModelProviderError:
    return sanitized_error(
        operation=operation,
        kind=ProviderErrorKind.INVALID_RESPONSE,
        provider_code=code,
    )


def parse_usage(value: Any) -> ProviderUsage:
    if not isinstance(value, Mapping):
        return ProviderUsage()
    return ProviderUsage(
        input_tokens=_first_nonnegative_int(
            value.get("prompt_tokens"),
            value.get("input_tokens"),
        ),
        output_tokens=_first_nonnegative_int(
            value.get("completion_tokens"),
            value.get("output_tokens"),
        ),
        total_tokens=_nonnegative_int(value.get("total_tokens")),
    )


def response_model(
    body: Mapping[str, Any],
    requested_model: str,
    *,
    sensitive_values: Sequence[str] = (),
) -> str:
    return safe_identifier(body.get("model"), sensitive_values=sensitive_values) or requested_model


def safe_identifier(
    value: Any,
    *,
    sensitive_values: Sequence[str] = (),
) -> str | None:
    if not isinstance(value, str) or not _SAFE_IDENTIFIER.fullmatch(value):
        return None
    if any(value in sensitive or sensitive in value for sensitive in sensitive_values if sensitive):
        return None
    return value


def extract_request_id(
    body: Mapping[str, Any],
    *,
    sensitive_values: Sequence[str] = (),
) -> str | None:
    return safe_identifier(
        body.get("id") or body.get("request_id"),
        sensitive_values=sensitive_values,
    )


def _nonnegative_int(value: Any) -> int | None:
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        return None
    return int(value)


def _first_nonnegative_int(*values: Any) -> int | None:
    for value in values:
        parsed = _nonnegative_int(value)
        if parsed is not None:
            return parsed
    return None


@dataclass(frozen=True, slots=True)
class JsonResponse:
    body: Mapping[str, Any]
    elapsed_ms: float


class BailianHttpAdapter:
    """Minimal async HTTP client with sanitized error translation."""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str,
        expected_path: str,
        http_client: httpx.AsyncClient | None,
        timeout_seconds: float,
        max_retries: int = 0,
        retry_base_seconds: float = 0.5,
    ) -> None:
        if not api_key or api_key != api_key.strip():
            raise invalid_request("configuration", "invalid_api_key_value")
        if timeout_seconds <= 0:
            raise invalid_request("configuration", "invalid_timeout")
        if isinstance(max_retries, bool) or not isinstance(max_retries, int) or max_retries < 0:
            raise invalid_request("configuration", "invalid_max_retries")
        if retry_base_seconds < 0:
            raise invalid_request("configuration", "invalid_retry_base")

        self._base_url = _validated_base_url(base_url, expected_path)
        self._api_key = api_key
        self._owns_client = http_client is None
        self._client = http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(timeout_seconds),
            follow_redirects=False,
        )
        self._max_retries = max_retries
        self._retry_base_seconds = retry_base_seconds

    async def aclose(self) -> None:
        if self._owns_client:
            await self._client.aclose()

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, *_: object) -> None:
        await self.aclose()

    def _url(self, path: str) -> str:
        return f"{self._base_url}/{path.lstrip('/')}"

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

    async def _post_json(
        self,
        *,
        operation: str,
        path: str,
        payload: Mapping[str, Any],
        sensitive_values: Sequence[str],
    ) -> JsonResponse:
        started = perf_counter()
        for attempt in range(self._max_retries + 1):
            try:
                response = await self._client.post(
                    self._url(path),
                    headers=self._headers(),
                    json=payload,
                )
            except httpx.TimeoutException:
                error = sanitized_error(
                    operation=operation,
                    kind=ProviderErrorKind.TIMEOUT,
                    provider_code="request_timeout",
                    retryable=True,
                )
                if await self._maybe_retry(error, attempt=attempt, response=None):
                    continue
                raise error from None
            except httpx.HTTPError:
                error = sanitized_error(
                    operation=operation,
                    kind=ProviderErrorKind.TRANSPORT,
                    provider_code="transport_error",
                    retryable=True,
                )
                if await self._maybe_retry(error, attempt=attempt, response=None):
                    continue
                raise error from None

            if response.status_code >= 400:
                try:
                    self._raise_http_error(
                        operation=operation,
                        response=response,
                        sensitive_values=(*sensitive_values, self._api_key),
                    )
                except ModelProviderError as error:
                    if await self._maybe_retry(error, attempt=attempt, response=response):
                        continue
                    raise

            try:
                body = response.json()
            except ValueError:
                raise invalid_response(operation, "invalid_json") from None
            if not isinstance(body, Mapping):
                raise invalid_response(operation, "invalid_json_object")
            return JsonResponse(body=body, elapsed_ms=(perf_counter() - started) * 1000)

        raise AssertionError("bounded provider retry loop exhausted unexpectedly")

    async def _maybe_retry(
        self,
        error: ModelProviderError,
        *,
        attempt: int,
        response: httpx.Response | None,
    ) -> bool:
        if not error.retryable or attempt >= self._max_retries:
            return False
        await asyncio.sleep(self._retry_delay(attempt=attempt, response=response))
        return True

    def _retry_delay(self, *, attempt: int, response: httpx.Response | None) -> float:
        if response is not None:
            retry_after = response.headers.get("Retry-After")
            if retry_after:
                parsed_delay = _parse_retry_after(retry_after)
                if parsed_delay is not None:
                    return min(max(parsed_delay, 0.0), 30.0)
        base_delay = min(self._retry_base_seconds * (2**attempt), 30.0)
        if base_delay == 0:
            return 0.0
        jitter = base_delay * (float(secrets.randbelow(251)) / 1000.0)
        return float(min(base_delay + jitter, 30.0))

    def _raise_http_error(
        self,
        *,
        operation: str,
        response: httpx.Response,
        sensitive_values: Sequence[str],
    ) -> None:
        body: Mapping[str, Any] = {}
        try:
            decoded = response.json()
            if isinstance(decoded, Mapping):
                body = decoded
        except (ValueError, httpx.ResponseNotRead):
            pass

        nested_error = body.get("error")
        error_body = nested_error if isinstance(nested_error, Mapping) else body
        code = safe_identifier(
            error_body.get("code"),
            sensitive_values=sensitive_values,
        )
        request_id = extract_request_id(body, sensitive_values=sensitive_values)
        kind, retryable = _kind_for_status(response.status_code)
        raise sanitized_error(
            operation=operation,
            kind=kind,
            status_code=response.status_code,
            provider_code=code,
            request_id=request_id,
            retryable=retryable,
        ) from None


def _validated_base_url(base_url: str, expected_path: str) -> str:
    if not base_url or base_url != base_url.strip():
        raise invalid_request("configuration", "invalid_base_url")
    parsed = urlsplit(base_url)
    normalized_path = parsed.path.rstrip("/")
    if (
        parsed.scheme != "https"
        or not parsed.hostname
        or not _BAILIAN_BEIJING_HOST.fullmatch(parsed.hostname)
        or parsed.username is not None
        or parsed.password is not None
        or parsed.query
        or parsed.fragment
        or normalized_path != expected_path
    ):
        raise invalid_request("configuration", "invalid_base_url")
    return base_url.rstrip("/")


def _kind_for_status(status_code: int) -> tuple[ProviderErrorKind, bool]:
    if status_code == 400:
        return ProviderErrorKind.INVALID_REQUEST, False
    if status_code == 401:
        return ProviderErrorKind.AUTHENTICATION, False
    if status_code == 403:
        return ProviderErrorKind.PERMISSION_DENIED, False
    if status_code == 404:
        return ProviderErrorKind.NOT_FOUND, False
    if status_code == 408:
        return ProviderErrorKind.TIMEOUT, True
    if status_code == 429:
        return ProviderErrorKind.RATE_LIMITED, True
    return ProviderErrorKind.UPSTREAM, status_code >= 500


def _parse_retry_after(value: str) -> float | None:
    try:
        return float(value)
    except ValueError:
        try:
            retry_at = parsedate_to_datetime(value)
        except (TypeError, ValueError, OverflowError):
            return None
        if retry_at.tzinfo is None:
            retry_at = retry_at.replace(tzinfo=UTC)
        return (retry_at - datetime.now(UTC)).total_seconds()