diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 97e831f5822..fe43019d8ff 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -6,12 +6,14 @@ import random import time import uuid from collections import Counter -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Iterable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable import httpx +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -48,11 +50,86 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch" _MAX_QUEUE_SIZE: Final = 10_000 _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0 -_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_MAPPING_ADAPTER: Final = TypeAdapter(dict[str, object]) +_SEQUENCE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) + + +def _as_optional_mapping(value: object) -> Mapping[str, object] | None: + """``value`` as a string-keyed mapping, or ``None`` when it is not one.""" + try: + return _MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _as_mapping(value: object) -> Mapping[str, object]: + """``value`` as a string-keyed mapping, empty when it is not one.""" + return _as_optional_mapping(value) or _EMPTY_MAPPING + + +def _as_sequence(value: object) -> tuple[object, ...]: + """``value`` as a sequence, empty when it is not one.""" + try: + return _SEQUENCE_ADAPTER.validate_python(value) + except ValidationError: + return () + + +@runtime_checkable +class _ToolCallLike(Protocol): + """A provider tool-call object exposing ``id`` and ``function`` attributes.""" + + id: str | None + function: Function + + +@runtime_checkable +class _HasType(Protocol): + """Any object carrying a ``type`` discriminator.""" + + type: str | None + + +@runtime_checkable +class _HasModel(Protocol): + """Any object carrying the model name of an LLM call.""" + + model: str | None + + +class _ToolCallFunctionPayload(BaseModel): + name: str | None = "" + arguments: str | dict[str, object] | None = "" + + +class _ToolCallPayload(BaseModel): + id: str | None = "" + type: str | None = "function" + function: _ToolCallFunctionPayload | None = None + + +class _ModerationToolCall(BaseModel): + id: str | None = None + + +class _ModerationMessage(BaseModel): + content: str | None = None + tool_calls: tuple[_ModerationToolCall, ...] | None = None + + +class _ModerationChoice(BaseModel): + message: _ModerationMessage | None = None + + +class _ModerationResponse(BaseModel): + """The OpenAI chat-completion shape both Rubrik moderation webhooks return.""" + + choices: tuple[_ModerationChoice, ...] | None = None class _MalformedToolBlockingResponseError(Exception): - """Raised when the response moderation service returns a structurally invalid + """Raised when a moderation service returns a structurally invalid response (e.g. empty ``choices``). Distinct from transient network/HTTP errors so callers can surface a @@ -61,6 +138,14 @@ class _MalformedToolBlockingResponseError(Exception): """ +def _parse_moderation_response(service_response: Mapping[str, object], service_name: str) -> _ModerationResponse: + """Validate a webhook response into the chat-completion shape this module reads.""" + try: + return _ModerationResponse.model_validate(service_response) + except ValidationError as e: + raise _MalformedToolBlockingResponseError(f"{service_name} returned an unreadable response: {e}") from e + + @dataclass class BlockedResponseResult: """Returned by _extract_response_block when the response was blocked @@ -143,7 +228,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): else {"Content-Type": "application/json"} ) - self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task() @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -191,7 +276,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop: Final = asyncio.get_running_loop() @@ -212,7 +297,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Closing them here would close the shared connection pool for every other logger instance; let LiteLLM manage their lifecycle instead. """ - task: Final = getattr(self, "_periodic_flush_task", None) + task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None) if task is not None: task.cancel() @@ -253,7 +338,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod async def _guarded( - coro: Any, + coro: Awaitable[GenericGuardrailAPIInputs], inputs: GenericGuardrailAPIInputs, label: str, ) -> GenericGuardrailAPIInputs: @@ -400,34 +485,36 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request_data["_rubrik_logging_obj"] = logging_obj @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + def _normalize_tool_calls(tool_calls: Iterable[object]) -> tuple[ChatCompletionMessageToolCall, ...]: """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) @staticmethod - def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + def _normalize_tool_call(tc: object) -> ChatCompletionMessageToolCall: if isinstance(tc, ChatCompletionMessageToolCall): return tc - if isinstance(tc, dict): - func: Final = tc.get("function") or _EMPTY_MAPPING + as_mapping: Final = _as_optional_mapping(tc) + if as_mapping is not None: + payload: Final = _ToolCallPayload.model_validate(as_mapping) + func: Final = payload.function or _ToolCallFunctionPayload() return ChatCompletionMessageToolCall( - id=tc.get("id", ""), - type=tc.get("type", "function"), + id=payload.id, + type=payload.type, function=Function( - name=func.get("name", ""), - arguments=func.get("arguments", ""), + name=func.name, + arguments=func.arguments, ), ) - if hasattr(tc, "id") and hasattr(tc, "function"): + if isinstance(tc, _ToolCallLike): return ChatCompletionMessageToolCall( id=tc.id or "", - type=getattr(tc, "type", None) or "function", + type=(tc.type if isinstance(tc, _HasType) else None) or "function", function=tc.function, ) raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") @staticmethod - def _join_texts(texts: Any) -> str: + def _join_texts(texts: Sequence[str] | None) -> str: """Join response text segments into the single content string the webhook evaluates. Empty when there is no assistant text.""" if not texts: @@ -439,14 +526,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): tool_calls: Sequence[ChatCompletionMessageToolCall], content: str, request_id: str | None, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """Build an OpenAI ChatCompletion-format dict (assistant text + tool calls) for the after_completion webhook. ``content`` is sent so the webhook can moderate the response text; ``None`` when the assistant produced no text (tool-call-only response). """ - message: Final[dict[str, Any]] = { + message: Final[dict[str, object]] = { "role": "assistant", "content": content or None, } @@ -467,7 +554,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, object], ...]: """Collapse each message's content to a plain string for the webhook. litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, @@ -483,31 +570,30 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "role": message.get("role"), "content": "\n".join(p for p in RubrikLogger._moderation_text_parts(message) if p), } - for message in messages or () - if isinstance(message, dict) + for message in (_as_optional_mapping(entry) for entry in messages or ()) + if message is not None ) @staticmethod - def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + def _moderation_text_parts(message: Mapping[str, object]) -> tuple[str, ...]: """Every attacker-controlled text segment of a message: its content plus the arguments of any tool call or deprecated function call.""" - fc: Final = message.get("function_call") + fc: Final = _as_mapping(message.get("function_call")) return ( # Base text content (flattens Anthropic content-part arrays) - convert_content_list_to_str(message), # pyright: ignore[reportArgumentType] # dict[str,Any] is AllMessageValues at runtime + convert_content_list_to_str(message), # pyright: ignore[reportArgumentType] # a str-keyed mapping is AllMessageValues at runtime *( - str((tc.get("function") or _EMPTY_MAPPING).get("arguments") or "") - for tc in message.get("tool_calls") or () - if isinstance(tc, dict) + str(_as_mapping(_as_mapping(tc).get("function")).get("arguments") or "") + for tc in _as_sequence(message.get("tool_calls")) ), - str((fc.get("arguments") if isinstance(fc, dict) else None) or ""), + str(fc.get("arguments") or ""), ) @staticmethod def _build_prompt_moderation_payload( inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + request_data: Mapping[str, object], + ) -> Mapping[str, object]: """Build the bare OpenAI request the before_prompt webhook consumes. Unlike the after_completion envelope, this endpoint takes a raw OpenAI @@ -516,7 +602,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ``/v1/messages`` requests too. Optional fields are sent only when present so the payload stays clean. """ - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "model": inputs.get("model") or request_data.get("model") or "", "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), } @@ -539,9 +625,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _extract_request_data( - call_details: Mapping[str, Any], - request_data: Mapping[str, Any] | None, - ) -> Mapping[str, Any]: + call_details: Mapping[str, object], + request_data: Mapping[str, object] | None, + ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -553,7 +639,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return _EMPTY_MAPPING call_details = call_details or _EMPTY_MAPPING request_data = request_data or _EMPTY_MAPPING - optional_params: Final = call_details.get("optional_params") or _EMPTY_MAPPING + optional_params: Final = _as_mapping(call_details.get("optional_params")) # Use ``in`` rather than truthy ``or`` so an explicit empty list # (caller declared the agent has NO tools) is forwarded as-is. @@ -576,27 +662,32 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + def _sanitize_proxy_server_request(proxy_server_request: object) -> object: """Allowlist only routing fields (``url``, ``method``) when forwarding ``proxy_server_request`` to an external webhook, dropping inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw request ``body`` so proxy credentials are not exfiltrated.""" - if not isinstance(proxy_server_request, dict): + parsed: Final = _as_optional_mapping(proxy_server_request) + if parsed is None: return proxy_server_request - return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} + return {key: parsed[key] for key in ("url", "method") if key in parsed} @staticmethod - def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: + def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, object]) -> str: """Get the model name for the ModifyResponseException.""" response: Final = request_data.get("response") - if response and hasattr(response, "model"): + if response and isinstance(response, _HasModel): return response.model or "unknown" - return call_details.get("model", "unknown") + model: Final = call_details.get("model", "unknown") + return model if isinstance(model, str) else str(model) # -- Logging hooks --------------------------------------------------------- @staticmethod - def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + def _correlation_id( + call_details: Mapping[str, object], + request_data: Mapping[str, object] | None = None, + ) -> str | None: """The id that joins a blocked request's two S3 logs by filename: the moderation (``_blocking``) log and the failure (response) log. @@ -607,10 +698,13 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): block fires before the response/logging object is populated, so the two logs correlate for every provider (OpenAI and Anthropic alike). """ - return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") + correlated: Final = call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get( + "litellm_call_id" + ) + return correlated if isinstance(correlated, str) else None @classmethod - def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, object]) -> None: """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log shares its S3 filename id with the moderation (``_blocking``) and failure logs for the same request -- for every provider. @@ -630,7 +724,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): payload["id"] = correlated @staticmethod - def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None: """Prepend ``source["system"]`` onto ``payload["messages"]``. Builds a NEW messages list rather than mutating ``payload["messages"]`` @@ -658,7 +752,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + async def _prepare_log_payload( + self, + kwargs: Mapping[str, object], + event_type: str, + ) -> StandardLoggingPayload | None: """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate) @@ -697,7 +795,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now - async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str): try: payload: Final = await self._prepare_log_payload(kwargs, event_type) if payload is None: @@ -857,7 +955,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): - time fields: ``call_details["start_time"]`` reused for all three; end/completion times are meaningless for a prompt block. """ - call_details: Final = logging_obj.model_call_details + call_details: Final = _as_mapping(logging_obj.model_call_details) exception_text: Final = f"{type(exception).__name__}: {exception.message}" base: Final = call_details.get("standard_logging_object") @@ -906,13 +1004,13 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @classmethod def _build_fallback_payload( cls, - call_details: Mapping[str, Any], + call_details: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", - ) -> dict[str, Any]: + ) -> dict[str, object]: # Convert datetime to a Unix float so json.dumps can serialize it. # httpx's json= parameter uses stdlib json.dumps with no custom encoder. _raw_start: Final = call_details.get("start_time") - _start: Final = _raw_start.timestamp() if _raw_start is not None else None + _start: Final = _raw_start.timestamp() if isinstance(_raw_start, datetime) else None return { "id": call_details.get("litellm_call_id"), "model": call_details.get("model") or "", @@ -996,7 +1094,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Webhook services ------------------------------------------------------ - async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, object]: """POST ``payload`` to a Rubrik webhook and return its dict response. Raises: @@ -1010,20 +1108,21 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): headers=self._headers, ) http_response.raise_for_status() - result: Final = http_response.json() - if not isinstance(result, dict): + result: Final[object] = http_response.json() # any-ok: httpx exposes the decoded JSON body as Any + parsed: Final = _as_optional_mapping(result) + if parsed is None: raise TypeError( f"{service_name} returned non-dict JSON " f"({type(result).__name__}); expected OpenAI chat completion " "shape or empty object." ) - return result + return parsed async def _post_to_response_moderation_endpoint( self, - response_data: Mapping[str, Any], - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + response_data: Mapping[str, object], + request_data: Mapping[str, object], + ) -> Mapping[str, object]: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1039,7 +1138,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "Response moderation service", ) - async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, object]: """Post a bare OpenAI request to the before_prompt webhook. Returns ``{}`` (passthrough) or a synthetic chat.completion (block). @@ -1047,23 +1146,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + def _extract_prompt_refusal(service_response: Mapping[str, object]) -> str | None: """Return the refusal text when the prompt was blocked, else None. The before_prompt webhook returns ``{}`` (passthrough) or a synthetic chat.completion whose ``choices[0].message.content`` is the refusal explanation. """ - choices: Final = service_response.get("choices") + choices: Final = _parse_moderation_response(service_response, "Prompt moderation service").choices or () if not choices: return None - message: Final = choices[0].get("message") or _EMPTY_MAPPING - content: Final = message.get("content") - return content or "Request blocked by policy." + message: Final = choices[0].message + return (message.content if message else None) or "Request blocked by policy." @staticmethod def _extract_response_block( - service_response: Mapping[str, Any], + service_response: Mapping[str, object], all_tool_calls: Sequence[ChatCompletionMessageToolCall], sent_content: str, ) -> BlockedResponseResult | None: @@ -1086,19 +1184,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices: Final = service_response.get("choices") or () + choices: Final = _parse_moderation_response(service_response, "Response moderation service").choices or () if not choices: raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") - message: Final = choices[0].get("message") or _EMPTY_MAPPING - returned_tool_calls: Final = message.get("tool_calls") or () - returned_content: Final = message.get("content") or "" + message: Final = choices[0].message + returned_tool_calls: Final = (message.tool_calls if message else None) or () + returned_content: Final = (message.content if message else None) or "" # Use Counter so duplicate IDs are handled correctly: if the model # emits two calls with the same ID (one allowed, one prohibited) and # the service returns only the allowed one, a set-based check would # miss the block. Counter preserves multiplicity. - returned_id_counts: Final[Counter[str]] = Counter(tc["id"] for tc in returned_tool_calls if tc.get("id")) + returned_id_counts: Final[Counter[str]] = Counter(tc.id for tc in returned_tool_calls if tc.id) required_id_counts: Final[Counter[str]] = Counter(tc.id for tc in all_tool_calls if tc.id) # Cardinality check catches ID-less tool calls (not counted in # required_id_counts because tc.id is falsy); Counter check catches diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 2e0ae30a192..a055fec84db 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -2,7 +2,8 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final import httpx -from httpx._types import RequestFiles +from httpx._types import FileContent, RequestFiles +from pydantic import BaseModel, ConfigDict, Field import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION @@ -31,6 +32,29 @@ else: LiteLLMLoggingObj = Any +class RunwayMLTaskResponse(BaseModel): + """RunwayML task object returned by ``/v1/image_to_video`` and ``/v1/tasks/{id}``.""" + + model_config = ConfigDict(populate_by_name=True) + + id: str = "" + status: str = "pending" + created_at: str | None = Field(default=None, alias="createdAt") + completed_at: str | None = Field(default=None, alias="completedAt") + output: tuple[str, ...] | str | None = None + progress: int | None = None + failure: str | None = None + failure_code: str | None = Field(default=None, alias="failureCode") + + +class RunwayMLTaskRequest(BaseModel): + """The RunwayML-shaped request fields this module reads back when building a ``VideoObject``.""" + + model: str | None = None + ratio: str | None = None + duration: int | None = None + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. @@ -78,7 +102,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, Any]] = {} + mapped_params: Final[dict[str, object]] = {} # Handle input_reference parameter - map to promptImage if "input_reference" in video_create_optional_params: @@ -180,7 +204,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): } """ # Build the request data - request_data: Final[dict[str, Any]] = { + request_data: Final[dict[str, object]] = { "model": model, "promptText": prompt, } @@ -189,7 +213,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): request_data.update(video_create_optional_request_params) # RunwayML uses JSON body, no files multipart - files_list: Final[list[tuple[str, Any]]] = [] + files_list: Final[RequestFiles] = [] # Append the specific endpoint for video generation full_api_base: Final = f"{api_base}/image_to_video" @@ -216,60 +240,58 @@ class RunwayMLVideoConfig(BaseVideoConfig): We map this to OpenAI VideoObject format. """ - response_data: Final = raw_response.json() + task: Final = RunwayMLTaskResponse.model_validate(raw_response.json()) + request_params: Final = RunwayMLTaskRequest.model_validate(request_data or {}) + seconds: Final = str(request_params.duration) if request_params.duration is not None else None - # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { - "id": response_data.get("id", ""), - "object": "video", - "status": self._map_runway_status(response_data.get("status", "pending")), - "created_at": self._parse_runway_timestamp(response_data.get("createdAt")), + return VideoObject( + id=self._encoded_video_id(task.id, custom_llm_provider, model), + object="video", + status=self._map_runway_status(task.status), + created_at=self._parse_runway_timestamp(task.created_at), + completed_at=self._completed_at(task), + error=self._build_error(task), + model=request_params.model, + size=self._ratio_to_size(request_params.ratio), + seconds=seconds, + usage=self._usage_from_seconds(seconds), + ) + + @staticmethod + def _encoded_video_id(video_id: str, custom_llm_provider: str | None, model: str | None) -> str: + if not (custom_llm_provider and video_id): + return video_id + return encode_video_id_with_provider(video_id, custom_llm_provider, model) + + @staticmethod + def _ratio_to_size(ratio: str | None) -> str | None: + if not ratio or ":" not in ratio: + return None + return ratio.replace(":", "x") + + @staticmethod + def _usage_from_seconds(seconds: str | None) -> dict[str, float]: + if not seconds: + return {} + try: + return {"duration_seconds": float(seconds)} + except ValueError: + return {} + + def _completed_at(self, task: RunwayMLTaskResponse) -> int | None: + if "completed_at" not in task.model_fields_set: + return None + return self._parse_runway_timestamp(task.completed_at) + + @staticmethod + def _build_error(task: RunwayMLTaskResponse) -> dict[str, str] | None: + if not {"failure", "failure_code"} & task.model_fields_set: + return None + return { + "code": task.failure_code if task.failure_code is not None else "unknown", + "message": task.failure if task.failure is not None else "Video generation failed", } - # Add optional fields if present - if "output" in response_data and response_data["output"]: - # RunwayML returns output as array of URLs when task succeeds - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) - - if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) - - if "failureCode" in response_data or "failure" in response_data: - video_data["error"] = { - "code": response_data.get("failureCode", "unknown"), - "message": response_data.get("failure", "Video generation failed"), - } - - # Add model and size info if available from request - if request_data: - if "model" in request_data: - video_data["model"] = request_data["model"] - if "ratio" in request_data: - # Convert ratio back to size format - ratio: Final = request_data["ratio"] - if isinstance(ratio, str) and ":" in ratio: - video_data["size"] = ratio.replace(":", "x") - if "duration" in request_data: - video_data["seconds"] = str(request_data["duration"]) - - video_obj: Final = VideoObject(**video_data) - - if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) - - # Add usage data for cost tracking - usage_data: Final = {} - if video_obj and hasattr(video_obj, "seconds") and video_obj.seconds: - try: - usage_data["duration_seconds"] = float(video_obj.seconds) - except (ValueError, TypeError): - pass - video_obj.usage = usage_data - - return video_obj - def _map_runway_status(self, runway_status: str) -> str: """ Map RunwayML status to OpenAI status format. @@ -326,33 +348,32 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, Any]] = {} + params: Final[dict[str, str]] = {} return url, params - def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str: + def _extract_video_url_from_response(self, task: RunwayMLTaskResponse) -> str: """ Helper method to extract video URL from RunwayML response. Shared between sync and async transforms. """ - # Extract video URL from the output field - video_url = None - if "output" in response_data and response_data["output"]: - output: Final = response_data["output"] - video_url = output[0] if isinstance(output, list) else output + video_url: Final = self._first_output_url(task.output) + if video_url: + return video_url - if not video_url: - # Check if the video generation failed or is still processing - status: Final = response_data.get("status", "UNKNOWN") - if status in ["PENDING", "RUNNING", "THROTTLED"]: - raise ValueError(f"Video is still processing (status: {status}). Please wait and try again.") - elif status == "FAILED": - failure_reason: Final = response_data.get("failure", "Unknown error") - raise ValueError(f"Video generation failed: {failure_reason}") - else: - raise ValueError("Video URL not found in response. Video may not be ready yet.") + status: Final = task.status + if status in ("PENDING", "RUNNING", "THROTTLED"): + raise ValueError(f"Video is still processing (status: {status}). Please wait and try again.") + if status == "FAILED": + failure_reason: Final = task.failure if task.failure is not None else "Unknown error" + raise ValueError(f"Video generation failed: {failure_reason}") + raise ValueError("Video URL not found in response. Video may not be ready yet.") - return video_url + @staticmethod + def _first_output_url(output: tuple[str, ...] | str | None) -> str | None: + if isinstance(output, tuple): + return output[0] if output else None + return output def transform_video_content_response( self, @@ -373,8 +394,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() - video_url: Final = self._extract_video_url_from_response(response_data) + task: Final = RunwayMLTaskResponse.model_validate(raw_response.json()) + video_url: Final = self._extract_video_url_from_response(task) # Download the video from the CloudFront URL synchronously httpx_client: Final[HTTPHandler] = _get_httpx_client() @@ -402,8 +423,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() - video_url: Final = self._extract_video_url_from_response(response_data) + task: Final = RunwayMLTaskResponse.model_validate(raw_response.json()) + video_url: Final = self._extract_video_url_from_response(task) # Download the video from the CloudFront URL asynchronously async_httpx_client: Final[AsyncHTTPHandler] = get_async_httpx_client( @@ -421,7 +442,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video remix request for RunwayML API. @@ -448,7 +469,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: dict[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video list request for RunwayML API. @@ -484,7 +505,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Construct the URL for task cancellation url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -494,17 +515,15 @@ class RunwayMLVideoConfig(BaseVideoConfig): logging_obj: LiteLLMLoggingObj, ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" - response_data: Final = raw_response.json() + task: Final = RunwayMLTaskResponse.model_validate(raw_response.json()) - video_obj: Final = VideoObject( - id=response_data.get("id", ""), + return VideoObject( + id=task.id, object="video", status="cancelled", - created_at=self._parse_runway_timestamp(response_data.get("createdAt")), + created_at=self._parse_runway_timestamp(task.created_at), ) - return video_obj - def transform_video_status_retrieve_request( self, video_id: str, @@ -524,7 +543,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): url: Final = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) - data: Final[dict[str, Any]] = {} + data: Final[dict[str, str]] = {} return url, data @@ -537,42 +556,19 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the RunwayML video status retrieve response. """ - response_data: Final = raw_response.json() + task: Final = RunwayMLTaskResponse.model_validate(raw_response.json()) - # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { - "id": response_data.get("id", ""), - "object": "video", - "status": self._map_runway_status(response_data.get("status", "pending")), - "created_at": self._parse_runway_timestamp(response_data.get("createdAt")), - } + return VideoObject( + id=self._encoded_video_id(task.id, custom_llm_provider, None), + object="video", + status=self._map_runway_status(task.status), + created_at=self._parse_runway_timestamp(task.created_at), + completed_at=self._completed_at(task), + progress=task.progress, + error=self._build_error(task), + ) - # Add optional fields if present - if "output" in response_data and response_data["output"]: - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) - - if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) - - if "progress" in response_data: - video_data["progress"] = response_data["progress"] - - if "failureCode" in response_data or "failure" in response_data: - video_data["error"] = { - "code": response_data.get("failureCode", "unknown"), - "message": response_data.get("failure", "Video generation failed"), - } - - video_obj: Final = VideoObject(**video_data) - - if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) - - return video_obj - - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: FileContent, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for RunwayML") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3538fc5b1a7..fbcd1de2cfd 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,12 +5,13 @@ import json import os import re import time -from collections.abc import Callable, Iterable, Iterator -from typing import Any, Final +from collections.abc import Callable, Iterable, Iterator, Mapping +from typing import Final, Protocol import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from pydantic import TypeAdapter import litellm from litellm._uuid import uuid @@ -46,13 +47,14 @@ from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContentRequest, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError @@ -61,6 +63,53 @@ from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_PURPOSE_ADAPTER: Final = TypeAdapter[OpenAIFilesPurpose](OpenAIFilesPurpose) +_JSON_DECODER: Final = json.JSONDecoder() + + +class _JsonDecoder(Protocol): + def decode(self, s: str, /) -> object: ... + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_object(value: object) -> dict[str, object]: + return _JSON_OBJECT_ADAPTER.validate_python(value) + + +def _json_object_or_empty(value: object) -> dict[str, object]: + return _json_object(value) if value else {} + + +def _parse_json_object(raw: str, decoder: _JsonDecoder = _JSON_DECODER) -> dict[str, object]: + return _json_object(decoder.decode(raw)) + + +def _response_json_object(response: _JsonResponse) -> dict[str, object]: + return _json_object(response.json()) + + +def _str_field(payload: Mapping[str, object], key: str, default: str = "") -> str: + value: Final = payload.get(key, default) + return value if isinstance(value, str) else default + + +def _int_field(payload: Mapping[str, object], key: str) -> int: + value: Final = payload.get(key, 0) + return int(value) if isinstance(value, (int, float, str)) else 0 + + +def _purpose_field(payload: Mapping[str, object], default: OpenAIFilesPurpose = "batch") -> OpenAIFilesPurpose: + return _PURPOSE_ADAPTER.validate_python(payload.get("purpose", default)) + + +def _gcs_file_id(payload: Mapping[str, object]) -> str: + gcs_id: Final = _str_field(payload, "id") + return "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" + def _sanitize_gcp_label_value(value: str) -> str: """ @@ -106,7 +155,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None: return None -def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None: +def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: object) -> None: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -122,7 +171,7 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: @@ -140,10 +189,15 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) +# any-ok: batch bodies must reach _transform_request_body unmodified, so no lossy revalidation here +def _batch_messages(openai_request_body) -> list[AllMessageValues]: + return openai_request_body.get("messages", []) + + def _openai_batch_jsonl_entry_to_vertex_wrapped_request( - openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: + openai_entry: Mapping[str, object], + map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], +) -> dict[str, object]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -151,10 +205,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( 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"}}]}]}} """ - openai_request_body: Final = openai_entry.get("body") or {} + openai_request_body: Final = _json_object_or_empty(openai_entry.get("body")) vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), + messages=_batch_messages(openai_request_body), + model=_str_field(openai_request_body, "model"), optional_params=map_openai_to_vertex_params(openai_request_body), custom_llm_provider="vertex_ai", litellm_params={}, @@ -186,9 +240,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited JSONL. """ - content: Any = openai_file_content - if isinstance(content, tuple): - content = content[1] + content: Final[object] = openai_file_content[1] if isinstance(openai_file_content, tuple) else openai_file_content if isinstance(content, (bytes, bytearray)): # Scan for newlines in place so a large in-memory payload is not copied @@ -220,14 +272,13 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: # object name, then the body stream), so it must rewind to 0. A # non-seekable handle would silently resume mid-stream and drop the # already-consumed first row, so reject it loudly instead. - seek: Final = getattr(content, "seek", None) - if seek is None: + if not hasattr(content, "seek"): raise ValueError( "Batch upload file handle must be seekable; got a non-seekable " "stream. Pass bytes, a path, or a seekable handle." ) try: - seek(0) + content.seek(0) except (OSError, ValueError) as e: raise ValueError( "Batch upload file handle must be seekable so it can be re-read " @@ -241,9 +292,9 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, Any]]: +) -> Iterator[dict[str, object]]: for line in _iter_openai_jsonl_lines(openai_file_content): - yield json.loads(line) + yield _parse_json_object(line) class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): @@ -257,7 +308,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -308,17 +359,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: list[dict[str, object]], ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + model_name: Final = _str_field(_json_object_or_empty(openai_jsonl_content[0].get("body")), "model") + qualified_model: Final = ( + model_name if "publishers/google/models" in model_name else f"publishers/google/models/{model_name}" + ) + safe_model_path: Final = sanitize_cloud_object_path(qualified_model, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name @@ -343,9 +395,11 @@ 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: bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + _str_field(litellm_params, "gcs_bucket_name") + or _str_field(litellm_params, "bucket_name") + or os.getenv("GCS_BUCKET_NAME") ) if not bucket_name: raise ValueError("GCS bucket_name is required") @@ -396,8 +450,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: dict[str, object], + ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ @@ -406,9 +460,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) config: Final = VertexGeminiConfig() - _model: Final = openai_request_body.get("model", "") - vertex_params: Final = config.map_openai_params( - model=_model, + vertex_params: Final[dict[str, object]] = config.map_openai_params( + model=_str_field(openai_request_body, "model"), non_default_params=openai_request_body, optional_params={}, drop_params=False, @@ -463,10 +516,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() - try: - response_object: Final = GcsBucketResponse(**response_json) + response_object: Final = _response_json_object(raw_response) except Exception as e: raise VertexAIError( status_code=raw_response.status_code, @@ -474,19 +525,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers=raw_response.headers, ) - gcs_id = response_object.get("id", "") - # Remove the last numeric ID from the path - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" - return OpenAIFileObject( - purpose=response_object.get("purpose", "batch"), - id=f"gs://{gcs_id}", - filename=response_object.get("name", ""), + purpose=_purpose_field(response_object), + id=f"gs://{_gcs_file_id(response_object)}", + filename=_str_field(response_object, "name"), created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=response_object.get("timeCreated", "") + vertex_datetime=_str_field(response_object, "timeCreated") ), status="uploaded", - bytes=int(response_object.get("size", 0)), + bytes=_int_field(response_object, "size"), object="file", ) @@ -523,18 +570,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() - gcs_id = response_json.get("id", "") - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" + response_json: Final = _response_json_object(raw_response) return OpenAIFileObject( - id=f"gs://{gcs_id}", - bytes=int(response_json.get("size", 0)), + id=f"gs://{_gcs_file_id(response_json)}", + bytes=_int_field(response_json, "size"), created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=response_json.get("timeCreated", "") + vertex_datetime=_str_field(response_json, "timeCreated") ), - filename=response_json.get("name", ""), + filename=_str_field(response_json, "name"), object="file", - purpose=response_json.get("metadata", {}).get("purpose", "batch"), + purpose=_purpose_field(_json_object_or_empty(response_json.get("metadata"))), status="processed", status_details=None, ) @@ -584,7 +629,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def transform_file_content_request( self, - file_content_request, + file_content_request: FileContentRequest, optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: @@ -682,14 +727,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # discriminating fields. Anything else (e.g. a binary file whose # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. - first_row: Final = json.loads(first_line) + first_row: Final = _parse_json_object(first_line) + first_row_response: Final = _json_object_or_empty(first_row.get("response")) is_vertex_batch_output: Final = ( "request" in first_row and "response" in first_row and "processed_time" in first_row and ( - "candidates" in first_row.get("response", {}) - or "promptFeedback" in first_row.get("response", {}) + "candidates" in first_row_response + or "promptFeedback" in first_row_response or bool(first_row.get("status")) ) ) @@ -723,7 +769,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): for line in itertools.chain([first_line], lines): try: openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), + vertex_output=_parse_json_object(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, @@ -742,22 +788,22 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: Mapping[str, object], vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data: Final = vertex_output.get("request", {}) - labels: Final = request_data.get("labels", {}) or {} + request_data: Final = _json_object_or_empty(vertex_output.get("request")) + labels: Final = _json_object_or_empty(request_data.get("labels")) custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels) # Check if there's an error - status: Final = vertex_output.get("status", "") + status: Final = _str_field(vertex_output, "status") has_error: Final = bool(status) if has_error: @@ -772,12 +818,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): } # Transform successful response using existing transformation - vertex_response: Final = vertex_output.get("response", {}) + vertex_response: Final = _json_object_or_empty(vertex_output.get("response")) # Extract model from response - model = vertex_response.get("modelVersion", "gemini-1.5-flash-001") - if "@" in model: - model = model.split("@")[0] + model_version: Final = _str_field(vertex_response, "modelVersion", "gemini-1.5-flash-001") + model: Final = model_version.split("@")[0] if "@" in model_version else model_version try: # Use existing VertexGeminiConfig transformation @@ -792,7 +837,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) # Convert ModelResponse to dict - response_dict: Final = transformed_response.model_dump() + response_dict: Final = _json_object(transformed_response.model_dump()) # Return in OpenAI batch format return { @@ -800,7 +845,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "custom_id": custom_id, "response": { "status_code": 200, - "request_id": response_dict.get("id", ""), + "request_id": _str_field(response_dict, "id"), "body": response_dict, }, "error": None, diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..9d0cda5adf4 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -11,36 +11,209 @@ Flow: [file_search_call output item] + [message output item with file_citation annotations] """ -import json -import time import uuid -from collections.abc import Iterable -from typing import Any, Final, cast +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Literal, Protocol, TypeAlias, TypedDict, TypeVar + +from pydantic import BaseModel, ConfigDict, ValidationError from litellm._internal_context import is_internal_call from litellm._logging import verbose_logger -from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse -from litellm.types.vector_stores import VectorStoreSearchResult - -# Keep ToolParam broad so we stay compatible with both dict and Pydantic forms -ToolParam = Any +from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.types.llms.openai import ( + ResponseInputParam, + ResponseOutputItem, + ResponsesAPIResponse, + ToolParam, +) +from litellm.types.vector_stores import ( + VectorStoreResultContent, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" +_ModelT = TypeVar("_ModelT", bound=BaseModel) + + +class SearchResultContentLike(Protocol): + """Attribute-based counterpart of ``VectorStoreResultContent``.""" + + @property + def text(self) -> str | None: ... + + +class SearchResultLike(Protocol): + """Attribute-based counterpart of ``VectorStoreSearchResult``.""" + + @property + def score(self) -> float | None: ... + + @property + def file_id(self) -> str | None: ... + + @property + def filename(self) -> str | None: ... + + @property + def content(self) -> Sequence[VectorStoreResultContent | SearchResultContentLike] | None: ... + + @property + def attributes(self) -> Mapping[str, object] | None: ... + + +class SearchResponseLike(Protocol): + """Attribute-based counterpart of ``VectorStoreSearchResponse``.""" + + @property + def data(self) -> Sequence["SearchResult"] | None: ... + + +SearchResult: TypeAlias = VectorStoreSearchResult | SearchResultLike + + +class FunctionToolQueriesSchema(TypedDict): + type: Literal["array"] + items: dict[str, str] + description: str + + +class FunctionToolVectorStoreIdSchema(TypedDict): + type: Literal["string"] + description: str + enum: list[str] + + +class FunctionToolProperties(TypedDict): + queries: FunctionToolQueriesSchema + vector_store_id: FunctionToolVectorStoreIdSchema + + +class FunctionToolParameters(TypedDict): + type: Literal["object"] + properties: FunctionToolProperties + required: list[str] + + +class EmulatedFileSearchTool(TypedDict): + type: Literal["function"] + name: str + description: str + parameters: FunctionToolParameters + + +class FileCitationAnnotation(TypedDict): + type: Literal["file_citation"] + index: int + file_id: str + filename: str + + +class OutputTextContent(TypedDict): + type: Literal["output_text"] + text: str + annotations: list[FileCitationAnnotation] + + +class MessageOutput(TypedDict): + type: Literal["message"] + role: Literal["assistant"] + content: list[OutputTextContent] + + +class FileSearchResultEntry(TypedDict): + file_id: str + filename: str + score: float | None + text: str + attributes: Mapping[str, object] + + +class FileSearchCallOutput(TypedDict): + type: Literal["file_search_call"] + id: str + status: Literal["completed"] + queries: list[str] + search_results: list[FileSearchResultEntry] | None + + +class FunctionCallOutput(TypedDict): + type: Literal["function_call_output"] + call_id: str + output: str + + +class FileSearchToolSpec(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["file_search"] + vector_store_ids: tuple[str, ...] | None = None + + +class FileSearchFunctionCall(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: Literal["function_call"] + name: str | None = None + call_id: str | None = None + id: str | None = None + arguments: str | Mapping[str, object] | None = None + + +class FileSearchArguments(BaseModel): + model_config = ConfigDict(extra="ignore") + + queries: tuple[str, ...] | str | None = None + query: str | None = None + vector_store_id: str | None = None + + +class MessageContentBlock(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: str | None = None + text: str | None = None + + +class MessageOutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: Literal["message"] + content: tuple[MessageContentBlock, ...] | None = None + + +class IncludeOptions(BaseModel): + model_config = ConfigDict(extra="ignore") + + include: tuple[str, ...] | None = None + + +def _validate_or_none(model: type[_ModelT], value: object) -> _ModelT | None: + try: + return model.model_validate(value) + except ValidationError: + return None + # --------------------------------------------------------------------------- # Detection # --------------------------------------------------------------------------- +def _as_file_search_tool(tool: object) -> FileSearchToolSpec | None: + return _validate_or_none(FileSearchToolSpec, tool) + + def should_use_emulated_file_search( tools: Iterable[ToolParam] | None, - provider_config: Any, # BaseResponsesAPIConfig + provider_config: BaseResponsesAPIConfig | None, ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" if not tools: return False - has_fs: Final = any(isinstance(t, dict) and t.get("type") == "file_search" for t in tools) + has_fs: Final = any(_as_file_search_tool(tool) is not None for tool in tools) if not has_fs: return False return provider_config is None or not provider_config.supports_native_file_search() @@ -51,7 +224,7 @@ def should_use_emulated_file_search( # --------------------------------------------------------------------------- -def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: +def _build_function_tool(vector_store_ids: list[str]) -> EmulatedFileSearchTool: """ Create a Responses API function-tool definition that describes file search. The function accepts one or more natural-language queries (like OpenAI's native @@ -62,63 +235,62 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: NOT Chat Completion format (nested under "function"), so that the LiteLLMCompletionResponsesConfig transformation picks up name and description. """ - return { - "type": "function", - "name": FILE_SEARCH_FUNCTION_NAME, - "description": ( + return EmulatedFileSearchTool( + type="function", + name=FILE_SEARCH_FUNCTION_NAME, + description=( "Search the knowledge base for information relevant to the query. " "Use this whenever you need to look up specific facts, documents, " "or content from the vector store. You can provide multiple queries " "to search for different aspects of the information." ), - "parameters": { - "type": "object", - "properties": { - "queries": { - "type": "array", - "items": {"type": "string"}, - "description": ( + parameters=FunctionToolParameters( + type="object", + properties=FunctionToolProperties( + queries=FunctionToolQueriesSchema( + type="array", + items={"type": "string"}, + description=( "One or more search queries to look up in the vector store. " "Multiple queries help find comprehensive information from " "different angles." ), - }, - "vector_store_id": { - "type": "string", - "description": "ID of the vector store to search.", - "enum": vector_store_ids, - }, - }, - "required": ["queries"], - }, - } + ), + vector_store_id=FunctionToolVectorStoreIdSchema( + type="string", + description="ID of the vector store to search.", + enum=vector_store_ids, + ), + ), + required=["queries"], + ), + ) def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[list[ToolParam | EmulatedFileSearchTool], list[str]]: """ Replace all file_search tools with a single function tool. Returns: (new_tools_list, all_vector_store_ids) """ - non_file_search: Final[list[dict[str, Any]]] = [] - vector_store_ids: Final[list[str]] = [] - - for tool in tools or []: - if isinstance(tool, dict) and tool.get("type") == "file_search": - ids = tool.get("vector_store_ids") or [] - vector_store_ids.extend(ids) - else: - non_file_search.append(tool) - - # Deduplicate while preserving order - unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids)) - if unique_ids: - non_file_search.append(_build_function_tool(unique_ids)) - - return non_file_search, unique_ids + parsed: Final = tuple((tool, _as_file_search_tool(tool)) for tool in tools or ()) + non_file_search: Final[list[ToolParam | EmulatedFileSearchTool]] = [ + tool for tool, file_search in parsed if file_search is None + ] + unique_ids: Final[list[str]] = list( + dict.fromkeys( + vector_store_id + for _, file_search in parsed + if file_search is not None + for vector_store_id in file_search.vector_store_ids or () + ) + ) + if not unique_ids: + return non_file_search, unique_ids + return [*non_file_search, _build_function_tool(unique_ids)], unique_ids # --------------------------------------------------------------------------- @@ -126,10 +298,31 @@ def _replace_file_search_tools( # --------------------------------------------------------------------------- +def _results_of(response: VectorStoreSearchResponse | SearchResponseLike) -> tuple[SearchResult, ...]: + results_data: Final = response.get("data") if isinstance(response, dict) else response.data + return tuple(results_data or ()) + + +async def _search_one_vector_store(vector_store_id: str, query: str) -> tuple[SearchResult, ...]: + """Run a single ``asearch`` call, returning no results when the search fails.""" + import litellm.vector_stores.main as vs_main + + try: + return _results_of(await vs_main.asearch(vector_store_id=vector_store_id, query=query)) + except Exception as exc: + verbose_logger.warning( + "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", + query, + vector_store_id, + exc, + ) + return () + + async def _run_vector_searches( queries: list[str], vector_store_ids: list[str], -) -> tuple[list[str], list[VectorStoreSearchResult]]: +) -> tuple[list[str], list[SearchResult]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -140,30 +333,12 @@ async def _run_vector_searches( Returns: (queries_list, combined_results) """ - import litellm.vector_stores.main as vs_main - - all_results: Final[list[VectorStoreSearchResult]] = [] - ids_to_search: Final = vector_store_ids - - # Execute each query against all vector stores - for query in queries: - for vs_id in ids_to_search: - try: - response = await vs_main.asearch( - vector_store_id=vs_id, - query=query, - ) - results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) - if results_data: - all_results.extend(results_data) - except Exception as exc: - verbose_logger.warning( - "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", - query, - vs_id, - exc, - ) - + all_results: Final[list[SearchResult]] = [ + result + for query in queries + for vector_store_id in vector_store_ids + for result in await _search_one_vector_store(vector_store_id=vector_store_id, query=query) + ] return queries, all_results @@ -172,78 +347,89 @@ async def _run_vector_searches( # --------------------------------------------------------------------------- -def _get_field(result: Any, key: str, default: Any = None) -> Any: - """Read a field from either a dict/TypedDict or an attribute-based object.""" +@dataclass(frozen=True, slots=True) +class SearchResultView: + """Normalized view over a dict- or attribute-shaped vector store search result.""" + + score: float | None + file_id: str + filename: str + text: str + attributes: Mapping[str, object] + + +def _content_text(item: VectorStoreResultContent | SearchResultContentLike) -> str: + text: Final = item.get("text") if isinstance(item, dict) else item.text + return text or "" + + +def _view_of(result: SearchResult) -> SearchResultView: if isinstance(result, dict): - return result.get(key, default) - return getattr(result, key, default) + score, file_id, filename = result.get("score"), result.get("file_id"), result.get("filename") + content, attributes = result.get("content"), result.get("attributes") + else: + score, file_id, filename = result.score, result.file_id, result.filename + content, attributes = result.content, result.attributes + text: Final = " ".join(chunk for chunk in (_content_text(item) for item in content or ()) if chunk) + return SearchResultView( + score=score, + file_id=file_id or "", + filename=filename or "", + text=text, + attributes=attributes or {}, + ) + + +def _result_header(index: int, view: SearchResultView) -> str: + segments: Final = ( + f"Result {index}", + view.filename or None, + f"file_id={view.file_id}" if view.file_id else None, + f"score={view.score:.3f}" if view.score is not None else None, + ) + return f"[{' | '.join(segment for segment in segments if segment)}]" def _format_search_results_as_tool_output( - results: list[VectorStoreSearchResult], + results: Sequence[SearchResult], ) -> str: """Serialize search results into a string to pass back as the tool's output.""" if not results: return "No results found in the vector store." - parts: Final[list[str]] = [] - for i, result in enumerate(results, 1): - score = _get_field(result, "score") - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - - header = f"[Result {i}" - if filename: - header += f" | {filename}" - if file_id: - header += f" | file_id={file_id}" - if score is not None: - header += f" | score={score:.3f}" - header += "]" - - parts.append(f"{header}\n{text}") - - return "\n\n".join(parts) + views: Final = tuple(_view_of(result) for result in results) + return "\n\n".join(f"{_result_header(index, view)}\n{view.text}" for index, view in enumerate(views, 1)) def _build_search_results_for_include( - results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: + results: Sequence[SearchResult], +) -> list[FileSearchResultEntry]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). - All chunks are returned — no deduplication by file_id — matching the + All chunks are returned, with no deduplication by file_id, matching the behaviour of OpenAI's native file_search which surfaces every relevant chunk even when multiple chunks originate from the same document. """ - formatted: Final[list[dict[str, Any]]] = [] - for result in results: - file_id = _get_field(result, "file_id") or "" - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - formatted.append( - { - "file_id": file_id, - "filename": _get_field(result, "filename") or "", - "score": _get_field(result, "score"), - "text": text, - "attributes": _get_field(result, "attributes") or {}, - } + return [ + FileSearchResultEntry( + file_id=view.file_id, + filename=view.filename, + score=view.score, + text=view.text, + attributes=view.attributes, ) - return formatted + for view in (_view_of(result) for result in results) + ] def _build_file_search_call_output( call_id: str, queries: list[str], - results: list[VectorStoreSearchResult] | None = None, + results: Sequence[SearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> FileSearchCallOutput: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -253,85 +439,89 @@ def _build_file_search_call_output( include_search_results: Populate search_results when the caller passed ``include=["file_search_call.results"]``. """ - search_results = None - if include_search_results and results: - search_results = _build_search_results_for_include(results) - return { - "type": "file_search_call", - "id": call_id, - "status": "completed", - "queries": queries, - "search_results": search_results, - } + search_results: Final = ( + _build_search_results_for_include(results) if include_search_results and results else None + ) + return FileSearchCallOutput( + type="file_search_call", + id=call_id, + status="completed", + queries=queries, + search_results=search_results, + ) def _build_file_citation_annotations( - results: list[VectorStoreSearchResult], + results: Sequence[SearchResult], text: str, -) -> list[dict[str, Any]]: +) -> list[FileCitationAnnotation]: """ Build file_citation annotations for the text. Each result with a file_id gets a citation at the end of the text. """ - annotations: Final[list[dict[str, Any]]] = [] - index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() - - for result in results: - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - if not file_id or file_id in seen_file_ids: - continue - seen_file_ids.add(file_id) - annotations.append( - { - "type": "file_citation", - "index": index, - "file_id": file_id, - "filename": filename or "", - } + views: Final = tuple(_view_of(result) for result in results) + ordered_file_ids: Final = tuple(dict.fromkeys(view.file_id for view in views if view.file_id)) + filename_by_file_id: Final = {view.file_id: view.filename for view in reversed(views) if view.file_id} + return [ + FileCitationAnnotation( + type="file_citation", + index=len(text), + file_id=file_id, + filename=filename_by_file_id.get(file_id, ""), ) - - return annotations + for file_id in ordered_file_ids + ] def _build_message_output( response_text: str, - results: list[VectorStoreSearchResult], -) -> dict[str, Any]: + results: Sequence[SearchResult], +) -> MessageOutput: """Build the message output item with optional file_citation annotations.""" - annotations: Final = _build_file_citation_annotations(results, response_text) - return { - "type": "message", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": response_text, - "annotations": annotations, - } + return MessageOutput( + type="message", + role="assistant", + content=[ + OutputTextContent( + type="output_text", + text=response_text, + annotations=_build_file_citation_annotations(results, response_text), + ) ], - } + ) + + +def _output_items_of(response: ResponsesAPIResponse) -> tuple[object, ...]: + return tuple(response.output) def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: """Pull the assistant's text from the provider's response.""" - for item in response.output: - item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) - if item_type == "message": - content = item.get("content") if isinstance(item, dict) else getattr(item, "content", []) - for block in content or []: - block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) - if block_type == "output_text": - raw = block.get("text") if isinstance(block, dict) else getattr(block, "text", "") - return str(raw) if raw is not None else "" - return "" + messages: Final = tuple( + message + for message in (_validate_or_none(MessageOutputItem, item) for item in _output_items_of(response)) + if message is not None + ) + texts: Final = tuple( + block.text or "" for message in messages for block in message.content or () if block.type == "output_text" + ) + return texts[0] if texts else "" + + +def _response_cost(hidden_params: Mapping[str, object]) -> float | None: + cost: Final = hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + +def _hidden_params_of(response: ResponsesAPIResponse) -> Mapping[str, object]: + raw: Final[Mapping[str, object] | None] = getattr(response, "_hidden_params", None) + return raw or {} def _synthesize_responses_api_response( original_response: ResponsesAPIResponse, - file_search_call_output: dict[str, Any], - message_output: dict[str, Any], + file_search_call_output: FileSearchCallOutput, + message_output: MessageOutput, first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,30 +533,25 @@ def _synthesize_responses_api_response( synthesized _hidden_params so that billing callbacks see the total cost of both provider calls that the emulated flow makes. """ - synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output] + synthesized_output: Final[list[ResponseOutputItem | dict[str, object]]] = [ + dict(file_search_call_output), + dict(message_output), + ] synthesized: Final = ResponsesAPIResponse( - id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), + id=original_response.id, object="response", - created_at=getattr(original_response, "created_at", int(time.time())), + created_at=original_response.created_at, status="completed", - model=getattr(original_response, "model", ""), - output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output), - usage=getattr(original_response, "usage", None), + model=original_response.model, + output=synthesized_output, + usage=original_response.usage, error=None, ) - if hasattr(original_response, "_hidden_params"): - hidden: Final = dict(getattr(original_response, "_hidden_params") or {}) - if first_response is not None and hasattr(first_response, "_hidden_params"): - first_hidden: Final = getattr(first_response, "_hidden_params") or {} - first_cost: Final = ( - first_hidden.get("response_cost") - if isinstance(first_hidden, dict) - else getattr(first_hidden, "response_cost", None) - ) - if first_cost is not None: - current_cost: Final = hidden.get("response_cost") if isinstance(hidden, dict) else 0 - hidden["response_cost"] = (current_cost or 0) + first_cost - synthesized._hidden_params = hidden + hidden: Final = _hidden_params_of(original_response) + first_cost: Final = _response_cost(_hidden_params_of(first_response)) if first_response is not None else None + synthesized._hidden_params = ( + {**hidden} if first_cost is None else {**hidden, "response_cost": (_response_cost(hidden) or 0) + first_cost} + ) return synthesized @@ -375,127 +560,118 @@ def _synthesize_responses_api_response( # --------------------------------------------------------------------------- -async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover – thin wrapper for patching in tests +# any-ok: thin, test-patchable seam forwarding provider pass-through params to aresponses +async def _call_aresponses(input, model: str, tools, **kwargs) -> ResponsesAPIResponse: from litellm.responses.main import aresponses - return await aresponses(input=input, model=model, tools=tools, **kwargs) + response: Final = await aresponses(input=input, model=model, tools=tools, **kwargs) + if isinstance(response, ResponsesAPIResponse): + return response + raise ValueError("emulated file_search does not support streaming responses") def _prepare_emulated_file_search_call( - kwargs: dict[str, Any], -) -> tuple[bool, dict[str, Any]]: - include_items: Final[list[str]] = list(kwargs.get("include") or []) - include_search_results: Final = "file_search_call.results" in include_items + kwargs: Mapping[str, object], +) -> tuple[bool, dict[str, object]]: + include_options: Final = _validate_or_none(IncludeOptions, kwargs) + include_search_results: Final = "file_search_call.results" in ((include_options and include_options.include) or ()) - original_stream: Final = kwargs.get("stream") - updated_kwargs = kwargs - if original_stream: - verbose_logger.debug( - "Streaming is not yet supported for emulated file_search. Disabling stream for this request." - ) - updated_kwargs = {**kwargs, "stream": False} + if not kwargs.get("stream"): + return include_search_results, {**kwargs} - return include_search_results, updated_kwargs + verbose_logger.debug("Streaming is not yet supported for emulated file_search. Disabling stream for this request.") + return include_search_results, {**kwargs, "stream": False} -def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]: - """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" - if isinstance(tool_call, dict): - call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) - raw_args = tool_call.get("arguments") or "{}" - else: - raw_call_id: Final = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) or fallback_call_id - call_id = str(raw_call_id) - raw_args = getattr(tool_call, "arguments", "{}") or "{}" - return call_id, raw_args - - -def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: +def _resolve_queries_from_args(args: FileSearchArguments, input: str | ResponseInputParam) -> list[str]: """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks.""" - queries_from_call: Final = args.get("queries") - if not queries_from_call: - # Fallback: check for single "query" field (backward compat) - single_query: Final = args.get("query") - return [single_query] if single_query else [str(input)] - if not isinstance(queries_from_call, list): - return [str(queries_from_call)] - return queries_from_call + if not args.queries: + return [args.query] if args.query else [str(input)] + if isinstance(args.queries, str): + return [args.queries] + return list(args.queries) + + +def _parse_file_search_arguments(raw_arguments: str | Mapping[str, object] | None) -> FileSearchArguments: + if isinstance(raw_arguments, str): + try: + return FileSearchArguments.model_validate_json(raw_arguments) + except ValidationError: + return FileSearchArguments() + return _validate_or_none(FileSearchArguments, raw_arguments) or FileSearchArguments() async def _execute_file_search_tool_calls( - file_search_calls: list[Any], + file_search_calls: Sequence[FileSearchFunctionCall], all_vs_ids: list[str], - input: Any, + input: str | ResponseInputParam, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[list[FunctionCallOutput], list[str], list[SearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] + tool_results: Final[list[FunctionCallOutput]] = [] all_queries: Final[list[str]] = [] - all_results: Final[list[VectorStoreSearchResult]] = [] + all_results: Final[list[SearchResult]] = [] for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) - - try: - args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - args = {} - - queries_from_call = _resolve_queries_from_args(args, input) - - vs_id_arg = args.get("vector_store_id") - vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - + call_id = str(tool_call.call_id or tool_call.id or file_search_call_id) + args = _parse_file_search_arguments(tool_call.arguments) queries, results = await _run_vector_searches( - queries=queries_from_call, - vector_store_ids=vs_ids_for_call, + queries=_resolve_queries_from_args(args, input), + vector_store_ids=[args.vector_store_id] if args.vector_store_id else all_vs_ids, ) all_queries.extend(queries) all_results.extend(results) tool_results.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": _format_search_results_as_tool_output(results), - } + FunctionCallOutput( + type="function_call_output", + call_id=call_id, + output=_format_search_results_as_tool_output(results), + ) ) return tool_results, all_queries, all_results +def _as_plain_input_item(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item + + def _build_follow_up_input( - input: Any, + input: str | ResponseInputParam, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: Sequence[FunctionCallOutput], +) -> list[object]: """Assemble the follow-up call input: original messages + first-response output + tool results. Including all output items (text blocks, reasoning, non-file-search calls) ensures providers like Anthropic that emit text before the tool call have complete conversation context. Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ - original_input_items: Final = ( + original_input_items: Final[list[object]] = ( list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] ) - first_response_output_items: Final[list[Any]] = [] - for _item in first_response.output: - if isinstance(_item, dict): - first_response_output_items.append(_item) - elif hasattr(_item, "model_dump"): - first_response_output_items.append(_item.model_dump(exclude_none=True)) - else: - first_response_output_items.append(_item) + return [ + *original_input_items, + *(_as_plain_input_item(item) for item in _output_items_of(first_response)), + *tool_results, + ] - return original_input_items + first_response_output_items + tool_results + +def _file_search_calls_in(response: ResponsesAPIResponse) -> tuple[FileSearchFunctionCall, ...]: + return tuple( + call + for call in (_validate_or_none(FileSearchFunctionCall, item) for item in _output_items_of(response)) + if call is not None and call.name == FILE_SEARCH_FUNCTION_NAME + ) async def aresponses_with_emulated_file_search( - input: Any, + input: str | ResponseInputParam, model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call - **kwargs: Any, + **kwargs: object, ) -> ResponsesAPIResponse: """ Emulated file_search for providers that don't support it natively. @@ -504,7 +680,7 @@ async def aresponses_with_emulated_file_search( runs vector search, and synthesizes an OpenAI-format response. """ # Determine whether caller wants search_results populated in the output. - _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) + include_search_results, forwarded_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool transformed_tools, all_vs_ids = _replace_file_search_tools(tools) @@ -512,42 +688,25 @@ async def aresponses_with_emulated_file_search( # 2. First provider call — provider will call the file_search function. # Mark as an internal sub-call so wrapper_async skips billing callbacks; # the parent litellm_logging_obj (propagated via kwargs) fires once at the end. - _prev_internal: Final = is_internal_call.get() + prev_internal: Final = is_internal_call.get() is_internal_call.set(True) try: - first_response: Final[ResponsesAPIResponse] = cast( - ResponsesAPIResponse, - await _call_aresponses( - input=input, - model=model, - tools=transformed_tools or None, - **kwargs, - ), + first_response: Final = await _call_aresponses( + input=input, + model=model, + tools=transformed_tools or None, + **forwarded_kwargs, ) finally: - is_internal_call.set(_prev_internal) + is_internal_call.set(prev_internal) # 3. Look for a file_search function_call in the output - file_search_calls: Final = [ - item - for item in first_response.output - if ( - isinstance(item, dict) - and item.get("type") == "function_call" - and item.get("name") == FILE_SEARCH_FUNCTION_NAME - ) - or ( - hasattr(item, "type") - and getattr(item, "type") == "function_call" - and getattr(item, "name", None) == FILE_SEARCH_FUNCTION_NAME - ) - ] + file_search_calls: Final = _file_search_calls_in(first_response) if not file_search_calls: # Provider answered without calling the tool (e.g. it had enough context). # Return as-is wrapped in OpenAI format. call_id: Final = f"fs_{uuid.uuid4().hex[:24]}" - response_text = _extract_text_from_responses_output(first_response) return _synthesize_responses_api_response( original_response=first_response, file_search_call_output=_build_file_search_call_output( @@ -556,7 +715,7 @@ async def aresponses_with_emulated_file_search( results=None, include_search_results=False, ), - message_output=_build_message_output(response_text, []), + message_output=_build_message_output(_extract_text_from_responses_output(first_response), []), ) # 4. Execute each file_search tool call @@ -579,29 +738,24 @@ async def aresponses_with_emulated_file_search( # Also an internal sub-call; billing is suppressed so the outer call fires once. is_internal_call.set(True) try: - final_response: Final[ResponsesAPIResponse] = cast( - ResponsesAPIResponse, - await _call_aresponses( - input=follow_up_input, - model=model, - tools=None, # no tools needed for the answer step - **kwargs, - ), + final_response: Final = await _call_aresponses( + input=follow_up_input, + model=model, + tools=None, # no tools needed for the answer step + **forwarded_kwargs, ) finally: - is_internal_call.set(_prev_internal) + is_internal_call.set(prev_internal) # 7. Synthesize OpenAI-format output - response_text = _extract_text_from_responses_output(final_response) - return _synthesize_responses_api_response( original_response=final_response, file_search_call_output=_build_file_search_call_output( call_id=file_search_call_id, queries=all_queries or [str(input)], results=all_results, - include_search_results=_include_search_results, + include_search_results=include_search_results, ), - message_output=_build_message_output(response_text, all_results), + message_output=_build_message_output(_extract_text_from_responses_output(final_response), all_results), first_response=first_response, ) diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py index 24879ce83f9..e02481bddc2 100644 --- a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py @@ -7,7 +7,10 @@ from unittest.mock import Mock import httpx import pytest -from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig +from litellm.llms.runwayml.videos.transformation import ( + RunwayMLTaskResponse, + RunwayMLVideoConfig, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoObject @@ -126,13 +129,17 @@ class TestRunwayMLVideoTransformation: "status": "SUCCEEDED", "output": ["https://dnznrvs05pmza.cloudfront.net/video.mp4"], } - video_url = self.config._extract_video_url_from_response(response_data) + video_url = self.config._extract_video_url_from_response( + RunwayMLTaskResponse.model_validate(response_data) + ) assert video_url == "https://dnznrvs05pmza.cloudfront.net/video.mp4" # Test error handling when video is still processing processing_response = {"id": "test-id", "status": "RUNNING", "output": None} with pytest.raises(ValueError, match="still processing"): - self.config._extract_video_url_from_response(processing_response) + self.config._extract_video_url_from_response( + RunwayMLTaskResponse.model_validate(processing_response) + ) def test_transform_video_status_encodes_video_id_path_segment(self): """Test task IDs are encoded before being appended to Runway URLs."""