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

"""Alibaba Cloud Model Studio OpenAI-compatible embedding adapter."""

from __future__ import annotations

import math
from collections.abc import Mapping, Sequence
from typing import Any

import httpx

from app.adapters.bailian._base import (
    BailianHttpAdapter,
    conservative_utf8_token_count,
    count_tokens,
    extract_request_id,
    invalid_request,
    invalid_response,
    parse_usage,
    response_model,
)
from app.ports.model_providers import EmbeddingResult, TokenCounter

EMBEDDING_MAX_BATCH_SIZE = 10
EMBEDDING_MAX_TOKENS_PER_TEXT = 8_192
EMBEDDING_MAX_TOKENS_PER_REQUEST = 33_000
EMBEDDING_DIMENSIONS = 1_024


class BailianEmbeddingAdapter(BailianHttpAdapter):
    """Call ``compatible-mode/v1/embeddings`` with strict input/output checks."""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str,
        model: str = "text-embedding-v4",
        dimensions: int = EMBEDDING_DIMENSIONS,
        token_counter: TokenCounter = conservative_utf8_token_count,
        http_client: httpx.AsyncClient | None = None,
        timeout_seconds: float = 30.0,
        max_retries: int = 0,
        retry_base_seconds: float = 0.5,
    ) -> None:
        if not model or model != model.strip():
            raise invalid_request("embedding.configuration", "invalid_model")
        if isinstance(dimensions, bool) or dimensions != EMBEDDING_DIMENSIONS:
            raise invalid_request("embedding.configuration", "unsupported_dimensions")
        if not callable(token_counter):
            raise invalid_request("embedding.configuration", "invalid_token_counter")

        super().__init__(
            api_key=api_key,
            base_url=base_url,
            expected_path="/compatible-mode/v1",
            http_client=http_client,
            timeout_seconds=timeout_seconds,
            max_retries=max_retries,
            retry_base_seconds=retry_base_seconds,
        )
        self._model = model
        self._dimensions = dimensions
        self._token_counter = token_counter

    async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
        operation = "embedding.create"
        validated_texts = self._validate_texts(texts, operation=operation)
        payload = {
            "model": self._model,
            "input": list(validated_texts),
            "dimensions": self._dimensions,
            "encoding_format": "float",
        }
        response = await self._post_json(
            operation=operation,
            path="embeddings",
            payload=payload,
            sensitive_values=validated_texts,
        )
        vectors = self._parse_vectors(
            response.body,
            expected_count=len(validated_texts),
            operation=operation,
        )
        return EmbeddingResult(
            vectors=vectors,
            model=response_model(
                response.body,
                self._model,
                sensitive_values=(*validated_texts, self._api_key),
            ),
            request_id=extract_request_id(
                response.body,
                sensitive_values=(*validated_texts, self._api_key),
            ),
            usage=parse_usage(response.body.get("usage")),
            elapsed_ms=response.elapsed_ms,
        )

    async def embed_query(self, text: str) -> EmbeddingResult:
        return await self.embed_documents((text,))

    def _validate_texts(
        self,
        texts: Sequence[str],
        *,
        operation: str,
    ) -> tuple[str, ...]:
        if isinstance(texts, (str, bytes)) or not isinstance(texts, Sequence):
            raise invalid_request(operation, "invalid_input_collection")
        validated = tuple(texts)
        if not validated:
            raise invalid_request(operation, "empty_input")
        if len(validated) > EMBEDDING_MAX_BATCH_SIZE:
            raise invalid_request(operation, "batch_size_exceeded")

        total_tokens = 0
        for text in validated:
            if not isinstance(text, str) or not text:
                raise invalid_request(operation, "invalid_input_text")
            token_count = count_tokens(
                self._token_counter,
                text,
                operation=operation,
            )
            if token_count > EMBEDDING_MAX_TOKENS_PER_TEXT:
                raise invalid_request(operation, "text_token_limit_exceeded")
            total_tokens += token_count
        if total_tokens > EMBEDDING_MAX_TOKENS_PER_REQUEST:
            raise invalid_request(operation, "request_token_limit_exceeded")
        return validated

    def _parse_vectors(
        self,
        body: Mapping[str, Any],
        *,
        expected_count: int,
        operation: str,
    ) -> tuple[tuple[float, ...], ...]:
        data = body.get("data")
        if not isinstance(data, list) or len(data) != expected_count:
            raise invalid_response(operation, "invalid_embedding_count")

        by_index: list[tuple[float, ...] | None] = [None] * expected_count
        for item in data:
            if not isinstance(item, Mapping):
                raise invalid_response(operation, "invalid_embedding_item")
            index = item.get("index")
            if (
                isinstance(index, bool)
                or not isinstance(index, int)
                or not 0 <= index < expected_count
                or by_index[index] is not None
            ):
                raise invalid_response(operation, "invalid_embedding_index")

            raw_vector = item.get("embedding")
            if not isinstance(raw_vector, list) or len(raw_vector) != self._dimensions:
                raise invalid_response(operation, "invalid_embedding_dimensions")

            vector: list[float] = []
            for component in raw_vector:
                if isinstance(component, bool) or not isinstance(component, (int, float)):
                    raise invalid_response(operation, "invalid_embedding_component")
                normalized = float(component)
                if not math.isfinite(normalized):
                    raise invalid_response(operation, "invalid_embedding_component")
                vector.append(normalized)

            norm = math.hypot(*vector)
            if not math.isfinite(norm) or norm <= 0:
                raise invalid_response(operation, "invalid_embedding_norm")
            by_index[index] = tuple(vector)

        if any(vector is None for vector in by_index):
            raise invalid_response(operation, "missing_embedding_index")
        return tuple(vector for vector in by_index if vector is not None)