mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(vertex_ai): forward agent engine attachments
This commit is contained in:
AkilanChithra
parent
c80e6e2474
commit
a8fab3990f
2 changed files with 353 additions and 17 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue