"""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)