mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(vertex_ai): batch Claude models in Anthropic's native row shape (#45715)
Backport of #45715 to lts/1.108.x.
Cherry-picked from a7030a5a2c (main).
Adapted to the line: the Gemini row mapper keeps the line's signature, since
deployment_drop_params on it comes from #31915, which is not on the line, and
the Claude row reads drop_params with the line's normalize_drop_params, since
drops_unsupported_params is not in core_helpers here.
This commit is contained in:
Gopal Kamineni
mateo-berri
parent
bf06bf1edb
commit
32469bcab9
7 changed files with 860 additions and 69 deletions
|
|
@ -14,8 +14,10 @@ from litellm.litellm_core_utils.url_utils import (
|
|||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
HTTPHandler,
|
||||
MaskedHTTPStatusError,
|
||||
get_async_httpx_client,
|
||||
get_httpx_client,
|
||||
mask_sensitive_info,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
|
|
@ -80,6 +82,15 @@ def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str
|
|||
return response.json()
|
||||
|
||||
|
||||
def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatusError:
|
||||
body: Final = mask_sensitive_info(error.response.text)
|
||||
return MaskedHTTPStatusError(
|
||||
httpx.HTTPStatusError(f"{error}\n{body}", request=error.request, response=error.response),
|
||||
message=body,
|
||||
text=body,
|
||||
)
|
||||
|
||||
|
||||
class VertexAIBatchPrediction(VertexLLM):
|
||||
def __init__(self, gcs_bucket_name: str, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
|
@ -165,11 +176,19 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
headers=headers,
|
||||
)
|
||||
|
||||
response: Final = sync_handler.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
data=json.dumps(vertex_batch_request),
|
||||
)
|
||||
try:
|
||||
response: Final = sync_handler.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
data=json.dumps(vertex_batch_request),
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch create failed: status=%s, body=%s",
|
||||
e.response.status_code,
|
||||
e.response.text[:1000],
|
||||
)
|
||||
raise _vertex_batch_create_error(e) from e
|
||||
|
||||
payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()}
|
||||
_json_response: Final = payload_view["payload"]
|
||||
|
|
@ -265,13 +284,12 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
data=json.dumps(vertex_batch_request),
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_body: Final = e.response.text
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch create failed: status=%s, body=%s",
|
||||
e.response.status_code,
|
||||
error_body[:1000],
|
||||
e.response.text[:1000],
|
||||
)
|
||||
raise
|
||||
raise _vertex_batch_create_error(e) from e
|
||||
|
||||
payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()}
|
||||
_json_response: Final = payload_view["payload"]
|
||||
|
|
|
|||
|
|
@ -9,14 +9,41 @@ from litellm._logging import verbose_logger
|
|||
from litellm._uuid import uuid
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
VertexAIError,
|
||||
VertexAIModelRoute,
|
||||
convert_vertex_datetime_to_openai_datetime,
|
||||
get_vertex_ai_model_route,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
|
||||
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
|
||||
from litellm.types.llms.vertex_ai import *
|
||||
from litellm.types.llms.vertex_ai import GenerateContentResponseBody
|
||||
from litellm.types.llms.vertex_ai import GenerateContentResponseBody, VertexPartnerProvider
|
||||
from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage
|
||||
|
||||
_NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody)
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
VERTEX_ANTHROPIC_VERSION: Final = "vertex-2023-10-16"
|
||||
_ANTHROPIC_PUBLISHER_MODEL_PATH: Final = "publishers/anthropic/models/"
|
||||
|
||||
|
||||
def is_vertex_anthropic_batch_model(model: str) -> bool:
|
||||
"""
|
||||
Whether a Vertex batch model (bare, `vertex_ai/`-prefixed or a publisher path) is a Claude model, decided the
|
||||
way Vertex chat routes Claude. Vertex batches Claude only as `publishers/anthropic/models/<model>` with rows in
|
||||
Anthropic's Messages shape
|
||||
"""
|
||||
bare_model: Final = model.rsplit("/", 1)[-1]
|
||||
return _ANTHROPIC_PUBLISHER_MODEL_PATH in model or (
|
||||
get_vertex_ai_model_route(bare_model) is VertexAIModelRoute.PARTNER_MODELS
|
||||
and VertexPartnerProvider.claude.value in bare_model
|
||||
)
|
||||
|
||||
|
||||
def as_json_object(value: object) -> Mapping[str, object] | None:
|
||||
return _JSON_OBJECT.validate_python(value) if isinstance(value, dict) else None
|
||||
|
||||
|
||||
def is_vertex_anthropic_batch_output_response(response_body: Mapping[str, object]) -> bool:
|
||||
return response_body.get("type") == "message" and isinstance(response_body.get("usage"), dict)
|
||||
|
||||
|
||||
def _int_field(mapping: Mapping[str, object], key: str) -> int:
|
||||
|
|
@ -43,8 +70,22 @@ def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object
|
|||
return isinstance(response_body.get("embedding"), dict)
|
||||
|
||||
|
||||
def is_vertex_anthropic_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool:
|
||||
"""
|
||||
Whether a Vertex batch output row came from a Claude (Anthropic Messages) request. Vertex echoes only the
|
||||
request's top-level `custom_id` on these rows, with an Anthropic `message` response on success and an empty
|
||||
response plus a `status` on failure, so a Claude input file read back is never mistaken for output
|
||||
"""
|
||||
if not ("custom_id" in vertex_output_row and "processed_time" in vertex_output_row):
|
||||
return False
|
||||
response: Final = as_json_object(vertex_output_row.get("response"))
|
||||
return (response is not None and is_vertex_anthropic_batch_output_response(response)) or bool(
|
||||
vertex_output_row.get("status")
|
||||
)
|
||||
|
||||
|
||||
def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool:
|
||||
return isinstance(row.get("request"), dict)
|
||||
return isinstance(row.get("request"), dict) or is_vertex_anthropic_batch_output_row(row)
|
||||
|
||||
|
||||
class NativeVertexBatchCostCalculator(Protocol):
|
||||
|
|
@ -72,6 +113,9 @@ def _native_vertex_row_usage(
|
|||
response_body: Mapping[str, object],
|
||||
calculate_usage: Callable[[GenerateContentResponseBody], Usage],
|
||||
) -> Usage | None:
|
||||
anthropic_usage: Final = as_json_object(response_body.get("usage"))
|
||||
if is_vertex_anthropic_batch_output_response(response_body) and anthropic_usage is not None:
|
||||
return VertexAIAnthropicConfig().calculate_usage(usage_object=anthropic_usage, reasoning_content=None)
|
||||
if "usageMetadata" not in response_body:
|
||||
if not is_vertex_embedding_batch_output_response(response_body):
|
||||
return None
|
||||
|
|
@ -97,21 +141,24 @@ def native_vertex_batch_row_stats(
|
|||
"""
|
||||
Usage and cost of one native Vertex predictions.jsonl row, a
|
||||
`{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}`
|
||||
generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}`
|
||||
embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`).
|
||||
`model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion`
|
||||
does, else the wildcard name so explicit deployment prices still apply; a row without a response, a
|
||||
generateContent row without `response.usageMetadata`, and a row whose response fails validation are
|
||||
generateContent object, a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}`
|
||||
embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`) or a
|
||||
`{"custom_id": ..., "request": ..., "response": {"type": "message", "model": ..., "usage": {...}}}` Claude object.
|
||||
`model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` (Claude
|
||||
rows: `model`) does, else the wildcard name so explicit deployment prices still apply; a row without a response,
|
||||
a generateContent row without `response.usageMetadata`, and a row whose response fails validation are
|
||||
None (failed).
|
||||
"""
|
||||
response_body: Final = row.get("response")
|
||||
if not isinstance(response_body, dict):
|
||||
response_body: Final = as_json_object(row.get("response"))
|
||||
if response_body is None:
|
||||
return None
|
||||
usage: Final = _native_vertex_row_usage(response_body, calculate_usage)
|
||||
if usage is None:
|
||||
return None
|
||||
total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens)
|
||||
model_version: Final = response_body.get("modelVersion")
|
||||
model_version: Final = response_body.get(
|
||||
"model" if is_vertex_anthropic_batch_output_response(response_body) else "modelVersion"
|
||||
)
|
||||
deployment_model: Final = model_name if model_name and "*" not in model_name else None
|
||||
model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name)
|
||||
if model is None:
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from typing_extensions import ReadOnly, Required
|
|||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES, RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
from litellm.files.utils import FilesAPIUtils
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
|
|
@ -30,11 +31,13 @@ from litellm.litellm_core_utils.cloud_storage_security import (
|
|||
split_configured_cloud_bucket_name,
|
||||
validate_managed_cloud_file_id,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_file_data,
|
||||
extract_file_metadata,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY
|
||||
from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body
|
||||
|
|
@ -43,7 +46,14 @@ from litellm.llms.base_llm.files.transformation import (
|
|||
BaseFileUploadStream,
|
||||
LiteLLMLoggingObj,
|
||||
)
|
||||
from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count
|
||||
from litellm.llms.vertex_ai.batches.transformation import (
|
||||
VERTEX_ANTHROPIC_VERSION,
|
||||
as_json_object,
|
||||
is_vertex_anthropic_batch_model,
|
||||
is_vertex_anthropic_batch_output_response,
|
||||
is_vertex_anthropic_batch_output_row,
|
||||
vertex_embedding_prompt_token_count,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
convert_vertex_datetime_to_openai_datetime,
|
||||
get_vertex_ai_fine_tuned_endpoint_id,
|
||||
|
|
@ -55,6 +65,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
|||
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
|
||||
transform_openai_input_gemini_embed_content,
|
||||
)
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
|
||||
from litellm.types.files import StreamingMediaUploadConfig
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -75,6 +86,7 @@ from litellm.types.utils import (
|
|||
ModelResponse,
|
||||
Usage,
|
||||
)
|
||||
from litellm.utils import get_non_default_completion_params, get_optional_params
|
||||
|
||||
from ..common_utils import VertexAIError
|
||||
from ..vertex_llm_base import VertexBase
|
||||
|
|
@ -83,6 +95,7 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63
|
|||
_CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_"
|
||||
_VERTEX_BATCH_KEY_FIELD: Final = "key"
|
||||
_MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)")
|
||||
_PUBLISHER_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/")
|
||||
_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = (
|
||||
("outputDimensionality", "output_dimensionality"),
|
||||
("taskType", "task_type"),
|
||||
|
|
@ -220,10 +233,14 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str
|
|||
"""
|
||||
Resolve the OpenAI `custom_id` for a Vertex batch output row.
|
||||
|
||||
Embedding rows carry it in the top-level `key` field that Vertex echoes back;
|
||||
`generateContent` rows have no such field, so it is smuggled through request
|
||||
labels instead (see `_set_litellm_batch_custom_id_labels`).
|
||||
Claude rows carry it in their own top-level `custom_id` field and embedding rows in
|
||||
the top-level `key` field that Vertex echoes back; `generateContent` rows have no
|
||||
such field, so it is smuggled through request labels instead (see
|
||||
`_set_litellm_batch_custom_id_labels`).
|
||||
"""
|
||||
anthropic_custom_id: Final = vertex_output_row.get("custom_id")
|
||||
if anthropic_custom_id is not None:
|
||||
return str(anthropic_custom_id)
|
||||
key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD)
|
||||
if key is not None:
|
||||
return unquote(str(key))
|
||||
|
|
@ -284,6 +301,12 @@ def _is_vertex_generate_content_batch_output_row(vertex_output_row: Mapping[str,
|
|||
)
|
||||
|
||||
|
||||
def _is_vertex_chat_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool:
|
||||
return _is_vertex_generate_content_batch_output_row(vertex_output_row) or is_vertex_anthropic_batch_output_row(
|
||||
vertex_output_row
|
||||
)
|
||||
|
||||
|
||||
def _try_parse_vertex_batch_output_row(line: bytes) -> _VertexBatchRow | None:
|
||||
try:
|
||||
row: Final = _parse_vertex_batch_output_row(line.decode("utf-8"))
|
||||
|
|
@ -405,6 +428,37 @@ def _openai_batch_output_row(
|
|||
}
|
||||
|
||||
|
||||
def _vertex_anthropic_batch_output_row_to_openai(
|
||||
custom_id: str,
|
||||
response: object,
|
||||
mock_httpx_response: httpx.Response,
|
||||
) -> _OpenAIBatchOutputRow:
|
||||
"""
|
||||
One OpenAI batch output row from a Vertex Claude row's Anthropic `message` response
|
||||
"""
|
||||
anthropic_response: Final = as_json_object(response)
|
||||
if anthropic_response is None or not is_vertex_anthropic_batch_output_response(anthropic_response):
|
||||
return _openai_batch_output_row(
|
||||
custom_id=custom_id,
|
||||
error_code="vertex_ai_error",
|
||||
error_message="Claude batch row has no message response",
|
||||
)
|
||||
try:
|
||||
model_response: Final = VertexAIAnthropicConfig().transform_parsed_response(
|
||||
completion_response=dict(anthropic_response),
|
||||
raw_response=mock_httpx_response,
|
||||
model_response=ModelResponse(),
|
||||
json_mode=True,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # one malformed row must fail alone, not the whole output file
|
||||
return _openai_batch_output_row(
|
||||
custom_id=custom_id,
|
||||
error_code="transformation_error",
|
||||
error_message=f"Failed to transform response: {e}",
|
||||
)
|
||||
return _openai_batch_output_row(custom_id=custom_id, body=model_response.model_dump())
|
||||
|
||||
|
||||
def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]:
|
||||
"""
|
||||
Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch
|
||||
|
|
@ -679,21 +733,123 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows(
|
|||
)
|
||||
|
||||
|
||||
_CLAUDE_BATCH_ROW_FIXED_PARAMS: Final = frozenset(
|
||||
("model", "messages", "custom_llm_provider", "drop_params", "stream", "stream_options")
|
||||
)
|
||||
_CLAUDE_BATCH_OPENAI_PARAMS: Final = frozenset(DEFAULT_CHAT_COMPLETION_PARAM_VALUES) - {"metadata"}
|
||||
_CLAUDE_BATCH_DEPLOYMENT_PARAMS: Final = (
|
||||
_CLAUDE_BATCH_OPENAI_PARAMS - _CLAUDE_BATCH_ROW_FIXED_PARAMS - {"extra_headers"}
|
||||
) | {"drop_params"}
|
||||
|
||||
|
||||
def _claude_batch_deployment_params(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{key: value for key, value in litellm_params.items() if key in _CLAUDE_BATCH_DEPLOYMENT_PARAMS}
|
||||
)
|
||||
|
||||
|
||||
def _names_the_response_format_tool(tool: object) -> bool:
|
||||
match tool:
|
||||
case {"function": {"name": str(name)}}:
|
||||
return name == RESPONSE_FORMAT_TOOL_NAME
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _openai_batch_jsonl_entry_to_vertex_anthropic_row(
|
||||
custom_id: object,
|
||||
messages: list[AllMessageValues],
|
||||
request_params: Mapping[str, object],
|
||||
model: str,
|
||||
deployment_params: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Transforms an OpenAI chat batch entry into the Anthropic Messages row Vertex batches Claude with,
|
||||
`{"custom_id": ..., "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}}`
|
||||
"""
|
||||
if custom_id is None:
|
||||
raise ValueError("custom_id is required on Claude lines, since Claude output rows are matched back only by it")
|
||||
config: Final = VertexAIAnthropicConfig()
|
||||
bare_model: Final = model.rsplit("/", 1)[-1]
|
||||
params: Final = {**deployment_params, **request_params}
|
||||
drop_params: Final = litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True
|
||||
chat_params: Final = {
|
||||
**get_non_default_completion_params(params),
|
||||
**{key: value for key, value in params.items() if key in _CLAUDE_BATCH_OPENAI_PARAMS},
|
||||
}
|
||||
if any(map(_names_the_response_format_tool, iter_items(chat_params.get("tools") or ()))):
|
||||
raise ValueError(
|
||||
f"tool name {RESPONSE_FORMAT_TOOL_NAME!r} is reserved for response_format on Claude, "
|
||||
"and batch output cannot tell a user tool by that name apart, so rename it"
|
||||
)
|
||||
optional_params: Final[Mapping[str, object]] = get_optional_params(
|
||||
model=bare_model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
drop_params=drop_params,
|
||||
messages=messages,
|
||||
**{ # pyright: ignore[reportArgumentType] # batch JSONL bodies are untyped; get_optional_params checks each param as chat does
|
||||
key: value for key, value in chat_params.items() if key not in _CLAUDE_BATCH_ROW_FIXED_PARAMS
|
||||
},
|
||||
)
|
||||
request_state: Final[dict[str, object]] = {"drop_params": drop_params}
|
||||
anthropic_request: Final[Mapping[str, object]] = config.transform_request(
|
||||
model=bare_model,
|
||||
messages=messages,
|
||||
optional_params={
|
||||
**{key: value for key, value in optional_params.items() if key not in ("json_mode", "extra_body")},
|
||||
"anthropic_version": VERTEX_ANTHROPIC_VERSION,
|
||||
},
|
||||
litellm_params=request_state,
|
||||
headers={},
|
||||
)
|
||||
renamed_tools: Final = as_str_mapping(request_state.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY))
|
||||
if renamed_tools:
|
||||
raise ValueError(
|
||||
f"tool names {sorted(map(str, renamed_tools.values()))} must match ^[a-zA-Z0-9_-]{{1,128}}$ for Claude on Vertex, "
|
||||
"and batch output cannot map a renamed tool back, so rename these tools"
|
||||
)
|
||||
extra_body: Final = as_json_object(optional_params.get("extra_body"))
|
||||
request: Final[Mapping[str, object]] = (
|
||||
anthropic_request
|
||||
if extra_body is None
|
||||
else {
|
||||
**anthropic_request,
|
||||
**config.transform_extra_body(
|
||||
extra_body=extra_body, request=anthropic_request, model=bare_model, litellm_params={}
|
||||
),
|
||||
}
|
||||
)
|
||||
return {"custom_id": str(custom_id), "request": request}
|
||||
|
||||
|
||||
def _openai_batch_jsonl_entry_to_vertex_rows(
|
||||
openai_entry: dict[str, Any],
|
||||
map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]],
|
||||
batch_model: str | None = None,
|
||||
deployment_params: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to.
|
||||
|
||||
`batch_model` is the deployment model the batch runs against and wins over the entry's own `body.model`,
|
||||
which picks the row shape: Anthropic Messages for Claude, `GenerateContentRequest` otherwise.
|
||||
|
||||
jsonl body for vertex is {"request": <request_body>}
|
||||
Example Vertex jsonl
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
|
||||
"""
|
||||
entry_body: Final = openai_entry.get("body")
|
||||
model: Final = batch_model or (str(entry_body.get("model") or "") if isinstance(entry_body, dict) else "")
|
||||
is_anthropic_model: Final = is_vertex_anthropic_batch_model(model)
|
||||
if _is_embeddings_batch_entry(openai_entry):
|
||||
if is_anthropic_model:
|
||||
raise VertexAIError(
|
||||
status_code=400,
|
||||
message=f"Vertex AI batch embeddings are not supported for Claude models ('{model}')",
|
||||
)
|
||||
return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry)
|
||||
|
||||
if not isinstance(openai_entry.get("body"), dict):
|
||||
if not isinstance(entry_body, dict):
|
||||
raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object")
|
||||
openai_request_body: Final = openai_entry["body"]
|
||||
chat_request_body: Final = (
|
||||
|
|
@ -701,8 +857,21 @@ def _openai_batch_jsonl_entry_to_vertex_rows(
|
|||
if _is_responses_batch_entry(openai_entry)
|
||||
else openai_request_body
|
||||
)
|
||||
chat_messages: Final = map_developer_role_to_system_role(chat_request_body.get("messages", []))
|
||||
if is_anthropic_model:
|
||||
return (
|
||||
_openai_batch_jsonl_entry_to_vertex_anthropic_row(
|
||||
custom_id=openai_entry.get("custom_id"),
|
||||
messages=chat_messages,
|
||||
request_params={
|
||||
key: value for key, value in chat_request_body.items() if key not in ("model", "messages")
|
||||
},
|
||||
model=model,
|
||||
deployment_params=deployment_params,
|
||||
),
|
||||
)
|
||||
vertex_request_body: Final = transform_request_body(
|
||||
messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])),
|
||||
messages=chat_messages,
|
||||
model=chat_request_body.get("model", ""),
|
||||
optional_params=map_openai_to_vertex_params(chat_request_body),
|
||||
custom_llm_provider="vertex_ai",
|
||||
|
|
@ -854,15 +1023,24 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream):
|
|||
self,
|
||||
openai_file_content: FileTypes,
|
||||
map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]],
|
||||
batch_model: str | None = None,
|
||||
deployment_params: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> None:
|
||||
self._openai_file_content = openai_file_content
|
||||
self._map_openai_to_vertex_params = map_openai_to_vertex_params
|
||||
self._batch_model = batch_model
|
||||
self._deployment_params = deployment_params
|
||||
|
||||
def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]:
|
||||
first = True
|
||||
for lineno, entry in _iter_openai_jsonl_entries(self._openai_file_content):
|
||||
try:
|
||||
vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params)
|
||||
vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(
|
||||
entry,
|
||||
self._map_openai_to_vertex_params,
|
||||
batch_model=self._batch_model,
|
||||
deployment_params=self._deployment_params,
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
status_code=400,
|
||||
|
|
@ -918,16 +1096,31 @@ class _RawFileUploadStream(BaseFileUploadStream):
|
|||
return _iter_raw_file_chunks(self._file_content)
|
||||
|
||||
|
||||
def _deployment_batch_model(litellm_params: Mapping[str, object]) -> str | None:
|
||||
model: Final = litellm_params.get("model")
|
||||
return model.removeprefix("vertex_ai/") if isinstance(model, str) and model else None
|
||||
|
||||
|
||||
def _publisher_model_path(raw_model: str) -> str:
|
||||
if _PUBLISHER_MODEL_PATH_PATTERN.search(raw_model):
|
||||
return raw_model
|
||||
publisher: Final = "anthropic" if is_vertex_anthropic_batch_model(raw_model) else "google"
|
||||
return f"publishers/{publisher}/models/{raw_model}"
|
||||
|
||||
|
||||
def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str:
|
||||
endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
|
||||
model_path: Final = (
|
||||
f"endpoints/{endpoint_id}"
|
||||
if endpoint_id is not None
|
||||
else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
|
||||
)
|
||||
safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
|
||||
model_path: Final = f"endpoints/{endpoint_id}" if endpoint_id is not None else _publisher_model_path(raw_model)
|
||||
prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX
|
||||
return f"{prefix}{safe_model_path}/{uuid.uuid4()}"
|
||||
return f"{prefix}{_safe_managed_model_path(model_path)}/{uuid.uuid4()}"
|
||||
|
||||
|
||||
def _safe_managed_model_path(model_path: str) -> str:
|
||||
if not is_vertex_anthropic_batch_model(model_path):
|
||||
return sanitize_cloud_object_path(model_path, fallback="model")
|
||||
return "@".join(
|
||||
sanitize_cloud_object_path(version_part, fallback="model") for version_part in model_path.split("@")
|
||||
)
|
||||
|
||||
|
||||
class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
||||
|
|
@ -1030,10 +1223,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
fallback_filename="file",
|
||||
)
|
||||
|
||||
def _get_configured_bucket_name(self, litellm_params: dict) -> str:
|
||||
def _get_configured_bucket_name(self, litellm_params: Mapping[str, object]) -> str:
|
||||
configured_bucket_name: Final = litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name")
|
||||
bucket_name: Final = (
|
||||
litellm_params.get("gcs_bucket_name")
|
||||
or litellm_params.get("bucket_name")
|
||||
(configured_bucket_name if isinstance(configured_bucket_name, str) else None)
|
||||
or os.getenv("GCS_BATCH_BUCKET_NAME")
|
||||
or os.getenv("GCS_BUCKET_NAME")
|
||||
)
|
||||
|
|
@ -1047,7 +1240,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
data: CreateFileRequest,
|
||||
) -> str:
|
||||
"""
|
||||
|
|
@ -1071,11 +1264,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
raise ValueError("file is required")
|
||||
if purpose is None:
|
||||
raise ValueError("purpose is required")
|
||||
configured_model: Final = litellm_params.get("model")
|
||||
object_name = self.get_object_name(
|
||||
file_data,
|
||||
purpose,
|
||||
deployment_model=configured_model if isinstance(configured_model, str) else None,
|
||||
deployment_model=_deployment_batch_model(litellm_params),
|
||||
passthrough=is_passthrough_batch_upload(data, litellm_params),
|
||||
)
|
||||
if object_prefix:
|
||||
|
|
@ -1126,7 +1318,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
model: str,
|
||||
create_file_data: CreateFileRequest,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> bytes | str | dict:
|
||||
"""
|
||||
2 Cases:
|
||||
|
|
@ -1157,6 +1349,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
body_stream=_OpenAIToVertexBatchUploadStream(
|
||||
file_data,
|
||||
self._map_openai_to_vertex_params,
|
||||
batch_model=_deployment_batch_model(litellm_params),
|
||||
deployment_params=_claude_batch_deployment_params(litellm_params),
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
|
@ -1207,7 +1401,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException:
|
||||
return VertexAIError(status_code=status_code, message=error_message, headers=headers)
|
||||
|
||||
def _parse_gcs_uri(self, file_id: str, litellm_params: dict | None = None) -> tuple[str, str]:
|
||||
def _parse_gcs_uri(self, file_id: str, litellm_params: Mapping[str, object] | None = None) -> tuple[str, str]:
|
||||
"""
|
||||
Validate a managed GCS file_id and return (bucket, url-encoded-object-path).
|
||||
"""
|
||||
|
|
@ -1225,7 +1419,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
self,
|
||||
file_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[str, dict]:
|
||||
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
|
||||
url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
|
||||
|
|
@ -1255,7 +1449,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
self,
|
||||
file_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[str, dict]:
|
||||
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
|
||||
url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
|
||||
|
|
@ -1298,7 +1492,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
self,
|
||||
file_content_request,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[str, dict]:
|
||||
file_id: Final = file_content_request.get("file_id", "")
|
||||
bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params)
|
||||
|
|
@ -1395,7 +1589,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
headers=MappingProxyType({**headers, "content-length": str(len(transformed_content))}),
|
||||
)
|
||||
|
||||
if not _is_vertex_generate_content_batch_output_row(first_row):
|
||||
if not _is_vertex_chat_batch_output_row(first_row):
|
||||
return FileContentStreamingResult(stream_iterator=replayed_stream, headers=headers)
|
||||
|
||||
return FileContentStreamingResult(
|
||||
|
|
@ -1479,10 +1673,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
# first line is not valid UTF-8/JSON) raises and falls through to the
|
||||
# passthrough below, leaving the content untouched.
|
||||
first_row: Final = _parse_vertex_batch_output_row(first_line)
|
||||
if not (
|
||||
_is_vertex_embeddings_batch_output_row(first_row)
|
||||
or _is_vertex_generate_content_batch_output_row(first_row)
|
||||
):
|
||||
if not (_is_vertex_embeddings_batch_output_row(first_row) or _is_vertex_chat_batch_output_row(first_row)):
|
||||
return content
|
||||
|
||||
context: Final = _new_vertex_batch_output_row_transform_context()
|
||||
|
|
@ -1547,6 +1738,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
error_message=status,
|
||||
)
|
||||
|
||||
if is_vertex_anthropic_batch_output_row(vertex_output):
|
||||
return _vertex_anthropic_batch_output_row_to_openai(
|
||||
custom_id=custom_id,
|
||||
response=vertex_output.get("response"),
|
||||
mock_httpx_response=mock_httpx_response,
|
||||
)
|
||||
|
||||
# Transform successful response using existing transformation
|
||||
vertex_response: Final = vertex_output.get("response", {})
|
||||
|
||||
|
|
|
|||
|
|
@ -2296,6 +2296,36 @@ async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeyp
|
|||
assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_vertex_claude_rows_are_priced_like_with_the_native_output_flag(monkeypatch):
|
||||
claude_message: Final = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "Blue."}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5},
|
||||
}
|
||||
rows: Final = [
|
||||
{"custom_id": "ok", "processed_time": "2026-10-29T03:00:00Z", "response": claude_message, "status": ""},
|
||||
{"custom_id": "bad", "processed_time": "2026-10-29T03:00:00Z", "response": {}, "status": "Bad Request"},
|
||||
]
|
||||
|
||||
async def cost_with_flag(flag: bool) -> bu.BatchCostUsageResult:
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", flag, raising=False)
|
||||
return await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="vertex_ai/claude-sonnet-4-6"
|
||||
)
|
||||
|
||||
unflagged: Final = await cost_with_flag(False)
|
||||
flagged: Final = await cost_with_flag(True)
|
||||
|
||||
assert unflagged.cost > 0
|
||||
assert (unflagged.successful_requests, unflagged.failed_requests) == (1, 1)
|
||||
assert (unflagged.cost, unflagged.usage.total_tokens) == (flagged.cost, flagged.usage.total_tokens)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -41,9 +42,11 @@ from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402
|
|||
VertexAIBatchPrediction,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402
|
||||
from litellm.proxy.utils import handle_exception_on_proxy # noqa: E402
|
||||
from litellm.types.utils import LiteLLMBatch # noqa: E402
|
||||
|
||||
HMOD = "litellm.llms.vertex_ai.batches.handler"
|
||||
_CLIENT_ERROR_429: Final = "Client error '429 Too Many Requests' for url 'https://x/batchPredictionJobs'"
|
||||
TOKEN = "ya29.fake-access-token"
|
||||
PROJECT = "my-project"
|
||||
LOCATION = "us-central1"
|
||||
|
|
@ -370,19 +373,19 @@ def test_create_batch_sync_endpoint_without_deployed_model_raises_400():
|
|||
client.post.assert_not_called()
|
||||
|
||||
|
||||
def test_create_batch_sync_httpstatuserror_propagates():
|
||||
"""``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the
|
||||
sync create path must surface that error, not swallow it."""
|
||||
def test_create_batch_sync_httpstatuserror_carries_vertex_status_body_and_headers():
|
||||
h = _make_handler()
|
||||
client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs")
|
||||
err_response = httpx.Response(status_code=500, request=request, text="boom")
|
||||
err_response: Final = httpx.Response(
|
||||
status_code=429, request=request, text="boom", headers={"retry-after": "30"}
|
||||
)
|
||||
client.post.side_effect = httpx.HTTPStatusError(
|
||||
"boom", request=request, response=err_response
|
||||
_CLIENT_ERROR_429, request=request, response=err_response
|
||||
)
|
||||
|
||||
with patch(f"{HMOD}.get_httpx_client", return_value=client):
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
||||
h.create_batch(
|
||||
_is_async=False,
|
||||
create_batch_data=CREATE_DATA,
|
||||
|
|
@ -393,6 +396,11 @@ def test_create_batch_sync_httpstatuserror_propagates():
|
|||
timeout=600.0,
|
||||
max_retries=None,
|
||||
)
|
||||
assert exc_info.value.response.status_code == 429
|
||||
assert exc_info.value.response.headers["retry-after"] == "30"
|
||||
assert handle_exception_on_proxy(exc_info.value).code == "429"
|
||||
assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom")
|
||||
assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"]
|
||||
|
||||
|
||||
def test_create_batch_input_file_id_without_model_raises_400_before_post():
|
||||
|
|
@ -800,14 +808,16 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised():
|
|||
client.get.assert_not_called()
|
||||
|
||||
|
||||
def test_create_batch_async_httpstatuserror_logged_and_reraised():
|
||||
def test_create_batch_async_httpstatuserror_carries_vertex_status_body_and_headers():
|
||||
h = _make_handler()
|
||||
async_client = MagicMock()
|
||||
request = httpx.Request("POST", "https://x/batchPredictionJobs")
|
||||
err_response = httpx.Response(status_code=500, request=request, text="boom")
|
||||
err_response: Final = httpx.Response(
|
||||
status_code=429, request=request, text="boom", headers={"retry-after": "30"}
|
||||
)
|
||||
async_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"boom", request=request, response=err_response
|
||||
_CLIENT_ERROR_429, request=request, response=err_response
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -825,8 +835,13 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised():
|
|||
timeout=600.0,
|
||||
max_retries=None,
|
||||
)
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
||||
_run(coro)
|
||||
assert exc_info.value.response.status_code == 429
|
||||
assert exc_info.value.response.headers["retry-after"] == "30"
|
||||
assert handle_exception_on_proxy(exc_info.value).code == "429"
|
||||
assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom")
|
||||
assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"]
|
||||
|
||||
|
||||
def test_async_retrieve_batch_non_200_raises():
|
||||
|
|
@ -1043,15 +1058,7 @@ async def test_vertex_list_batches(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_async_create_batch_logs_error_body_on_http_error():
|
||||
"""
|
||||
When Vertex AI returns an HTTP error (e.g. 400), _async_create_batch should
|
||||
re-raise httpx.HTTPStatusError (not swallow it) and log the response body.
|
||||
|
||||
Before the fix the error body was lost because AsyncHTTPHandler.post()
|
||||
calls raise_for_status() internally, raising before the handler's own
|
||||
status-code check could log the body.
|
||||
"""
|
||||
async def test_vertex_async_create_batch_httpstatuserror_carries_vertex_error_body():
|
||||
from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction
|
||||
|
||||
handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
|
||||
|
|
@ -1081,4 +1088,4 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error():
|
|||
)
|
||||
|
||||
assert exc_info.value.response.status_code == 400
|
||||
assert "gemini-2.0-flash" in exc_info.value.response.text
|
||||
assert "Do not support publisher model gemini-2.0-flash" in str(exc_info.value)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ import pytest
|
|||
|
||||
from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402
|
||||
VertexAIBatchTransformation,
|
||||
is_vertex_anthropic_batch_model,
|
||||
native_vertex_batch_row_stats,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402
|
||||
from litellm.types.utils import LiteLLMBatch # noqa: E402
|
||||
|
|
@ -493,3 +495,61 @@ def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_inp
|
|||
T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE)
|
||||
== "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash"
|
||||
)
|
||||
|
||||
|
||||
CLAUDE_INPUT_FILE = (
|
||||
"gs://litellm-testing-bucket/litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/uuid-1"
|
||||
)
|
||||
|
||||
|
||||
def test_claude_batch_job_targets_the_anthropic_publisher_model():
|
||||
job: Final = T.transform_openai_batch_request_to_vertex_ai_batch_request(
|
||||
{"input_file_id": CLAUDE_INPUT_FILE}, vertex_project="p", vertex_location="us-east5"
|
||||
)
|
||||
assert job["model"] == "publishers/anthropic/models/claude-sonnet-4-6"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected"),
|
||||
[
|
||||
("claude-sonnet-4-6", True),
|
||||
("vertex_ai/claude-opus-4-1@20250805", True),
|
||||
("publishers/anthropic/models/claude-haiku-4-5", True),
|
||||
("gemini-2.5-flash", False),
|
||||
("publishers/google/models/gemini-2.5-flash", False),
|
||||
],
|
||||
)
|
||||
def test_is_vertex_anthropic_batch_model(model: str, expected: bool):
|
||||
assert is_vertex_anthropic_batch_model(model) is expected
|
||||
|
||||
|
||||
def test_claude_output_row_is_costed_from_anthropic_usage_and_its_own_model():
|
||||
def cost_calculator(usage, model, custom_llm_provider=None, model_info=None, vertex_location=None) -> tuple[float, float]:
|
||||
if model != "claude-sonnet-4-6":
|
||||
return 0.0, 0.0
|
||||
return usage.prompt_tokens * 0.001, usage.completion_tokens * 0.002
|
||||
|
||||
stats: Final = native_vertex_batch_row_stats(
|
||||
{
|
||||
"custom_id": "r1",
|
||||
"request": {"messages": [{"role": "user", "content": "hi"}], "anthropic_version": "vertex-2023-10-16"},
|
||||
"response": {
|
||||
"id": "msg_vrtx_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5},
|
||||
},
|
||||
},
|
||||
None,
|
||||
model_info=None,
|
||||
calculate_usage=lambda _: pytest.fail("Claude rows must not go through the Gemini usage parser"),
|
||||
cost_calculator=cost_calculator,
|
||||
)
|
||||
|
||||
assert stats is not None, "a successful Claude row must be costed, not counted as failed"
|
||||
assert (stats.usage.prompt_tokens, stats.usage.completion_tokens, stats.total_tokens) == (12, 5, 17)
|
||||
assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((0.012, 0.010))
|
||||
assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((12 * 0.001, 5 * 0.002))
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@ Includes tests for Vertex AI batch output transformation to OpenAI format.
|
|||
|
||||
import json
|
||||
import urllib.parse
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
from types import MappingProxyType
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
|
|
@ -13,6 +14,8 @@ import httpx
|
|||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileUploadStream
|
||||
from litellm.llms.vertex_ai.files.transformation import (
|
||||
VertexAIFilesConfig,
|
||||
_get_litellm_batch_custom_id_from_labels,
|
||||
|
|
@ -165,6 +168,17 @@ class TestBatchObjectNaming:
|
|||
object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}])
|
||||
assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "stored_model"),
|
||||
[
|
||||
("multimodalembedding@001", "publishers/google/models/multimodalembedding_001/"),
|
||||
("claude-haiku-4-5@20251001", "publishers/anthropic/models/claude-haiku-4-5@20251001/"),
|
||||
],
|
||||
)
|
||||
def test_only_claude_models_keep_the_at_version_in_their_object_name(self, config, model, stored_model):
|
||||
object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": model}}])
|
||||
assert object_name.startswith(f"litellm-vertex-files/{stored_model}")
|
||||
|
||||
def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config):
|
||||
"""A numeric endpoint id must not be filed under publishers/google/models/gemini/<id>,
|
||||
which the batch transformation later mangles into a nonexistent publisher model (LIT-6899)."""
|
||||
|
|
@ -1948,3 +1962,420 @@ class TestVertexEmbeddingsBatchOutputTranslation:
|
|||
content = json.dumps(legacy_row).encode("utf-8")
|
||||
|
||||
assert config._try_transform_vertex_batch_output_to_openai(content) == content
|
||||
|
||||
|
||||
CLAUDE_DEPLOYMENT = {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/claude-sonnet-4-6"}
|
||||
|
||||
|
||||
def _claude_chat_entry(custom_id: str | None, content: str, **body: object) -> dict[str, object]:
|
||||
entry: Final[dict[str, object]] = {
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "claude-sonnet-4-6", "messages": [{"role": "user", "content": content}], **body},
|
||||
}
|
||||
return entry if custom_id is None else {"custom_id": custom_id, **entry}
|
||||
|
||||
|
||||
def _uploaded_rows(
|
||||
config: VertexAIFilesConfig, litellm_params: Mapping[str, object], entries: Sequence[Mapping[str, object]]
|
||||
) -> list:
|
||||
upload: Final = config.transform_create_file_request(
|
||||
model="",
|
||||
create_file_data={
|
||||
"file": ("batch.jsonl", "\n".join(json.dumps(entry) for entry in entries).encode(), "application/jsonl"),
|
||||
"purpose": "batch",
|
||||
},
|
||||
optional_params={},
|
||||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
match upload:
|
||||
case {"streaming_media_upload": {"body_stream": BaseFileUploadStream() as stream}}:
|
||||
return [json.loads(line) for line in b"".join(stream.iter_bytes()).splitlines()]
|
||||
case _:
|
||||
pytest.fail(f"expected a streamed upload, got {upload!r}")
|
||||
|
||||
|
||||
def _claude_output_row(custom_id: str, text: str, content: object = None) -> dict[str, object]:
|
||||
return {
|
||||
"custom_id": custom_id,
|
||||
"status": "",
|
||||
"processed_time": "2026-10-28T20:37:47.634+00:00",
|
||||
"response": {
|
||||
"id": "msg_vrtx_01",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": [{"type": "text", "text": text}] if content is None else content,
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestVertexClaudeBatch:
|
||||
"""Claude on Vertex batches only as `publishers/anthropic/models/<model>` with Anthropic Messages rows
|
||||
(https://github.com/BerriAI/litellm/issues/45671)"""
|
||||
|
||||
@pytest.mark.parametrize("claude_model", ["claude-sonnet-4-6", "claude-haiku-4-5@20251001"])
|
||||
def test_claude_upload_creates_a_batch_job_for_the_anthropic_publisher_model(self, config, claude_model):
|
||||
from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation
|
||||
|
||||
url: Final = config.get_complete_file_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="",
|
||||
optional_params={},
|
||||
litellm_params={**CLAUDE_DEPLOYMENT, "model": f"vertex_ai/{claude_model}"},
|
||||
data={"file": ("batch.jsonl", json.dumps(_claude_chat_entry("r1", "hi")).encode()), "purpose": "batch"},
|
||||
)
|
||||
object_name: Final = parse_qs(urlparse(url).query)["name"][0]
|
||||
|
||||
job: Final = VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
|
||||
{"input_file_id": f"gs://my-bucket/{object_name}"}, vertex_project="p", vertex_location="us-east5"
|
||||
)
|
||||
|
||||
assert job["model"] == f"publishers/anthropic/models/{claude_model}", object_name
|
||||
|
||||
def test_claude_body_model_without_deployment_is_filed_under_anthropic_publisher(self, config):
|
||||
object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "claude-haiku-4-5"}}])
|
||||
assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-haiku-4-5/")
|
||||
|
||||
def test_publisher_path_model_is_not_nested_under_google(self, config):
|
||||
object_name: Final = config._get_gcs_object_name_from_batch_jsonl(
|
||||
[{"body": {}}], deployment_model="vertex_ai/publishers/anthropic/models/claude-sonnet-4-6"
|
||||
)
|
||||
assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/")
|
||||
|
||||
def test_claude_rows_are_uploaded_in_anthropic_messages_shape(self, config):
|
||||
rows: Final = _uploaded_rows(
|
||||
config,
|
||||
CLAUDE_DEPLOYMENT,
|
||||
[_claude_chat_entry("request-1", "Hello!", max_tokens=50, temperature=0.2)],
|
||||
)
|
||||
|
||||
assert rows == [
|
||||
{
|
||||
"custom_id": "request-1",
|
||||
"request": {
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}],
|
||||
"max_tokens": 50,
|
||||
"temperature": 0.2,
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
def test_claude_system_message_becomes_the_anthropic_system_field(self, config):
|
||||
entry: Final = _claude_chat_entry(
|
||||
"r1",
|
||||
"unused",
|
||||
max_tokens=10,
|
||||
messages=[{"role": "system", "content": "be brief"}, {"role": "user", "content": "Hello!"}],
|
||||
)
|
||||
|
||||
(row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry])
|
||||
|
||||
assert row["request"]["system"] == [{"type": "text", "text": "be brief"}]
|
||||
assert [message["role"] for message in row["request"]["messages"]] == ["user"]
|
||||
|
||||
def test_claude_row_drops_an_unsupported_param_under_deployment_drop_params(self, config, monkeypatch):
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5)
|
||||
|
||||
(row,) = _uploaded_rows(
|
||||
config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "drop_params": True}, [entry]
|
||||
)
|
||||
|
||||
assert row["request"]["max_tokens"] == 10
|
||||
assert "top_p" not in row["request"]
|
||||
|
||||
def test_claude_row_drops_a_param_the_deployment_lists_in_additional_drop_params(self, config, monkeypatch):
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5)
|
||||
|
||||
(row,) = _uploaded_rows(
|
||||
config,
|
||||
{**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "additional_drop_params": ["top_p"]},
|
||||
[entry],
|
||||
)
|
||||
|
||||
assert row["request"]["max_tokens"] == 10
|
||||
assert "top_p" not in row["request"]
|
||||
|
||||
def test_claude_rows_take_deployment_params_as_defaults_the_row_overrides(self, config):
|
||||
rows: Final = _uploaded_rows(
|
||||
config,
|
||||
{**CLAUDE_DEPLOYMENT, "max_tokens": 64, "temperature": 0.3},
|
||||
[_claude_chat_entry("r1", "hi"), _claude_chat_entry("r2", "hi", max_tokens=10)],
|
||||
)
|
||||
|
||||
assert [(row["request"]["max_tokens"], row["request"]["temperature"]) for row in rows] == [(64, 0.3), (10, 0.3)]
|
||||
|
||||
def test_claude_row_with_an_unsupported_param_is_refused_without_drop_params(self, config, monkeypatch):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5)
|
||||
|
||||
with pytest.raises(VertexAIError, match=r"line 1 .*top_p=0.5.*drop_params") as exc_info:
|
||||
_uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}, [entry])
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_claude_deployment_wins_over_a_gemini_body_model(self, config):
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, model="gemini-2.5-flash")
|
||||
|
||||
(row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry])
|
||||
|
||||
assert row["request"]["anthropic_version"] == "vertex-2023-10-16"
|
||||
assert "contents" not in row["request"]
|
||||
|
||||
def test_gemini_deployment_still_uploads_generate_content_rows(self, config):
|
||||
(row,) = _uploaded_rows(
|
||||
config,
|
||||
{"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash"},
|
||||
[_claude_chat_entry("r1", "hi")],
|
||||
)
|
||||
assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "hi"}]}]
|
||||
assert "custom_id" not in row
|
||||
|
||||
def test_claude_embeddings_batch_line_is_rejected_with_400(self):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
entry: Final = {
|
||||
"custom_id": "e1",
|
||||
"url": "/v1/embeddings",
|
||||
"body": {"model": "claude-sonnet-4-6", "input": "hi"},
|
||||
}
|
||||
with pytest.raises(VertexAIError) as exc_info:
|
||||
_openai_batch_jsonl_entry_to_vertex_rows(entry, VertexAIFilesConfig()._map_openai_to_vertex_params)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_claude_output_rows_become_openai_batch_rows(self, config):
|
||||
content: Final = b"\n".join(
|
||||
json.dumps(row).encode() for row in (_claude_output_row("a", "first"), _claude_output_row("b", "second"))
|
||||
)
|
||||
|
||||
rows: Final = [
|
||||
json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines()
|
||||
]
|
||||
|
||||
assert [row["custom_id"] for row in rows] == ["a", "b"]
|
||||
assert [row["error"] for row in rows] == [None, None]
|
||||
bodies: Final = [row["response"]["body"] for row in rows]
|
||||
assert [body["choices"][0]["message"]["content"] for body in bodies] == ["first", "second"]
|
||||
assert [body["choices"][0]["finish_reason"] for body in bodies] == ["stop", "stop"]
|
||||
assert [(body["usage"]["prompt_tokens"], body["usage"]["completion_tokens"]) for body in bodies] == [
|
||||
(12, 5),
|
||||
(12, 5),
|
||||
]
|
||||
assert {body["model"] for body in bodies} == {"claude-sonnet-4-6"}
|
||||
|
||||
@pytest.mark.parametrize("failed_first", [False, True])
|
||||
def test_failed_claude_output_row_is_an_openai_error_row(self, config, failed_first):
|
||||
failed_row: Final = {
|
||||
"custom_id": "bad",
|
||||
"status": 'Bad Request: {"type": "error", "error": {"type": "invalid_request_error", "message": "too long"}}',
|
||||
"processed_time": "2026-10-28T20:37:48.413+00:00",
|
||||
"response": {},
|
||||
}
|
||||
rows_in: Final = (failed_row, _claude_output_row("ok", "fine"))
|
||||
content: Final = b"\n".join(json.dumps(row).encode() for row in (rows_in if failed_first else rows_in[::-1]))
|
||||
|
||||
rows_out: Final = {
|
||||
row["custom_id"]: row
|
||||
for row in (
|
||||
json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines()
|
||||
)
|
||||
}
|
||||
|
||||
assert rows_out["ok"]["response"]["body"]["choices"][0]["message"]["content"] == "fine"
|
||||
assert rows_out["bad"]["response"] is None
|
||||
assert rows_out["bad"]["error"]["code"] == "vertex_ai_error"
|
||||
assert "too long" in rows_out["bad"]["error"]["message"]
|
||||
|
||||
def test_claude_row_keeps_provider_only_params_and_leaves_out_litellm_ones(self, config):
|
||||
(row,) = _uploaded_rows(
|
||||
config,
|
||||
CLAUDE_DEPLOYMENT,
|
||||
[_claude_chat_entry("r1", "hi", max_tokens=10, top_k=5, drop_params=True, metadata={"team": "a"})],
|
||||
)
|
||||
|
||||
assert row["request"]["top_k"] == 5
|
||||
assert not {"drop_params", "metadata", "extra_body"} & set(row["request"])
|
||||
|
||||
def test_claude_row_top_k_follows_drop_params_on_a_model_without_sampling_params(self, config, monkeypatch):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False})
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
deployment: Final = {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_k=5)
|
||||
|
||||
with pytest.raises(VertexAIError, match=r"line 1 .*top_k") as exc_info:
|
||||
_uploaded_rows(config, deployment, [entry])
|
||||
(row,) = _uploaded_rows(config, {**deployment, "drop_params": True}, [entry])
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "top_k" not in row["request"]
|
||||
|
||||
@pytest.mark.parametrize("param", [{"n": 2}, {"seed": 7}, {"presence_penalty": 0.5}, {"logprobs": True}])
|
||||
def test_claude_row_param_claude_does_not_support_follows_drop_params(self, config, monkeypatch, param):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
(name,) = param
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, **param)
|
||||
|
||||
with pytest.raises(VertexAIError, match=rf"line 1 .*{name}") as exc_info:
|
||||
_uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry])
|
||||
(row,) = _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "drop_params": True}, [entry])
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert name not in row["request"]
|
||||
|
||||
def test_claude_row_refuses_a_tool_name_claude_would_rename(self, config):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
def tool(name: str) -> dict[str, object]:
|
||||
return {"type": "function", "function": {"name": name, "parameters": {"type": "object", "properties": {}}}}
|
||||
|
||||
renamed: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool("get.weather")])
|
||||
kept: Final = _claude_chat_entry("r2", "hi", max_tokens=10, tools=[tool("get_weather")])
|
||||
|
||||
with pytest.raises(VertexAIError, match=r"line 1 .*get\.weather") as exc_info:
|
||||
_uploaded_rows(config, CLAUDE_DEPLOYMENT, [renamed])
|
||||
(row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [kept])
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert [t["name"] for t in row["request"]["tools"]] == ["get_weather"]
|
||||
|
||||
def test_claude_row_refuses_a_user_tool_named_like_the_response_format_tool(self, config):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
tool: Final = {"type": "function", "function": {"name": "json_tool_call", "parameters": {"type": "object"}}}
|
||||
entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool])
|
||||
|
||||
with pytest.raises(VertexAIError, match=r"line 1 .*json_tool_call") as exc_info:
|
||||
_uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry])
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_claude_row_leaves_out_stream(self, config):
|
||||
entry: Final = _claude_chat_entry(
|
||||
"r1", "hi", max_tokens=10, stream=True, stream_options={"include_usage": True}
|
||||
)
|
||||
|
||||
(row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry])
|
||||
|
||||
assert "stream" not in row["request"]
|
||||
assert "stream_options" not in row["request"]
|
||||
|
||||
def test_claude_row_merges_extra_body_into_the_request(self, config):
|
||||
(row,) = _uploaded_rows(
|
||||
config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10, extra_body={"top_k": 3})]
|
||||
)
|
||||
|
||||
assert row["request"]["top_k"] == 3
|
||||
assert "extra_body" not in row["request"]
|
||||
|
||||
def test_claude_response_format_row_round_trips_as_json_content(self, config):
|
||||
schema: Final = {"type": "object", "properties": {"color": {"type": "string"}}, "required": ["color"]}
|
||||
(row,) = _uploaded_rows(
|
||||
config,
|
||||
CLAUDE_DEPLOYMENT,
|
||||
[
|
||||
_claude_chat_entry(
|
||||
"r1",
|
||||
"sky color?",
|
||||
max_tokens=50,
|
||||
response_format={"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}},
|
||||
)
|
||||
],
|
||||
)
|
||||
assert "json_mode" not in row["request"]
|
||||
(json_tool,) = row["request"]["tools"]
|
||||
output_row: Final = _claude_output_row(
|
||||
"r1",
|
||||
"unused",
|
||||
content=[{"type": "tool_use", "id": "toolu_1", "name": json_tool["name"], "input": {"color": "blue"}}],
|
||||
)
|
||||
|
||||
(out,) = (
|
||||
json.loads(line)
|
||||
for line in config._try_transform_vertex_batch_output_to_openai(
|
||||
json.dumps(output_row).encode()
|
||||
).splitlines()
|
||||
)
|
||||
|
||||
message: Final = out["response"]["body"]["choices"][0]["message"]
|
||||
assert json.loads(message["content"]) == {"color": "blue"}
|
||||
assert not message.get("tool_calls")
|
||||
|
||||
def test_claude_output_keeps_a_user_tool_call_as_a_tool_call(self, config):
|
||||
output_row: Final = _claude_output_row(
|
||||
"r1",
|
||||
"unused",
|
||||
content=[{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}],
|
||||
)
|
||||
|
||||
(out,) = (
|
||||
json.loads(line)
|
||||
for line in config._try_transform_vertex_batch_output_to_openai(
|
||||
json.dumps(output_row).encode()
|
||||
).splitlines()
|
||||
)
|
||||
|
||||
(tool_call,) = out["response"]["body"]["choices"][0]["message"]["tool_calls"]
|
||||
assert tool_call["function"]["name"] == "get_weather"
|
||||
assert json.loads(tool_call["function"]["arguments"]) == {"city": "Paris"}
|
||||
|
||||
def test_claude_input_file_read_back_is_returned_unchanged(self, config):
|
||||
(row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10)])
|
||||
content: Final = json.dumps(row).encode()
|
||||
|
||||
assert config._try_transform_vertex_batch_output_to_openai(content) == content
|
||||
|
||||
def test_malformed_claude_message_fails_only_its_own_row(self, config):
|
||||
malformed: Final = _claude_output_row("bad", "unused", content=5)
|
||||
content: Final = json.dumps(_claude_output_row("ok", "fine")).encode() + b"\n" + json.dumps(malformed).encode()
|
||||
|
||||
ok_row, bad_row = (
|
||||
json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines()
|
||||
)
|
||||
|
||||
assert ok_row["response"]["body"]["choices"][0]["message"]["content"] == "fine"
|
||||
assert bad_row["custom_id"] == "bad"
|
||||
assert bad_row["response"] is None
|
||||
assert bad_row["error"]["code"] == "transformation_error"
|
||||
|
||||
def test_claude_line_without_custom_id_is_a_400_at_upload(self, config):
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
with pytest.raises(VertexAIError, match=r"line 2 .*custom_id") as exc_info:
|
||||
_uploaded_rows(
|
||||
config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10), _claude_chat_entry(None, "hi")]
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_claude_output_is_transformed(self, config):
|
||||
async def chunks():
|
||||
yield json.dumps(_claude_output_row("a", "first")).encode() + b"\n"
|
||||
|
||||
result: Final = await config.transform_file_content_stream(
|
||||
stream_iterator=chunks(),
|
||||
headers={},
|
||||
request_url="https://storage.googleapis.com/storage/v1/b/my-bucket/o/x?alt=media",
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={},
|
||||
)
|
||||
body: Final = b"".join([chunk async for chunk in result.stream_iterator])
|
||||
|
||||
(row,) = (json.loads(line) for line in body.splitlines())
|
||||
assert row["custom_id"] == "a"
|
||||
assert row["response"]["body"]["choices"][0]["message"]["content"] == "first"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue