diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 7874201f7f0..166f876ba04 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -33,6 +33,7 @@ class BaseRerankConfig(ABC): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: return {} diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 531b94d1805..64ae8e8ffa7 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Cohere rerank") diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 60d22ff4be0..4c800d6455d 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Cohere rerank") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad908..b8635085bcc 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -978,6 +978,7 @@ class BaseLLMHTTPHandler: api_key: Optional[str] = None, api_base: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + litellm_params: Optional[Dict[str, Any]] = None, ) -> RerankResponse: # get config from model, custom llm provider headers = provider_config.validate_environment( @@ -997,6 +998,7 @@ class BaseLLMHTTPHandler: model=model, optional_rerank_params=optional_rerank_params, headers=headers, + litellm_params=litellm_params, ) ## LOGGING diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 71e300d258c..276735f4758 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: # Convert OptionalRerankParams to dict as expected by parent class if optional_rerank_params is None: diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index eb92399a058..4a7b64b9b77 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to Fireworks AI rerank format diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 8066e53afc7..60b6dc7d23d 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Hosted VLLM rerank") diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 3f83b8e422d..2c847b617ef 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Union[OptionalRerankParams, dict], headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for HuggingFace rerank") diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 48d876f8ea2..ad4416925a6 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig): return cleaned_base def transform_rerank_request( - self, model: str, optional_rerank_params: Dict, headers: Dict + self, + model: str, + optional_rerank_params: Dict, + headers: Dict, + litellm_params: Optional[dict] = None, ) -> Dict: return {"model": model, **optional_rerank_params} diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 757d874bf31..b9a46b8ac2b 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request, using clean model name without 'ranking/' prefix. @@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): model=clean_model, optional_rerank_params=optional_rerank_params, headers=headers, + litellm_params=litellm_params, ) diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index bd5abac60c8..fc317293acc 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to Nvidia NIM format. diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5895a91f3aa..a8cc5e7bdec 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -27,6 +27,46 @@ class VertexAIError(BaseLLMException): super().__init__(message=message, status_code=status_code, headers=headers) +def vertex_request_labels_from_litellm_params( + litellm_params: Optional[dict], +) -> Optional[Dict[str, str]]: + """ + Build Vertex/GCP billing labels from LiteLLM ``litellm_params["metadata"]``, + using ``requester_metadata`` string key-value pairs (same convention as Gemini). + """ + if not litellm_params or "metadata" not in litellm_params: + return None + metadata = litellm_params["metadata"] + if metadata is None or not isinstance(metadata, dict): + return None + if "requester_metadata" not in metadata: + return None + rm = metadata["requester_metadata"] + if not isinstance(rm, dict): + return None + labels = {k: v for k, v in rm.items() if isinstance(v, str)} + return labels if labels else None + + +def pop_vertex_request_labels( + optional_params: Optional[dict], + litellm_params: Optional[dict], +) -> Optional[Dict[str, str]]: + """ + Resolve labels from optional ``labels`` (Gemini-style) and/or + ``litellm_params["metadata"]["requester_metadata"]``. Pops ``labels`` from + optional_params when present. + """ + labels: Optional[Dict[str, str]] = None + if optional_params is not None and "labels" in optional_params: + raw = optional_params.pop("labels") + if isinstance(raw, dict): + labels = {k: v for k, v in raw.items() if isinstance(v, str)} + if labels is None: + labels = vertex_request_labels_from_litellm_params(litellm_params) + return labels if labels else None + + class VertexAIModelRoute(str, Enum): """Enum for Vertex AI model routing""" diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 7945c44d44c..723fa2b1616 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( response_schema_prompt, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( get_file_mime_type_for_file_type, get_file_type_from_extension, @@ -711,16 +712,8 @@ def _transform_request_body( # noqa: PLR0915 optional_params.pop("output_config", None) config_fields = GenerationConfig.__annotations__.keys() - # If the LiteLLM client sends Gemini-supported parameter "labels", add it - # as "labels" field to the request sent to the Gemini backend. - labels: Optional[dict[str, str]] = optional_params.pop("labels", None) - # If the LiteLLM client sends OpenAI-supported parameter "metadata", add it - # as "labels" field to the request sent to the Gemini backend. - if labels is None and "metadata" in litellm_params: - metadata = litellm_params["metadata"] - if metadata is not None and "requester_metadata" in metadata: - rm = metadata["requester_metadata"] - labels = {k: v for k, v in rm.items() if isinstance(v, str)} + # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) + labels = pop_vertex_request_labels(optional_params, litellm_params) filtered_params = { k: v diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 1c7696d55a2..05ebd685d91 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -7,7 +7,10 @@ import litellm from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import ( + get_vertex_base_url, + pop_vertex_request_labels, +) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): "sampleCount": 1, } - # Merge with optional params + labels = pop_vertex_request_labels(optional_params, litellm_params) + # Merge with optional params (after popping labels so they are not sent as Imagen parameters) parameters = {**default_params, **optional_params} - request_body = { + request_body: dict = { "instances": [{"prompt": prompt}], "parameters": parameters, } + if labels: + request_body["labels"] = labels return request_body diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 53651839671..3b84972e946 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -11,12 +11,15 @@ import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.llms.vertex_ai.common_utils import ( + vertex_request_labels_from_litellm_params, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + RerankBilledUnits, RerankResponse, RerankResponseMeta, - RerankBilledUnits, RerankResponseResult, ) @@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform the request from Cohere format to Vertex AI Discovery Engine format @@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # When return_documents is False, we want to ignore record details (return only IDs) request_data["ignoreRecordDetailsInResponse"] = not return_documents + user_labels = vertex_request_labels_from_litellm_params(litellm_params) + if user_labels: + request_data["userLabels"] = user_labels + return request_data def transform_rerank_response( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 5fffd983c24..729cc9c3ead 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional, Union +from typing import Dict, Literal, Optional, Union import httpx @@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase): vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None, gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, + litellm_params: Optional[Dict] = None, ) -> EmbeddingResponse: if aembedding is True: return self.async_embedding( # type: ignore @@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase): vertex_credentials=vertex_credentials, gemini_api_key=gemini_api_key, extra_headers=extra_headers, + litellm_params=litellm_params, ) should_use_v1beta1_features = self.is_using_v1beta1_features( @@ -91,7 +93,10 @@ class VertexEmbedding(VertexBase): ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) _client_params = {} @@ -154,6 +159,7 @@ class VertexEmbedding(VertexBase): gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, encoding=None, + litellm_params: Optional[Dict] = None, ) -> EmbeddingResponse: """ Async embedding implementation @@ -185,7 +191,10 @@ class VertexEmbedding(VertexBase): ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) _async_client_params = {} diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 132f29987af..a3018c5bd6b 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union from pydantic import BaseModel +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.utils import EmbeddingResponse, Usage from .types import * @@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel): return optional_params def transform_openai_request_to_vertex_embedding_request( - self, input: Union[list, str], optional_params: dict, model: str + self, + input: Union[list, str], + optional_params: dict, + model: str, + litellm_params: Optional[dict] = None, ) -> VertexEmbeddingRequest: """ Transforms an openai request to a vertex embedding request. @@ -108,14 +113,24 @@ class VertexAITextEmbeddingConfig(BaseModel): # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + labels = pop_vertex_request_labels(optional_params, litellm_params) + if model.isdigit(): - return self._transform_openai_request_to_fine_tuned_embedding_request( - input, optional_params, model + vertex_request = ( + self._transform_openai_request_to_fine_tuned_embedding_request( + input, optional_params, model + ) ) + if labels: + vertex_request["labels"] = labels + return vertex_request if VertexBGEConfig.is_bge_model(model): - return VertexBGEConfig.transform_request( + vertex_request = VertexBGEConfig.transform_request( input=input, optional_params=optional_params, model=model ) + if labels: + vertex_request["labels"] = labels + return vertex_request vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest() vertex_text_embedding_input_list: List[TextEmbeddingInput] = [] @@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel): vertex_request["instances"] = vertex_text_embedding_input_list vertex_request["parameters"] = EmbeddingParameters(**optional_params) + if labels: + vertex_request["labels"] = labels return vertex_request diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index 317b9c4fb81..bf73f4d193a 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests """ from enum import Enum -from typing import List, Optional, Union +from typing import Dict, List, Optional, Union from typing_extensions import TypedDict @@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False): List[TextEmbeddingFineTunedInput], ] parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]] + labels: Optional[Dict[str, str]] # Example usage: diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 521dae980d5..d64450a1211 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig): return api_base def transform_rerank_request( - self, model: str, optional_rerank_params: Dict, headers: Dict + self, + model: str, + optional_rerank_params: Dict, + headers: Dict, + litellm_params: Optional[dict] = None, ) -> Dict: return {"model": model, **optional_rerank_params} diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 4f8e196f254..202760f68a6 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to IBM watsonx.ai rerank format diff --git a/litellm/main.py b/litellm/main.py index 14dac7a6727..47a13d1be1e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5310,6 +5310,7 @@ def embedding( # noqa: PLR0915 api_key=api_key, api_base=api_base, client=client, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "oobabooga": response = oobabooga.embedding( diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 9868634362f..0dd18b6a5bf 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -163,19 +163,21 @@ def rerank( # noqa: PLR0915 model_response = RerankResponse() + rerank_litellm_params = { + "litellm_call_id": litellm_call_id, + "proxy_server_request": proxy_server_request, + "model_info": model_info, + "preset_cache_key": None, + "stream_response": {}, + **optional_params.model_dump(exclude_unset=True), + } + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, user=user, optional_params=dict(optional_rerank_params), - litellm_params={ - "litellm_call_id": litellm_call_id, - "proxy_server_request": proxy_server_request, - "model_info": model_info, - "preset_cache_key": None, - "stream_response": {}, - **optional_params.model_dump(exclude_unset=True), - }, + litellm_params=rerank_litellm_params, custom_llm_provider=_custom_llm_provider, ) @@ -214,6 +216,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI: api_base = ( @@ -235,6 +238,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.INFINITY: # Implement Infinity rerank logic @@ -265,6 +269,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI: # Implement Together AI rerank logic @@ -318,6 +323,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM: if dynamic_api_key is None: @@ -346,6 +352,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.BEDROCK: api_base = ( @@ -409,6 +416,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA: @@ -442,6 +450,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.FIREWORKS_AI: api_key = ( @@ -472,6 +481,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.VOYAGE: api_key = ( @@ -500,6 +510,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.WATSONX: credentials = IBMWatsonXMixin.get_watsonx_credentials( @@ -527,6 +538,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) else: # Generic handler for all providers that use base_llm_http_handler @@ -559,6 +571,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) # Placeholder return diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 350fd75d3d8..61b5e3c9249 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -67,7 +67,9 @@ class TestVertexAIGeminiImageGenerationConfig: def test_get_supported_openai_params_includes_native_gemini_params(self): """Test that native Gemini imageConfig params are supported""" - supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") + supported = self.config.get_supported_openai_params( + "gemini-3-pro-image-preview" + ) assert "aspectRatio" in supported assert "aspect_ratio" in supported assert "imageSize" in supported @@ -188,11 +190,11 @@ class TestVertexAIGeminiImageGenerationConfig: { "modality": "IMAGE", "tokenCount": 39, - } + }, ], "candidatesTokenCount": 17, "totalTokenCount": 110, - } + }, } mock_response.headers = {} @@ -219,7 +221,6 @@ class TestVertexAIGeminiImageGenerationConfig: assert result.usage.output_tokens == 17 assert result.usage.total_tokens == 110 - def test_transform_image_generation_response_multiple_images(self): """Test response transformation with multiple images""" mock_response = MagicMock(spec=httpx.Response) @@ -305,7 +306,10 @@ class TestVertexAIGeminiImageGenerationConfig: assert len(result.data) == 1 assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" + assert ( + result.data[0].provider_specific_fields["thought_signature"] + == "test_signature_abc123" + ) class TestVertexAIImagenImageGenerationConfig: @@ -369,14 +373,26 @@ class TestVertexAIImagenImageGenerationConfig: assert request["parameters"]["sampleCount"] == 2 assert request["parameters"]["aspectRatio"] == "16:9" + def test_transform_image_generation_request_labels_from_metadata(self): + """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={ + "metadata": {"requester_metadata": {"team": "platform", "env": "prod"}} + }, + headers={}, + ) + assert request["labels"] == {"team": "platform", "env": "prod"} + assert "labels" not in request["parameters"] + def test_transform_image_generation_response(self): """Test response transformation""" mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 mock_response.json.return_value = { - "predictions": [ - {"bytesBase64Encoded": "base64_encoded_image_data"} - ] + "predictions": [{"bytesBase64Encoded": "base64_encoded_image_data"}] } mock_response.headers = {} @@ -453,9 +469,7 @@ class TestGetVertexAIImageGenerationConfig: config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") assert isinstance(config, VertexAIImagenImageGenerationConfig) - config = get_vertex_ai_image_generation_config( - "vertex_ai/imagegeneration@006" - ) + config = get_vertex_ai_image_generation_config("vertex_ai/imagegeneration@006") assert isinstance(config, VertexAIImagenImageGenerationConfig) def test_get_non_gemini_model_config(self): @@ -548,4 +562,3 @@ class TestVertexAIImageGenerationIntegration: assert "us-central1" in url assert "imagegeneration@006" in url assert "predict" in url - diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 5bf2cb97fa9..fb5b0c76eaa 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -63,13 +63,16 @@ class TestVertexAIRerankTransform: import litellm # Set vertex_project attribute if it doesn't exist - if not hasattr(litellm, 'vertex_project'): + if not hasattr(litellm, "vertex_project"): litellm.vertex_project = None original_project = litellm.vertex_project litellm.vertex_project = "litellm-project-456" # Reset mock call count mock_ensure_access_token.reset_mock() - mock_ensure_access_token.return_value = ("mock-token", "litellm-project-456") + mock_ensure_access_token.return_value = ( + "mock-token", + "litellm-project-456", + ) try: url = self.config.get_complete_url(api_base=None, model=self.model) expected_url = "https://discoveryengine.googleapis.com/v1/projects/litellm-project-456/locations/global/rankingConfigs/default_ranking_config:rank" @@ -82,15 +85,19 @@ class TestVertexAIRerankTransform: import litellm # Set vertex_project to None to ensure no project ID is available - if not hasattr(litellm, 'vertex_project'): + if not hasattr(litellm, "vertex_project"): litellm.vertex_project = None original_project = litellm.vertex_project litellm.vertex_project = None # Reset mock and set it to raise an error mock_ensure_access_token.reset_mock() - mock_ensure_access_token.side_effect = ValueError("Vertex AI project ID is required") + mock_ensure_access_token.side_effect = ValueError( + "Vertex AI project ID is required" + ) try: - with pytest.raises(ValueError, match="Vertex AI project ID is required"): + with pytest.raises( + ValueError, match="Vertex AI project ID is required" + ): self.config.get_complete_url(api_base=None, model=self.model) finally: litellm.vertex_project = original_project @@ -109,15 +116,13 @@ class TestVertexAIRerankTransform: self.config._ensure_access_token = mock_ensure_access_token headers = self.config.validate_environment( - headers={}, - model=self.model, - api_key=None + headers={}, model=self.model, api_key=None ) expected_headers = { "Authorization": "Bearer test-access-token", "Content-Type": "application/json", - "X-Goog-User-Project": "test-project-123" + "X-Goog-User-Project": "test-project-123", } assert headers == expected_headers @@ -127,24 +132,22 @@ class TestVertexAIRerankTransform: "query": "What is Google Gemini?", "documents": [ "Gemini is a cutting edge large language model created by Google.", - "The Gemini zodiac symbol often depicts two figures standing side-by-side." + "The Gemini zodiac symbol often depicts two figures standing side-by-side.", ], - "top_n": 2 + "top_n": 2, } - + request_data = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params, - headers={} + model=self.model, optional_rerank_params=optional_params, headers={} ) - + # Verify basic structure assert request_data["model"] == self.model assert request_data["query"] == "What is Google Gemini?" assert request_data["topN"] == 2 assert "records" in request_data assert len(request_data["records"]) == 2 - + # Verify record structure for i, record in enumerate(request_data["records"]): assert "id" in record @@ -158,20 +161,25 @@ class TestVertexAIRerankTransform: optional_params = { "query": "What is Google Gemini?", "documents": [ - {"text": "Gemini is a cutting edge large language model created by Google.", "title": "Custom Title 1"}, - {"text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."} - ] + { + "text": "Gemini is a cutting edge large language model created by Google.", + "title": "Custom Title 1", + }, + { + "text": "The Gemini zodiac symbol often depicts two figures standing side-by-side." + }, + ], } - + request_data = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params, - headers={} + model=self.model, optional_rerank_params=optional_params, headers={} ) - + # Verify record structure with custom titles assert request_data["records"][0]["title"] == "Custom Title 1" - assert request_data["records"][1]["title"] == "The Gemini zodiac" # First 3 words + assert ( + request_data["records"][1]["title"] == "The Gemini zodiac" + ) # First 3 words def test_transform_rerank_request_return_documents_mapping(self): """Test return_documents to ignoreRecordDetailsInResponse mapping.""" @@ -179,43 +187,50 @@ class TestVertexAIRerankTransform: optional_params_true = { "query": "test query", "documents": ["doc1", "doc2"], - "return_documents": True + "return_documents": True, } - + request_data_true = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params_true, - headers={} + model=self.model, optional_rerank_params=optional_params_true, headers={} ) assert request_data_true["ignoreRecordDetailsInResponse"] == False - + # Test return_documents=False optional_params_false = { "query": "test query", "documents": ["doc1", "doc2"], - "return_documents": False + "return_documents": False, } - + request_data_false = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params_false, - headers={} + model=self.model, optional_rerank_params=optional_params_false, headers={} ) assert request_data_false["ignoreRecordDetailsInResponse"] == True - + # Test return_documents not specified (should default to True) - optional_params_default = { - "query": "test query", - "documents": ["doc1", "doc2"] - } - + optional_params_default = {"query": "test query", "documents": ["doc1", "doc2"]} + request_data_default = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params_default, - headers={} + model=self.model, optional_rerank_params=optional_params_default, headers={} ) assert request_data_default["ignoreRecordDetailsInResponse"] == False + def test_transform_rerank_request_user_labels_from_metadata(self): + """Discovery Engine Rank API uses userLabels (string map) for billing.""" + optional_params = { + "query": "q", + "documents": ["a", "b"], + } + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params=optional_params, + headers={}, + litellm_params={ + "metadata": {"requester_metadata": {"app": "litellm", "tier": "1"}} + }, + ) + assert request_data["userLabels"] == {"app": "litellm", "tier": "1"} + def test_transform_rerank_request_missing_required_params(self): """Test that transform_rerank_request handles missing required parameters.""" # Test missing query @@ -223,15 +238,17 @@ class TestVertexAIRerankTransform: self.config.transform_rerank_request( model=self.model, optional_rerank_params={"documents": ["doc1"]}, - headers={} + headers={}, ) - + # Test missing documents - with pytest.raises(ValueError, match="documents is required for Vertex AI rerank"): + with pytest.raises( + ValueError, match="documents is required for Vertex AI rerank" + ): self.config.transform_rerank_request( model=self.model, optional_rerank_params={"query": "test query"}, - headers={} + headers={}, ) def test_transform_rerank_response_success(self): @@ -243,34 +260,34 @@ class TestVertexAIRerankTransform: "id": "1", "score": 0.98, "title": "The Science of a Blue Sky", - "content": "The sky appears blue due to a phenomenon called Rayleigh scattering." + "content": "The sky appears blue due to a phenomenon called Rayleigh scattering.", }, { "id": "0", "score": 0.64, "title": "The Color of the Sky: A Poem", - "content": "A canvas stretched across the day, Where sunlight learns to dance and play." - } + "content": "A canvas stretched across the day, Where sunlight learns to dance and play.", + }, ] } - + # Create mock httpx response mock_response = MagicMock(spec=httpx.Response) mock_response.json.return_value = response_data mock_response.text = json.dumps(response_data) - + # Create mock logging object mock_logging = MagicMock() - + model_response = RerankResponse() - + result = self.config.transform_rerank_response( model=self.model, raw_response=mock_response, model_response=model_response, logging_obj=mock_logging, ) - + # Verify response structure assert result.id == f"vertex_ai_rerank_{self.model}" assert len(result.results) == 2 @@ -278,34 +295,29 @@ class TestVertexAIRerankTransform: assert result.results[0]["relevance_score"] == 0.98 assert result.results[1]["index"] == 0 assert result.results[1]["relevance_score"] == 0.64 - + # Verify metadata assert result.meta["billed_units"]["search_units"] == 2 def test_transform_rerank_response_with_ignore_record_details(self): """Test response transformation when ignoreRecordDetailsInResponse=true.""" # Mock response with only IDs (when ignoreRecordDetailsInResponse=true) - response_data = { - "records": [ - {"id": "1"}, - {"id": "0"} - ] - } - + response_data = {"records": [{"id": "1"}, {"id": "0"}]} + mock_response = MagicMock(spec=httpx.Response) mock_response.json.return_value = response_data mock_response.text = json.dumps(response_data) - + mock_logging = MagicMock() model_response = RerankResponse() - + result = self.config.transform_rerank_response( model=self.model, raw_response=mock_response, model_response=model_response, logging_obj=mock_logging, ) - + # Verify response structure with default scores assert len(result.results) == 2 assert result.results[0]["index"] == 1 # 0-based index @@ -318,10 +330,10 @@ class TestVertexAIRerankTransform: mock_response = MagicMock(spec=httpx.Response) mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "doc", 0) mock_response.text = "Invalid JSON response" - + mock_logging = MagicMock() model_response = RerankResponse() - + with pytest.raises(ValueError, match="Failed to parse response"): self.config.transform_rerank_response( model=self.model, @@ -345,14 +357,14 @@ class TestVertexAIRerankTransform: query="test query", documents=["doc1", "doc2"], top_n=2, - return_documents=True + return_documents=True, ) - + expected_params = { "query": "test query", "documents": ["doc1", "doc2"], "top_n": 2, - "return_documents": True + "return_documents": True, } assert params == expected_params @@ -363,34 +375,32 @@ class TestVertexAIRerankTransform: "documents": [ "This is a very long document with many words that should be truncated to only the first three words for the title", "Short doc", - "Another document with multiple words here" - ] + "Another document with multiple words here", + ], } - + request_data = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params, - headers={} + model=self.model, optional_rerank_params=optional_params, headers={} ) - + # Verify title generation assert request_data["records"][0]["title"] == "This is a" # First 3 words assert request_data["records"][1]["title"] == "Short doc" # Less than 3 words - assert request_data["records"][2]["title"] == "Another document with" # First 3 words + assert ( + request_data["records"][2]["title"] == "Another document with" + ) # First 3 words def test_record_id_generation(self): """Test that record IDs are generated correctly with 0-based indexing.""" optional_params = { "query": "test query", - "documents": ["doc1", "doc2", "doc3", "doc4"] + "documents": ["doc1", "doc2", "doc3", "doc4"], } - + request_data = self.config.transform_rerank_request( - model=self.model, - optional_rerank_params=optional_params, - headers={} + model=self.model, optional_rerank_params=optional_params, headers={} ) - + # Verify 0-based indexing for i, record in enumerate(request_data["records"]): assert record["id"] == str(i) @@ -402,9 +412,9 @@ class TestVertexAIRerankTransform: "documents": ["doc1", "doc2"], "vertex_credentials": "path/to/credentials.json", "vertex_project": "my-project-id", - "vertex_location": "us-central1" + "vertex_location": "us-central1", } - + params = self.config.map_cohere_rerank_params( non_default_params=non_default_params, model=self.model, @@ -412,14 +422,14 @@ class TestVertexAIRerankTransform: query="test query", documents=["doc1", "doc2"], top_n=2, - return_documents=True + return_documents=True, ) - + # Verify vertex-specific parameters are preserved assert params["vertex_credentials"] == "path/to/credentials.json" assert params["vertex_project"] == "my-project-id" assert params["vertex_location"] == "us-central1" - + # Verify standard params are still present assert params["query"] == "test query" assert params["documents"] == ["doc1", "doc2"] @@ -428,10 +438,8 @@ class TestVertexAIRerankTransform: def test_map_cohere_rerank_params_without_vertex_credentials(self): """Test that map_cohere_rerank_params works when vertex credentials are not provided.""" - non_default_params = { - "documents": ["doc1", "doc2"] - } - + non_default_params = {"documents": ["doc1", "doc2"]} + params = self.config.map_cohere_rerank_params( non_default_params=non_default_params, model=self.model, @@ -439,14 +447,14 @@ class TestVertexAIRerankTransform: query="test query", documents=["doc1", "doc2"], top_n=2, - return_documents=True + return_documents=True, ) - + # Verify no vertex-specific parameters are added when not provided assert "vertex_credentials" not in params assert "vertex_project" not in params assert "vertex_location" not in params - + # Verify standard params are still present assert params["query"] == "test query" assert params["documents"] == ["doc1", "doc2"] @@ -470,14 +478,11 @@ class TestVertexAIRerankTransform: "vertex_credentials": "path/to/credentials.json", "vertex_project": "custom-project-id", "query": "test query", - "documents": ["doc1"] + "documents": ["doc1"], } headers = self.config.validate_environment( - headers={}, - model=self.model, - api_key=None, - optional_params=optional_params + headers={}, model=self.model, api_key=None, optional_params=optional_params ) # Verify that _ensure_access_token was called with the credentials from optional_params @@ -490,7 +495,7 @@ class TestVertexAIRerankTransform: expected_headers = { "Authorization": "Bearer test-access-token", "Content-Type": "application/json", - "X-Goog-User-Project": "test-project-123" + "X-Goog-User-Project": "test-project-123", } assert headers == expected_headers @@ -527,7 +532,10 @@ class TestVertexAIRerankTransform: assert optional_params["vertex_project"] == "custom-project-id" # get_complete_url should still be able to access the vertex params - with patch('litellm.llms.vertex_ai.rerank.transformation.get_secret_str', return_value=None): + with patch( + "litellm.llms.vertex_ai.rerank.transformation.get_secret_str", + return_value=None, + ): url = self.config.get_complete_url( api_base=None, model=self.model, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index d483a81a349..8a5c397d6f4 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -15,7 +15,9 @@ from litellm.llms.vertex_ai.common_utils import ( convert_anyof_null_to_nullable, get_vertex_location_from_url, get_vertex_project_id_from_url, + pop_vertex_request_labels, set_schema_property_ordering, + vertex_request_labels_from_litellm_params, ) @@ -440,7 +442,9 @@ def test_vertex_ai_complex_response_schema(): optional_params = {} v.apply_response_schema_transformation( - value=non_default_params["response_format"], optional_params=optional_params, model="gemini-1.5-pro-preview-0409" + value=non_default_params["response_format"], + optional_params=optional_params, + model="gemini-1.5-pro-preview-0409", ) # Assertions for the transformed schema @@ -558,7 +562,6 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): assert url == expected_url - @pytest.mark.parametrize( "model_cost_entry, vertex_region, expected_region", [ @@ -571,9 +574,17 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): # Model with supported_regions=["us-west2"], no user region -> use "us-west2" ({"supported_regions": ["us-west2"]}, None, "us-west2"), # Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it - ({"supported_regions": ["us-west2", "us-central1"]}, "us-central1", "us-central1"), + ( + {"supported_regions": ["us-west2", "us-central1"]}, + "us-central1", + "us-central1", + ), # Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override - ({"supported_regions": ["us-west2", "us-central1"]}, "europe-west1", "us-west2"), + ( + {"supported_regions": ["us-west2", "us-central1"]}, + "europe-west1", + "us-west2", + ), # No model_cost entry, no user region -> default us-central1 ({}, None, "us-central1"), # No model_cost entry, user specifies region -> use specified region @@ -656,11 +667,12 @@ def test_vertex_filter_format_uri(): assert "uri" not in json.dumps(new_parameters) + def test_convert_schema_types_type_array_conversion(): """ Test _convert_schema_types function handles type arrays and case conversion. - - This test verifies the fix for the issue where type arrays like ["string", "number"] + + This test verifies the fix for the issue where type arrays like ["string", "number"] would raise an exception in Vertex AI schema validation. Relevant issue: https://github.com/BerriAI/litellm/issues/14091 @@ -673,12 +685,12 @@ def test_convert_schema_types_type_array_conversion(): "properties": { "studio": { "type": ["string", "number"], - "description": "The studio ID or name" + "description": "The studio ID or name", } }, "required": ["studio"], "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#" + "$schema": "http://json-schema.org/draft-07/schema#", } # Expected output: Vertex AI compatible schema with anyOf and uppercase types @@ -686,16 +698,13 @@ def test_convert_schema_types_type_array_conversion(): "type": "object", "properties": { "studio": { - "anyOf": [ - {"type": "string"}, - {"type": "number"} - ], - "description": "The studio ID or name" + "anyOf": [{"type": "string"}, {"type": "number"}], + "description": "The studio ID or name", } }, "required": ["studio"], "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#" + "$schema": "http://json-schema.org/draft-07/schema#", } # Apply the transformation @@ -718,15 +727,17 @@ def test_convert_schema_types_type_array_conversion(): assert anyof_types[1]["type"] == "number" # 4. Other properties preserved - assert input_schema["properties"]["studio"]["description"] == "The studio ID or name" + assert ( + input_schema["properties"]["studio"]["description"] == "The studio ID or name" + ) assert input_schema["required"] == ["studio"] def test_fix_enum_empty_strings(): """ Test _fix_enum_empty_strings function replaces empty strings with None in enum arrays. - - This test verifies the fix for the issue where Gemini rejects tool definitions + + This test verifies the fix for the issue where Gemini rejects tool definitions with empty strings in enum values, causing API failures. Relevant issue: Gemini does not accept empty strings in enum values @@ -740,23 +751,23 @@ def test_fix_enum_empty_strings(): "user_agent_type": { "enum": ["", "desktop", "mobile", "tablet"], "type": "string", - "description": "Device type for user agent" + "description": "Device type for user agent", } }, - "required": ["user_agent_type"] + "required": ["user_agent_type"], } # Expected output: Empty strings replaced with None expected_output = { - "type": "object", + "type": "object", "properties": { "user_agent_type": { "enum": [None, "desktop", "mobile", "tablet"], "type": "string", - "description": "Device type for user agent" + "description": "Device type for user agent", } }, - "required": ["user_agent_type"] + "required": ["user_agent_type"], } # Apply the transformation @@ -859,7 +870,7 @@ def test_construct_target_url_with_version_prefix(): def test_fix_enum_types(): """ Test _fix_enum_types function removes enum fields when type is not string. - + This test verifies the fix for the issue where Gemini rejects cached content with function parameter enums on non-string types, causing API failures. @@ -874,38 +885,41 @@ def test_fix_enum_types(): "truncateMode": { "enum": ["auto", "none", "start", "end"], "type": "string", # This should keep the enum - "description": "How to truncate content" + "description": "How to truncate content", }, "maxLength": { "enum": [100, 200, 500], # This should be removed "type": "integer", - "description": "Maximum length" + "description": "Maximum length", }, "enabled": { "enum": [True, False], # This should be removed "type": "boolean", - "description": "Whether feature is enabled" + "description": "Whether feature is enabled", }, "nested": { "type": "object", "properties": { "innerEnum": { "enum": ["a", "b", "c"], # This should be kept - "type": "string" + "type": "string", }, "innerNonStringEnum": { "enum": [1, 2, 3], # This should be removed - "type": "integer" - } - } + "type": "integer", + }, + }, }, "anyOfField": { "anyOf": [ - {"type": "string", "enum": ["option1", "option2"]}, # This should be kept - {"type": "integer", "enum": [1, 2, 3]} # This should be removed + { + "type": "string", + "enum": ["option1", "option2"], + }, # This should be kept + {"type": "integer", "enum": [1, 2, 3]}, # This should be removed ] - } - } + }, + }, } # Expected output: Non-string enums removed, string enums kept @@ -919,31 +933,32 @@ def test_fix_enum_types(): }, "maxLength": { # enum removed "type": "integer", - "description": "Maximum length" + "description": "Maximum length", }, "enabled": { # enum removed "type": "boolean", - "description": "Whether feature is enabled" + "description": "Whether feature is enabled", }, "nested": { "type": "object", "properties": { "innerEnum": { "enum": ["a", "b", "c"], # Kept - string type - "type": "string" + "type": "string", }, - "innerNonStringEnum": { # enum removed - "type": "integer" - } - } + "innerNonStringEnum": {"type": "integer"}, # enum removed + }, }, "anyOfField": { "anyOf": [ - {"type": "string", "enum": ["option1", "option2"]}, # Kept - has string type - {"type": "integer"} # enum removed + { + "type": "string", + "enum": ["option1", "option2"], + }, # Kept - has string type + {"type": "integer"}, # enum removed ] - } - } + }, + }, } # Apply the transformation @@ -955,15 +970,27 @@ def test_fix_enum_types(): # Verify specific transformations: # 1. String enums are preserved assert "enum" in input_schema["properties"]["truncateMode"] - assert input_schema["properties"]["truncateMode"]["enum"] == ["auto", "none", "start", "end"] - + assert input_schema["properties"]["truncateMode"]["enum"] == [ + "auto", + "none", + "start", + "end", + ] + assert "enum" in input_schema["properties"]["nested"]["properties"]["innerEnum"] - assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == ["a", "b", "c"] + assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == [ + "a", + "b", + "c", + ] # 2. Non-string enums are removed assert "enum" not in input_schema["properties"]["maxLength"] assert "enum" not in input_schema["properties"]["enabled"] - assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] + assert ( + "enum" + not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] + ) # 3. anyOf with string type keeps enum, non-string removes it assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0] @@ -1003,8 +1030,6 @@ def test_get_token_url(): print("url=", url) - - should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features( optional_params={"temperature": 0.1} ) @@ -1210,9 +1235,7 @@ def test_vertex_ai_minimax_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "minimaxai/minimax-m2-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("minimaxai/minimax-m2-maas") def test_vertex_ai_moonshot_uses_openai_handler(): @@ -1236,9 +1259,7 @@ def test_vertex_ai_zai_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "zai-org/glm-4.7-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("zai-org/glm-4.7-maas") def test_vertex_ai_zai_is_partner_model(): @@ -1255,14 +1276,14 @@ def test_vertex_ai_zai_is_partner_model(): def test_build_vertex_schema_empty_properties(): """ Test _build_vertex_schema handles empty properties objects correctly. - - This test verifies the fix for the issue where Gemini rejects schemas + + This test verifies the fix for the issue where Gemini rejects schemas with empty properties objects like {"properties": {}, "type": "object"}. - + Error from Gemini: "GenerateContentRequest.generation_config.response_schema - .properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties: + .properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties: should be non-empty for OBJECT type" - + The fix removes empty properties objects and their associated type/required fields. """ from litellm.llms.vertex_ai.common_utils import _build_vertex_schema @@ -1281,20 +1302,20 @@ def test_build_vertex_schema_empty_properties(): "type": "object", "additionalProperties": False, "description": "Go back", - "required": [] + "required": [], } }, "required": ["go_back"], "type": "object", - "additionalProperties": False + "additionalProperties": False, } ] }, - "type": "array" + "type": "array", } }, "type": "object", - "additionalProperties": False + "additionalProperties": False, } # Apply the transformation @@ -1302,24 +1323,36 @@ def test_build_vertex_schema_empty_properties(): # Verify the transformation removed empty properties # Navigate to the go_back schema - go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"]["go_back"] - + go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"][ + "go_back" + ] + # Verify empty properties was removed assert "properties" not in go_back_schema, "Empty properties should be removed" - + # Verify type is kept as object (Gemini requires type: object even without properties) - assert go_back_schema.get("type") == "object", "Type should be kept as object when properties is empty" - + assert ( + go_back_schema.get("type") == "object" + ), "Type should be kept as object when properties is empty" + # Verify required was also removed - assert "required" not in go_back_schema, "Required should be removed when properties is empty" - + assert ( + "required" not in go_back_schema + ), "Required should be removed when properties is empty" + # Verify description is preserved - assert go_back_schema.get("description") == "Go back", "Description should be preserved" - + assert ( + go_back_schema.get("description") == "Go back" + ), "Description should be preserved" + # Verify parent schema still has proper structure parent_schema = result["properties"]["action"]["items"]["anyOf"][0] - assert parent_schema["type"] == "object", "Parent schema should still have object type" - assert "go_back" in parent_schema["properties"], "go_back should still be in parent properties" + assert ( + parent_schema["type"] == "object" + ), "Parent schema should still have object type" + assert ( + "go_back" in parent_schema["properties"] + ), "go_back should still be in parent properties" def test_add_object_type_schema_with_no_properties_and_no_type(): @@ -1330,9 +1363,7 @@ def test_add_object_type_schema_with_no_properties_and_no_type(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with no properties and no type (the problematic case) - input_schema = { - "$schema": "https://json-schema.org/draft/2020-12/schema" - } + input_schema = {"$schema": "https://json-schema.org/draft/2020-12/schema"} # Apply the transformation add_object_type(input_schema) @@ -1351,10 +1382,7 @@ def test_add_object_type_does_not_override_existing_type(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with existing type - input_schema = { - "type": "string", - "description": "A string field" - } + input_schema = {"type": "string", "description": "A string field"} # Apply the transformation add_object_type(input_schema) @@ -1370,15 +1398,42 @@ def test_add_object_type_does_not_add_type_when_anyof_present(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with anyOf but no type - input_schema = { - "anyOf": [ - {"type": "string"}, - {"type": "null"} - ] - } + input_schema = {"anyOf": [{"type": "string"}, {"type": "null"}]} # Apply the transformation add_object_type(input_schema) # Verify type was not added (anyOf handles the type) assert "type" not in input_schema, "type should not be added when anyOf is present" + + +def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata(): + assert vertex_request_labels_from_litellm_params(None) is None + assert vertex_request_labels_from_litellm_params({}) is None + assert vertex_request_labels_from_litellm_params({"metadata": None}) is None + lp = {"metadata": {"requester_metadata": {"team": "analytics", "count": 3}}} + assert vertex_request_labels_from_litellm_params(lp) == {"team": "analytics"} + + +def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata(): + optional = {"labels": {"env": "prod"}} + litellm_params = {"metadata": {"requester_metadata": {"team": "x"}}} + assert pop_vertex_request_labels(optional, litellm_params) == {"env": "prod"} + assert "labels" not in optional + + optional2: dict = {} + assert pop_vertex_request_labels(optional2, litellm_params) == {"team": "x"} + + +def test_vertex_text_embedding_request_includes_labels_from_metadata(): + import litellm + + req = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input="hi", + optional_params={}, + model="text-embedding-004", + litellm_params={ + "metadata": {"requester_metadata": {"project_id": "cost-center-1"}} + }, + ) + assert req.get("labels") == {"project_id": "cost-center-1"}