""" Volcengine Embedding Transformation Transforms OpenAI embedding requests to Volcengine format """ from typing import List, Optional, Union, Dict, Any import httpx from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from ..common_utils import get_volcengine_base_url, get_volcengine_headers class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): """ Configuration class for Volcengine embedding models. Reference: https://ark.cn-beijing.volces.com/api/v3/embeddings """ def __init__( self, encoding_format: Optional[str] = None, ) -> None: locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): return super().get_config() def get_supported_openai_params(self, model: str) -> List[str]: """ Get the list of OpenAI parameters supported by Volcengine embedding models. Args: model: The model name Returns: List of supported parameter names """ return [ "encoding_format", "user", "extra_headers", ] def get_complete_url( self, api_base: Optional[str], api_key: Optional[str], model: str, optional_params: dict, litellm_params: dict, stream: Optional[bool] = None, ) -> str: """ Get the complete URL for volcengine embedding API calls. Args: api_base: Optional custom API base URL api_key: API key (not used for URL construction) model: Model name (not used for URL construction) optional_params: Optional parameters (not used for URL construction) litellm_params: LiteLLM parameters (not used for URL construction) stream: Stream parameter (not used for URL construction) Returns: Complete URL for the embedding API endpoint """ base_url = get_volcengine_base_url(api_base) # Construct the complete URL with /embeddings endpoint if base_url.endswith("/api/v3"): return f"{base_url}/embeddings" else: return f"{base_url}/api/v3/embeddings" def map_openai_params( self, non_default_params: Dict[str, Any], optional_params: Dict[str, Any], model: str, drop_params: bool, ) -> Dict[str, Any]: """ Map OpenAI embedding parameters to Volcengine format. Args: non_default_params: Parameters that are not default values optional_params: Optional parameters dict to update model: The model name drop_params: Whether to drop unsupported parameters Returns: Updated optional_params dict """ for param, value in non_default_params.items(): if param == "encoding_format": # Volcengine supports: float, base64, null if value in ["float", "base64", None]: optional_params["encoding_format"] = value else: if not drop_params: raise ValueError( f"Unsupported encoding_format: {value}. Volcengine supports: float, base64, null" ) elif param == "user": # Keep user parameter as-is optional_params["user"] = value elif param in self.get_supported_openai_params(model): optional_params[param] = value elif not drop_params: raise ValueError(f"Unsupported parameter for Volcengine: {param}") return optional_params def transform_embedding_request( self, model: str, input: AllEmbeddingInputValues, optional_params: dict, headers: dict, ) -> dict: """Transform embedding request to Volcengine format""" # Prepare request data (only the JSON body, not the full request) data = { "model": model, "input": input if isinstance(input, list) else [input], } # Add optional parameters from optional_params if "encoding_format" in optional_params: encoding_format = optional_params["encoding_format"] if encoding_format is not None: data["encoding_format"] = encoding_format if "user" in optional_params: user = optional_params["user"] if user is not None: data["user"] = user return data def transform_embedding_response( self, model: str, raw_response: httpx.Response, model_response: EmbeddingResponse, logging_obj: LiteLLMLoggingObj, api_key: Optional[str], request_data: dict, optional_params: dict, litellm_params: dict, ) -> EmbeddingResponse: """Transform Volcengine response to EmbeddingResponse""" try: response_json = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Volcengine response as JSON: {str(e)}") # Volcengine response format matches OpenAI format closely # Just need to ensure all required fields are present transformed_response = { "object": "list", "data": response_json.get("data", []), "model": response_json.get("model", model), "usage": response_json.get("usage", {}), } # Add id if present if "id" in response_json: transformed_response["id"] = response_json["id"] # Create EmbeddingResponse from transformed data return EmbeddingResponse(**transformed_response) def validate_environment( self, headers: dict, model: str, messages: List[AllMessageValues], optional_params: dict, litellm_params: dict, api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: """Validate environment and return headers""" # Get Volcengine headers if api_key is None: raise ValueError("api_key is required for Volcengine authentication") volcengine_headers = get_volcengine_headers(api_key) return {**headers, **volcengine_headers} def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: """Get error class for Volcengine errors""" from ..common_utils import VolcEngineError # Convert dict to httpx.Headers if needed if isinstance(headers, dict): headers = httpx.Headers(headers) return VolcEngineError( status_code=status_code, message=error_message, headers=headers, )