From a8fab3990ff38e61eb62779c5676d2579ecc59ae Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Sat, 3 Oct 2026 21:58:01 -0500 Subject: [PATCH] fix(vertex_ai): forward agent engine attachments --- .../vertex_ai/agent_engine/transformation.py | 156 ++++++++++++- .../agent_engine/test_transformation.py | 214 +++++++++++++++++- 2 files changed, 353 insertions(+), 17 deletions(-) diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index fd1ff24fd52..7a118366e77 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -9,24 +9,34 @@ API Reference: - :streamQuery endpoint - for actual queries (stream_query method) """ +import base64 +import binascii import json +import mimetypes +import re from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Optional, Union, cast +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, Union, cast +from urllib.parse import urlparse import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - convert_content_list_to_str, -) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionFileObject, + ChatCompletionImageObject, + ChatCompletionTextObject, +) +from litellm.types.llms.vertex_ai import ContentType, PartType from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: @@ -52,6 +62,131 @@ class VertexAgentEngineError(BaseLLMException): super().__init__(message=message, status_code=status_code) +_AgentEngineContentBlock: TypeAlias = ChatCompletionTextObject | ChatCompletionImageObject | ChatCompletionFileObject +_CONTENT_ADAPTER: Final = TypeAdapter(tuple[_AgentEngineContentBlock, ...]) + + +@dataclass(frozen=True, slots=True) +class _AgentEngineInputError: + message: str + + +def _data_url_content(source: str, mime_hint: str | None) -> tuple[str, str] | _AgentEngineInputError: + header, separator, encoded = source.partition(",") + if not separator or not header.endswith(";base64"): + return _AgentEngineInputError("Agent Engine requires base64-encoded data URLs") + declared_mime: Final = header[5:-7] + if mime_hint and mime_hint != declared_mime: + return _AgentEngineInputError("Agent Engine media format conflicts with the data URL MIME type") + return declared_mime, encoded + + +def _media_to_part( + source: str, + mime_hint: str | None, + filename: str | None, + is_image: bool, + allow_inline: bool = True, +) -> PartType | _AgentEngineInputError: + try: + parsed: Final = urlparse(source) + except ValueError: + return _AgentEngineInputError("Agent Engine requires a valid media URI") + is_uri: Final = parsed.scheme in ("gs", "https") + if is_uri and (not parsed.netloc or not parsed.path.strip("/")): + return _AgentEngineInputError("Agent Engine requires a valid gs:// or https:// media URI") + if not allow_inline and not is_uri: + return _AgentEngineInputError( + "Agent Engine cannot resolve file IDs; provide a gs:// or https:// URI or file_data" + ) + if parsed.scheme not in ("", "gs", "https", "data"): + return _AgentEngineInputError("Agent Engine supports gs://, https:// and base64-encoded media only") + decoded: Final = ( + _data_url_content(source, mime_hint) + if source.startswith("data:") + else (mime_hint or mimetypes.guess_type(filename or parsed.path)[0], source) + ) + if isinstance(decoded, _AgentEngineInputError): + return decoded + mime_type, data = decoded + if not mime_type or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9.+_-]*/[A-Za-z0-9][A-Za-z0-9.+_-]*", mime_type): + return _AgentEngineInputError( + "Agent Engine requires a MIME type; provide format or a filename with an extension" + ) + if is_image and not mime_type.startswith("image/"): + return _AgentEngineInputError("Agent Engine image_url requires an image MIME type; use file for documents") + if is_uri: + return {"file_data": {"file_uri": source, "mime_type": mime_type}} + if not data: + return _AgentEngineInputError("Agent Engine media data cannot be empty") + try: + base64.b64decode(data, validate=True) + except (binascii.Error, ValueError): + return _AgentEngineInputError("Agent Engine media data must be valid base64") + return {"inline_data": {"data": data, "mime_type": mime_type}} + + +def _content_block_to_part(block: _AgentEngineContentBlock) -> PartType | _AgentEngineInputError: + match block["type"]: + case "text": + return {"text": block["text"]} + case "image_url": + image: Final = block["image_url"] + return _media_to_part( + source=image if isinstance(image, str) else image["url"], + mime_hint=None if isinstance(image, str) else image.get("format"), + filename=None, + is_image=True, + ) + case "file": + file: Final = block["file"] + if bool(file.get("file_data")) == bool(file.get("file_id")): + return _AgentEngineInputError("Agent Engine requires exactly one of file_data or file_id") + return _media_to_part( + source=file.get("file_data") or file.get("file_id") or "", + mime_hint=file.get("format"), + filename=file.get("filename"), + is_image=False, + allow_inline=bool(file.get("file_data")), + ) + + +def _message_to_agent_input(message: Mapping[str, object]) -> str | ContentType | _AgentEngineInputError: + content: Final = message.get("content") + if content is None: + return "" + if isinstance(content, str): + return content + try: + blocks: Final = _CONTENT_ADAPTER.validate_python(content) + except ValidationError: + return _AgentEngineInputError("Agent Engine supports text, image_url and file parts with valid required fields") + converted: Final = tuple(_content_block_to_part(block) for block in blocks) + error: Final = next((part for part in converted if isinstance(part, _AgentEngineInputError)), None) + if error is not None: + return error + parts: Final = [part for part in converted if not isinstance(part, _AgentEngineInputError)] + if all(block["type"] == "text" for block in blocks): + return "".join(part.get("text", "") for part in parts) + if message["role"] != "user": + return _AgentEngineInputError("Agent Engine media must be in the final user message") + return {"role": "user", "parts": parts} + + +def _messages_to_agent_input(messages: list[AllMessageValues]) -> str | ContentType | _AgentEngineInputError: + if not messages: + return _AgentEngineInputError("Agent Engine requires at least one message") + previous: Final = tuple(_message_to_agent_input(message) for message in messages[:-1]) + error: Final = next((part for part in previous if isinstance(part, _AgentEngineInputError)), None) + if error is not None: + return error + if any(not isinstance(part, str) for part in previous): + return _AgentEngineInputError( + "Agent Engine media must be in the final user message; earlier attachments are not forwarded" + ) + return _message_to_agent_input(messages[-1]) + + class VertexAgentEngineConfig(BaseConfig, VertexBase): """ Configuration for Vertex AI Agent Engine (Reasoning Engines). @@ -197,22 +332,21 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): } } """ - # Use the last message content as the prompt - prompt: Final = convert_content_list_to_str(messages[-1]) + prompt: Final = _messages_to_agent_input(messages) + if isinstance(prompt, _AgentEngineInputError): + raise VertexAgentEngineError(status_code=400, message=prompt.message) # Get user_id and session_id user_id: Final = self._get_user_id(optional_params) session_id: Final = self._get_session_id(optional_params) # Build the input - input_data: Final[dict[str, str]] = { + input_data: Final[dict[str, str | ContentType]] = { "message": prompt, "user_id": user_id, + **({"session_id": session_id} if session_id else {}), } - if session_id: - input_data["session_id"] = session_id - # Build the request payload # Note: stream_query is used for both streaming and non-streaming # The difference is the endpoint (:streamQuery vs :query) diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py index 19616682c59..881c7bb98ca 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -4,14 +4,19 @@ Tests for Vertex AI Agent Engine transformation. Tests the request transformation and streaming chunk parsing without making real API calls. """ +from copy import deepcopy +from typing import Final, Literal import pytest - +from litellm import BadRequestError +from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) -from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig +from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig, VertexAgentEngineError +from litellm.types.llms.openai import AllMessageValues, OpenAIMessageContent, OpenAIMessageContentListBlock +from litellm.types.llms.vertex_ai import ContentType class TestVertexAgentEngineTransformRequest: @@ -94,10 +99,7 @@ class TestVertexAgentEngineChunkParser: result = iterator.chunk_parser(chunk) - assert ( - result.choices[0].delta.content - == "Hello! I can help you with financial analysis." - ) + assert result.choices[0].delta.content == "Hello! I can help you with financial analysis." assert result.choices[0].delta.role == "assistant" assert result.choices[0].finish_reason == "stop" assert result.usage["prompt_tokens"] == 100 @@ -125,3 +127,203 @@ class TestVertexAgentEngineChunkParser: assert result.choices[0].delta.content == "Partial response..." assert result.choices[0].finish_reason is None assert result.usage is None + + +@pytest.mark.parametrize( + "content, expected", + [ + ([{"type": "text", "text": "one"}, {"type": "text", "text": "two"}], "onetwo"), + ( + [{"type": "image_url", "image_url": "gs://bucket/image.png"}], + {"role": "user", "parts": [{"file_data": {"file_uri": "gs://bucket/image.png", "mime_type": "image/png"}}]}, + ), + ( + [{"type": "image_url", "image_url": {"url": "https://example.com/photo.jpg?signature=123"}}], + { + "role": "user", + "parts": [ + { + "file_data": { + "file_uri": "https://example.com/photo.jpg?signature=123", + "mime_type": "image/jpeg", + } + } + ], + }, + ), + ( + [{"type": "image_url", "image_url": {"url": "gs://bucket/no-extension", "format": "image/webp"}}], + { + "role": "user", + "parts": [{"file_data": {"file_uri": "gs://bucket/no-extension", "mime_type": "image/webp"}}], + }, + ), + ( + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}], + {"role": "user", "parts": [{"inline_data": {"data": "AA==", "mime_type": "image/png"}}]}, + ), + ( + [{"type": "file", "file": {"file_data": "gs://bucket/report.pdf"}}], + { + "role": "user", + "parts": [{"file_data": {"file_uri": "gs://bucket/report.pdf", "mime_type": "application/pdf"}}], + }, + ), + ( + [{"type": "file", "file": {"file_id": "gs://bucket/report.pdf"}}], + { + "role": "user", + "parts": [{"file_data": {"file_uri": "gs://bucket/report.pdf", "mime_type": "application/pdf"}}], + }, + ), + ( + [{"type": "file", "file": {"file_data": "JVBERi0xLjQ=", "filename": "report.pdf"}}], + {"role": "user", "parts": [{"inline_data": {"data": "JVBERi0xLjQ=", "mime_type": "application/pdf"}}]}, + ), + ( + [{"type": "file", "file": {"file_data": "data:application/pdf;base64,JVBERi0xLjQ="}}], + {"role": "user", "parts": [{"inline_data": {"data": "JVBERi0xLjQ=", "mime_type": "application/pdf"}}]}, + ), + ( + [ + {"type": "text", "text": "before"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + {"type": "text", "text": "after"}, + ], + { + "role": "user", + "parts": [ + {"text": "before"}, + {"inline_data": {"data": "AA==", "mime_type": "image/png"}}, + {"text": "after"}, + ], + }, + ), + ], +) +def test_attachment_parts_reach_agent_engine_without_changing_text_path( + content: OpenAIMessageContent, expected: str | ContentType +) -> None: + config: Final = VertexAgentEngineConfig() + original: Final = deepcopy(content) + messages: Final[list[AllMessageValues]] = [{"role": "user", "content": content}] + result: Final = config.transform_request( + model="agent_engine/123", + messages=messages, + optional_params={"user_id": "user-123", "session_id": "session-456"}, + litellm_params={}, + headers={}, + ) + assert result == { + "class_method": "stream_query", + "input": { + "message": expected, + "user_id": "user-123", + "session_id": "session-456", + }, + } + assert messages == [{"role": "user", "content": original}] + + +@pytest.mark.parametrize( + "part", + [ + {"type": "input_audio", "input_audio": {"data": "AA==", "format": "wav"}}, + {"type": "video_url", "video_url": {"url": "gs://bucket/movie.mp4"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,!!!"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,"}}, + {"type": "image_url", "image_url": {"url": "data:image/png,raw"}}, + {"type": "image_url", "image_url": {"url": "file:///tmp/image.png"}}, + {"type": "image_url", "image_url": {"url": "gs://bucket/no-extension"}}, + {"type": "image_url", "image_url": {"url": "data:application/pdf;base64,AA=="}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA==", "format": "image/jpeg"}}, + {"type": "image_url", "image_url": {}}, + {"type": "image_url", "image_url": {"url": "https://[invalid/image.png"}}, + {"type": "file", "file": {"file_id": "file-123"}}, + {"type": "file", "file": {"file_data": "AA=="}}, + {"type": "file", "file": {"file_data": "gs://bucket/a.pdf", "file_id": "gs://bucket/b.pdf"}}, + {"type": "file", "file": {"file_data": "AA==", "format": "not-a-mime-type"}}, + {"type": "file", "file": {}}, + {"type": "text"}, + {"type": "unknown", "data": "AA=="}, + ], +) +def test_unsupported_or_invalid_parts_raise_instead_of_disappearing(part: OpenAIMessageContentListBlock) -> None: + config: Final = VertexAgentEngineConfig() + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [{"type": "text", "text": "attachment"}, part]}, + ] + with pytest.raises(VertexAgentEngineError) as error: + config.transform_request( + model="agent_engine/123", + messages=messages, + optional_params={"user_id": "user"}, + litellm_params={}, + headers={}, + ) + assert error.value.status_code == 400 + assert "Agent Engine" in error.value.message + + +@pytest.mark.parametrize("role, following", [("assistant", False), ("user", True)]) +def test_media_is_rejected_when_it_cannot_be_forwarded(role: Literal["assistant", "user"], following: bool) -> None: + config: Final = VertexAgentEngineConfig() + media: Final[AllMessageValues] = { + "role": role, + "content": [ + {"type": "image_url", "image_url": {"url": "gs://bucket/image.png"}}, + ], + } + messages: Final[list[AllMessageValues]] = [media, {"role": "user", "content": "next"}] if following else [media] + with pytest.raises(VertexAgentEngineError) as error: + config.transform_request( + model="agent_engine/123", + messages=messages, + optional_params={"user_id": "user"}, + litellm_params={}, + headers={}, + ) + assert error.value.status_code == 400 + + +def test_invalid_media_maps_to_public_bad_request_error() -> None: + config: Final = VertexAgentEngineConfig() + messages: Final[list[AllMessageValues]] = [ + { + "role": "user", + "content": [ + {"type": "input_audio", "input_audio": {"data": "AA==", "format": "wav"}}, + ], + } + ] + with pytest.raises(VertexAgentEngineError) as provider_error: + config.transform_request( + model="agent_engine/123", + messages=messages, + optional_params={"user_id": "user"}, + litellm_params={}, + headers={}, + ) + with pytest.raises(BadRequestError) as public_error: + exception_type( + model="agent_engine/123", + original_exception=provider_error.value, + custom_llm_provider="vertex_ai", + completion_kwargs={}, + extra_kwargs={}, + ) + assert public_error.value.status_code == 400 + assert "Agent Engine" in public_error.value.message + + +def test_text_history_keeps_the_last_message_string_path() -> None: + config: Final = VertexAgentEngineConfig() + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [{"type": "text", "text": "earlier"}]}, + {"role": "assistant", "content": "response"}, + {"role": "user", "content": "last"}, + ] + result: Final = config.transform_request( + model="agent_engine/123", messages=messages, optional_params={"user_id": "user"}, litellm_params={}, headers={} + ) + assert result == {"class_method": "stream_query", "input": {"message": "last", "user_id": "user"}}