mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge 87c6a47f85 into 527d882508
This commit is contained in:
Akilan Chithra Sathish
GitHub
commit
4c083994ed
2 changed files with 458 additions and 20 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]] = []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue