diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 0339533e1f5..8d0c724d106 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -9,18 +9,22 @@ API Reference: - :streamQuery endpoint - for actual queries (stream_query method) """ +import base64 +import binascii import json -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Optional, Union +import mimetypes +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, Union +from urllib.parse import urlparse import httpx -from pydantic import TypeAdapter +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.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.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( @@ -28,7 +32,13 @@ from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( ) 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: @@ -60,6 +70,124 @@ 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: 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") + search_text: Final = extract_search_results_text(message.get("search_results")) + if content is None: + return search_text + if isinstance(content, str): + return content + search_text + 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 = tuple(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) + search_text + if message["role"] != "user": + return _AgentEngineInputError("Agent Engine media must be in the final user message") + return {"role": "user", "parts": [*parts, {"text": search_text}] if search_text else list(parts)} + + +def _messages_to_agent_input(messages: Sequence[AllMessageValues]) -> str | ContentType | _AgentEngineInputError: + if not messages: + return _AgentEngineInputError("Agent Engine requires at least one message") + return _message_to_agent_input(messages[-1]) + + class VertexAgentEngineConfig(BaseConfig, VertexBase): """ Configuration for Vertex AI Agent Engine (Reasoning Engines). @@ -205,22 +333,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 20333fd0465..07dc3adea08 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -4,21 +4,32 @@ Tests for Vertex AI Agent Engine transformation. Tests the request transformation and streaming chunk parsing without making real API calls. """ - import json +from copy import deepcopy from datetime import datetime from typing import Final import httpx import pytest +from typing_extensions import ReadOnly - +from litellm import BadRequestError +from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, get_async_httpx_client 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.bedrock import SearchResultBlock +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionUserMessage, + OpenAIChatCompletionAssistantMessage, + OpenAIMessageContent, + OpenAIMessageContentListBlock, +) +from litellm.types.llms.vertex_ai import ContentType from litellm.types.utils import LlmProviders from litellm.utils import CustomStreamWrapper @@ -103,10 +114,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 @@ -136,6 +144,309 @@ class TestVertexAgentEngineChunkParser: 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": "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"}}, + {"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 + + +def test_media_in_final_assistant_message_is_rejected() -> None: + config: Final = VertexAgentEngineConfig() + 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", + messages=messages, + optional_params={"user_id": "user"}, + litellm_params={}, + 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: + 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"}} + + +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"}} + + +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=[], optional_params={"user_id": "user"}, litellm_params={}, headers={} + ) + assert error.value.status_code == 400 + 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"}} + + async def test_async_stream_wrapper_without_a_client_posts_through_the_cached_vertex_ai_client(): api_base: Final = "https://us-central1-aiplatform.googleapis.com/v1/reasoningEngines/123:streamQuery" sent_requests: Final[list[httpx.Request]] = []