scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(vertex_ai): keep new batch handler code within the mutable-collection budget

Browse source
This commit is contained in:
mubashir1osmani 2026-09-03 19:58:37 -04:00
parent 1d2ed0bdac
commit 14d5ff5c21
1 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/vertex_ai/batches/handler.py
View file

@ -1,5 +1,5 @@
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Sequence
from typing import TYPE_CHECKING, Final, Protocol
import httpx
@ -61,7 +61,7 @@ class _VertexEndpointDeployedModel(TypedDict, total=False):
class _VertexEndpointResponse(TypedDict, total=False):
deployedModels: ReadOnly[list[_VertexEndpointDeployedModel]]
deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
class _VertexEndpointPayloadView(TypedDict):
@ -179,7 +179,7 @@ class VertexAIBatchPrediction(VertexLLM):
def _resolve_fine_tuned_endpoint_model(
self,
vertex_batch_request: VertexAIBatchPredictionJob,
headers: dict[str, str],
headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers
sync_handler: HTTPHandler,
api_base: str | None,
vertex_location: str,
@ -214,7 +214,7 @@ class VertexAIBatchPrediction(VertexLLM):
)
payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
deployed_models: Final = payload_view["payload"].get("deployedModels") or []
deployed_models: Final = payload_view["payload"].get("deployedModels") or ()
deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
if not deployed_model:
raise VertexAIError(
@ -224,7 +224,8 @@ class VertexAIBatchPrediction(VertexLLM):
"resource to run batch predictions against"
),
)
return {**vertex_batch_request, "model": deployed_model}
resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
return resolved_request
async def _async_create_batch(
self,