From a8fab3990ff38e61eb62779c5676d2579ecc59ae Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Sat, 3 Oct 2026 21:58:01 -0500 Subject: [PATCH 1/3] 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"}} From d57c79ea6d4f50f65d51cd95bab7700fc44255c0 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Sat, 3 Oct 2026 23:57:00 -0500 Subject: [PATCH 2/3] fix(vertex_ai): cover agent engine input edge cases --- .../vertex_ai/agent_engine/transformation.py | 2 +- .../agent_engine/test_transformation.py | 41 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 7a118366e77..4795a6e50ff 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -138,7 +138,7 @@ def _content_block_to_part(block: _AgentEngineContentBlock) -> PartType | _Agent filename=None, is_image=True, ) - case "file": + case _: 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") 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 881c7bb98ca..0e0bbe195f5 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -132,6 +132,7 @@ class TestVertexAgentEngineChunkParser: @pytest.mark.parametrize( "content, expected", [ + ([], ""), ([{"type": "text", "text": "one"}, {"type": "text", "text": "two"}], "onetwo"), ( [{"type": "image_url", "image_url": "gs://bucket/image.png"}], @@ -239,6 +240,8 @@ def test_attachment_parts_reach_agent_engine_without_changing_text_path( {"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": "image_url", "image_url": {"url": "gs://bucket"}}, + {"type": "image_url", "image_url": {"url": "https:///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"}}, @@ -327,3 +330,41 @@ def test_text_history_keeps_the_last_message_string_path() -> None: 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"}} + + +def test_assistant_without_content_preserves_empty_string_path() -> None: + config: Final = VertexAgentEngineConfig() + messages: Final[list[AllMessageValues]] = [{"role": "assistant", "content": None}] + 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": "", "user_id": "user"}} + + +@pytest.mark.parametrize( + "messages, expected_error", + [ + ([], "Agent Engine requires at least one message"), + ( + [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "gs://bucket"}}]}, + {"role": "user", "content": "next"}, + ], + "Agent Engine requires a valid gs:// or https:// media URI", + ), + ], +) +def test_invalid_messages_fail_before_forwarding_the_last_text( + messages: list[AllMessageValues], expected_error: str +) -> None: + config: Final = VertexAgentEngineConfig() + 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 error.value.message == expected_error From 2d115a275fe0142eb45863405f05c0adbac23cf4 Mon Sep 17 00:00:00 2001 From: AkilanChithra Date: Sun, 4 Oct 2026 00:21:46 -0500 Subject: [PATCH 3/3] fix(vertex_ai): preserve agent engine text context --- .../vertex_ai/agent_engine/transformation.py | 19 +-- .../agent_engine/test_transformation.py | 136 +++++++++++++----- 2 files changed, 110 insertions(+), 45 deletions(-) diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 4795a6e50ff..724334ece5d 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -24,6 +24,7 @@ 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 extract_search_results_text from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, @@ -153,10 +154,11 @@ def _content_block_to_part(block: _AgentEngineContentBlock) -> PartType | _Agent def _message_to_agent_input(message: Mapping[str, object]) -> str | ContentType | _AgentEngineInputError: content: Final = message.get("content") + search_text: Final = extract_search_results_text(message.get("search_results")) if content is None: - return "" + return search_text if isinstance(content, str): - return content + return content + search_text try: blocks: Final = _CONTENT_ADAPTER.validate_python(content) except ValidationError: @@ -167,23 +169,16 @@ def _message_to_agent_input(message: Mapping[str, object]) -> str | ContentType 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) + return "".join(part.get("text", "") for part in parts) + search_text if message["role"] != "user": return _AgentEngineInputError("Agent Engine media must be in the final user message") - return {"role": "user", "parts": parts} + search_parts: Final[list[PartType]] = [{"text": search_text}] if search_text else [] + return {"role": "user", "parts": parts + search_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]) 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 0e0bbe195f5..3f643f7c1bc 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -5,9 +5,10 @@ Tests the request transformation and streaming chunk parsing without making real """ from copy import deepcopy -from typing import Final, Literal +from typing import Final import pytest +from typing_extensions import ReadOnly from litellm import BadRequestError from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -15,7 +16,14 @@ from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig, VertexAgentEngineError -from litellm.types.llms.openai import AllMessageValues, OpenAIMessageContent, OpenAIMessageContentListBlock +from litellm.types.llms.bedrock import SearchResultBlock +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionUserMessage, + OpenAIChatCompletionAssistantMessage, + OpenAIMessageContent, + OpenAIMessageContentListBlock, +) from litellm.types.llms.vertex_ai import ContentType @@ -268,16 +276,11 @@ def test_unsupported_or_invalid_parts_raise_instead_of_disappearing(part: OpenAI 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: +def test_media_in_final_assistant_message_is_rejected() -> 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] + messages: Final[list[AllMessageValues]] = [ + {"role": "assistant", "content": [{"type": "image_url", "image_url": {"url": "gs://bucket/image.png"}}]}, + ] with pytest.raises(VertexAgentEngineError) as error: config.transform_request( model="agent_engine/123", @@ -287,6 +290,7 @@ def test_media_is_rejected_when_it_cannot_be_forwarded(role: Literal["assistant" headers={}, ) assert error.value.status_code == 400 + assert error.value.message == "Agent Engine media must be in the final user message" def test_invalid_media_maps_to_public_bad_request_error() -> None: @@ -341,30 +345,96 @@ def test_assistant_without_content_preserves_empty_string_path() -> None: assert result == {"class_method": "stream_query", "input": {"message": "", "user_id": "user"}} -@pytest.mark.parametrize( - "messages, expected_error", - [ - ([], "Agent Engine requires at least one message"), - ( - [ - {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "gs://bucket"}}]}, - {"role": "user", "content": "next"}, - ], - "Agent Engine requires a valid gs:// or https:// media URI", - ), - ], -) -def test_invalid_messages_fail_before_forwarding_the_last_text( - messages: list[AllMessageValues], expected_error: str -) -> None: +def test_empty_message_list_returns_a_clear_error() -> None: config: Final = VertexAgentEngineConfig() with pytest.raises(VertexAgentEngineError) as error: config.transform_request( - model="agent_engine/123", - messages=messages, - optional_params={"user_id": "user"}, - litellm_params={}, - headers={}, + model="agent_engine/123", messages=[], optional_params={"user_id": "user"}, litellm_params={}, headers={} ) assert error.value.status_code == 400 - assert error.value.message == expected_error + assert error.value.message == "Agent Engine requires at least one message" + + +@pytest.mark.parametrize( + "part", + [ + {"type": "image_url", "image_url": {"url": "gs://bucket/image.png"}}, + {"type": "image_url", "image_url": {"url": "gs://bucket"}}, + {"type": "input_audio", "input_audio": {"data": "AA==", "format": "wav"}}, + ], +) +def test_earlier_media_keeps_existing_final_message_behavior(part: OpenAIMessageContentListBlock) -> None: + config: Final = VertexAgentEngineConfig() + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [part]}, + {"role": "user", "content": "next"}, + ] + 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": "next", "user_id": "user"}} + + +class _UserMessageWithSearchResults(ChatCompletionUserMessage): + search_results: ReadOnly[list[SearchResultBlock]] + + +class _AssistantMessageWithSearchResults(OpenAIChatCompletionAssistantMessage): + search_results: ReadOnly[list[SearchResultBlock]] + + +_SEARCH_RESULTS: Final[list[SearchResultBlock]] = [ + { + "source": "https://example.com/result", + "title": "Finding", + "content": [{"type": "text", "text": "Search context"}], + "citations": {"enabled": False}, + } +] +_SEARCH_TEXT: Final = 'https://example.com/resultFindingSearch context{"enabled":false}' + + +@pytest.mark.parametrize( + "content, expected", + [ + ([], _SEARCH_TEXT), + ("question", "question" + _SEARCH_TEXT), + ([{"type": "text", "text": "question"}], "question" + _SEARCH_TEXT), + ( + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}], + { + "role": "user", + "parts": [ + {"inline_data": {"data": "AA==", "mime_type": "image/png"}}, + {"text": _SEARCH_TEXT}, + ], + }, + ), + ], +) +def test_search_result_context_is_preserved_with_text_and_media( + content: OpenAIMessageContent, expected: str | ContentType +) -> None: + config: Final = VertexAgentEngineConfig() + message: Final[_UserMessageWithSearchResults] = { + "role": "user", + "content": content, + "search_results": _SEARCH_RESULTS, + } + result: Final = config.transform_request( + model="agent_engine/123", messages=[message], optional_params={"user_id": "user"}, litellm_params={}, headers={} + ) + assert result == {"class_method": "stream_query", "input": {"message": expected, "user_id": "user"}} + + +def test_search_results_are_preserved_without_message_content() -> None: + config: Final = VertexAgentEngineConfig() + message: Final[_AssistantMessageWithSearchResults] = { + "role": "assistant", + "content": None, + "search_results": _SEARCH_RESULTS, + } + result: Final = config.transform_request( + model="agent_engine/123", messages=[message], optional_params={"user_id": "user"}, litellm_params={}, headers={} + ) + assert result == {"class_method": "stream_query", "input": {"message": _SEARCH_TEXT, "user_id": "user"}}