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"}}