mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
Merge pull request #24963 from BerriAI/litellm_vertex_request_metadata_labels
feat(vertex_ai): propagate metadata labels to embedding, Imagen, rerank
This commit is contained in:
Sameer Kankute
GitHub
commit
edb42fc107
25 changed files with 426 additions and 241 deletions
|
|
@ -33,6 +33,7 @@ class BaseRerankConfig(ABC):
|
|||
model: str,
|
||||
optional_rerank_params: Dict,
|
||||
headers: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> dict:
|
||||
return {}
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue