diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 8a9349f38cd..71081e9fd5d 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,13 +6,13 @@ import json import mimetypes import re import xml.etree.ElementTree as ET -from collections.abc import Container, Iterator, Mapping, Sequence +from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from types import MappingProxyType from typing import Any, Final, TypeAlias, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm import litellm.types @@ -1440,6 +1440,45 @@ def convert_to_gemini_tool_call_invoke( raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e}") +def _json_children(node: JsonValue) -> Iterable[JsonValue]: + if isinstance(node, dict): + return node.values() + if isinstance(node, list): + return (child for child in node if not isinstance(child, list)) + return () + + +def _next_json_level(level: tuple[JsonValue, ...], _: None) -> tuple[JsonValue, ...]: + return tuple(itertools.chain.from_iterable(_json_children(node) for node in level)) + + +def _is_part_reference(node: JsonValue) -> bool: + if not isinstance(node, dict): + return False + ref: Final[JsonValue] = node.get("$ref") + return isinstance(ref, str) and ref != "" + + +_REF_CHAR_ESCAPE: Final = re.compile(r"\\u00(?:24|72|65|66)") + + +def _may_spell_ref_key(raw_json: str) -> bool: + """JSON text can only spell the key $ref literally or by \\u-escaping one of its four characters""" + return '"$ref"' in raw_json or _REF_CHAR_ESCAPE.search(raw_json) is not None + + +def _contains_part_reference(payload: JsonValue) -> bool: + """Gemini reads a {"$ref": ""} object in a function_response as a pointer to one of its parts, + unless it sits under a list nested directly in a list + + Google does not document this. It was observed on the Gemini API and Vertex AI on 2026-10-15, see issue #38223 + """ + levels: Final = itertools.takewhile( + bool, itertools.accumulate(itertools.repeat(None), _next_json_level, initial=(payload,)) + ) + return any(_is_part_reference(node) for node in itertools.chain.from_iterable(levels)) + + def convert_to_gemini_tool_call_result( message: ChatCompletionToolMessage | ChatCompletionFunctionMessage, last_message_with_tool_calls: dict | None, @@ -1584,8 +1623,8 @@ def convert_to_gemini_tool_call_result( try: if content_str.strip().startswith("{") or content_str.strip().startswith("["): # Try to parse as JSON (for Computer Use structured responses) - parsed: Final = json.loads(content_str) - if isinstance(parsed, dict): + parsed: Final[JsonValue] = json.loads(content_str) + if isinstance(parsed, dict) and not (_may_spell_ref_key(content_str) and _contains_part_reference(parsed)): response_data = parsed # Use the parsed JSON directly else: response_data = {"content": content_str} diff --git a/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py new file mode 100644 index 00000000000..0194f138689 --- /dev/null +++ b/tests/integration/providers/test_gemini_tool_result_part_reference_chaos.py @@ -0,0 +1,491 @@ +import asyncio +import json +import re +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-2.5-flash" +_API_KEY: Final = "synthetic-gemini-key" +_CONFIG_MODEL: Final = "gemini-tool-result-part-reference-chaos" +_TOOL: Final = "inspect_schema" +_PARAMETERS: Final[Mapping[str, JsonValue]] = { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], +} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_GENERATE: Final = f"/models/{_BACKEND}:generateContent" +_STREAM_GENERATE: Final = f"/models/{_BACKEND}:streamGenerateContent?alt=sse" +_USAGE: Final = {"promptTokenCount": 30, "candidatesTokenCount": 5, "totalTokenCount": 35} + +Endpoint: TypeAlias = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _Call: + endpoint: Endpoint + stream: bool + marker: str + reference: bool + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _tool_text(call: _Call) -> str: + value: Final = {"$ref": "#/$defs/humanScalar"} if call.reference else {"type": "string"} + return json.dumps( + { + "title": f"marker-{call.marker}", + "properties": {"value": value}, + "$defs": {"humanScalar": {"type": "string"}}, + }, + separators=(",", ":"), + ) + + +def _answer(marker: str) -> str: + return f"answer marker-{marker}" + + +def _response_id(marker: str) -> str: + return f"gemini-reply-{marker}" + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body(model: str, call: _Call) -> Mapping[str, JsonValue]: + question: Final = f"Question marker-{call.marker}" + arguments: Final = json.dumps({"name": "Config"}) + common: Final[Mapping[str, JsonValue]] = { + "model": model, + "stream": call.stream, + "num_retries": 0, + "cache": {"no-cache": True}, + } + match call.endpoint: + case "chat": + return { + **common, + "tools": [{"type": "function", "function": {"name": _TOOL, "parameters": dict(_PARAMETERS)}}], + "messages": [ + {"role": "user", "content": question}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_0", "type": "function", "function": {"name": _TOOL, "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_0", "content": _tool_text(call)}, + ], + } + case "messages": + return { + **common, + "max_tokens": 64, + "tools": [{"name": _TOOL, "input_schema": dict(_PARAMETERS)}], + "messages": [ + {"role": "user", "content": question}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_0", "name": _TOOL, "input": {"name": "Config"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": _tool_text(call)}], + }, + ], + } + case "responses": + return { + **common, + "tools": [{"type": "function", "name": _TOOL, "parameters": dict(_PARAMETERS)}], + "input": [ + {"role": "user", "content": question}, + {"type": "function_call", "call_id": "call_0", "name": _TOOL, "arguments": arguments}, + {"type": "function_call_output", "call_id": "call_0", "output": _tool_text(call)}, + ], + } + + +def _frame(marker: str, text: str, finished: bool) -> Mapping[str, JsonValue]: + return { + "responseId": _response_id(marker), + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": text}]}, + "index": 0, + **({"finishReason": "STOP"} if finished else {}), + } + ], + "usageMetadata": dict(_USAGE), + "modelVersion": _BACKEND, + } + + +def _gemini_reply(marker: str, stream: bool, abort_after: int | None = None, pause: float = 0) -> Reply: + if not stream: + return Reply(body=json.dumps(_frame(marker, _answer(marker), finished=True)).encode()) + frames: Final = (_frame(marker, "answer ", finished=False), _frame(marker, f"marker-{marker}", finished=True)) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in frames), + abort_after=abort_after, + pause_between_chunks=pause, + ) + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _function_response(request: Request) -> Mapping[str, JsonValue]: + contents: Final = _JSON_OBJECT.validate_json(request.body)["contents"] + assert isinstance(contents, list), contents + last: Final = contents[-1] + assert isinstance(last, dict) and isinstance(last["parts"], list), last + (part,) = last["parts"] + assert isinstance(part, dict), part + found: Final = part["function_response"] + assert isinstance(found, dict), part + return found + + +def _items_gemini_walks(items: list[JsonValue]) -> Iterator[JsonValue]: + return (item for item in items if not isinstance(item, list)) + + +def _references(node: JsonValue) -> Sequence[str]: + if isinstance(node, dict): + own: Final = (node["$ref"],) if isinstance(node.get("$ref"), str) and node["$ref"] else () + return (*own, *(ref for child in node.values() for ref in _references(child))) + if isinstance(node, list): + return tuple(ref for child in _items_gemini_walks(node) for ref in _references(child)) + return () + + +def _vendor_rejection(request: Request) -> Reply | None: + unresolved: Final = _references(_function_response(request)["response"]) + if not unresolved: + return None + message: Final = ( + f"The referenced name `{unresolved[0]}` in function_response.response does not match to a display_name in " + "the function_response.parts." + ) + return Reply( + status=400, + body=json.dumps({"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}).encode(), + ) + + +def _echo(request: Request) -> Reply: + assert request.target in (_GENERATE, _STREAM_GENERATE), request.target + rejection: Final = _vendor_rejection(request) + if rejection is not None: + return rejection + return _gemini_reply(_marker_of(request), stream=request.target == _STREAM_GENERATE) + + +def _assert_each_in_its_own_form(received: tuple[Request, ...], calls: Sequence[_Call]) -> None: + by_marker: Final = MappingProxyType({call.marker: call for call in calls}) + forwarded: Final = tuple((_marker_of(request), _function_response(request)["response"]) for request in received) + assert sorted(marker for marker, _ in forwarded) == sorted(by_marker), forwarded + expected: Final = tuple( + (marker, {"content": _tool_text(by_marker[marker])}) + if by_marker[marker].reference + else (marker, json.loads(_tool_text(by_marker[marker]))) + for marker, _ in forwarded + ) + assert forwarded == expected + + +def _spend_rows_by_model(model: str, expected: int) -> frozenset[str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=60, + ) + assert [row["status"] for row in rows] == ["success"] * len(rows), rows + assert len(rows) == expected, rows + return frozenset(str(row["request_id"]) for row in rows) + + +def _spend_status_by_request(request_ids: frozenset[str]) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array + ), + lambda found: len(found) >= len(request_ids), + seconds=60, + ) + assert len(rows) == len(request_ids), rows + return {str(row["request_id"]): row["status"] for row in rows} + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _calls( + count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool], reference: Callable[[int], bool] +) -> tuple[_Call, ...]: + return tuple( + _Call( + endpoint=endpoints[index % len(endpoints)], + stream=stream(index), + marker=uuid.uuid4().hex, + reference=reference(index), + ) + for index in range(count) + ) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +async def test_a_concurrent_mixed_burst_reaches_gemini_with_each_tool_result_in_its_own_form(gateway: Gateway) -> None: + calls: Final = _calls( + 30, ("chat", "messages", "responses"), lambda index: index % 2 == 0, lambda index: (index // 2) % 2 == 0 + ) + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_each_in_its_own_form(wire.drain(), calls) + landed: Final = _spend_rows_by_model(model, 30) + assert landed >= { + _response_id(call.marker) for call in calls if not (call.endpoint == "messages" and call.stream) + } + + +async def test_upstream_stream_drops_reach_callers_and_a_later_ref_tool_result_still_reaches_gemini( + gateway: Gateway, +) -> None: + calls: Final = _calls(12, ("chat", "messages", "responses"), lambda _: True, lambda _: True) + dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0) + + def respond(request: Request) -> Reply: + rejection: Final = _vendor_rejection(request) + if rejection is not None: + return rejection + marker: Final = _marker_of(request) + return _gemini_reply(marker, stream=True, abort_after=0 if marker in dropped else None) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 12 + for item in served: + if item.call.marker in dropped: + assert item.status == 500, item.text + assert "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(endpoint="chat", stream=True, marker=uuid.uuid4().hex, reference=True) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + assert recovered.text.rstrip().endswith("data: [DONE]"), recovered.text + _assert_each_in_its_own_form(wire.drain(), (*calls, recovery)) + + +async def test_a_slow_upstream_completes_every_ref_tool_result_stream_exactly_once(gateway: Gateway) -> None: + calls: Final = _calls(12, ("chat", "messages", "responses"), lambda _: True, lambda _: True) + + def respond(request: Request) -> Reply: + rejection: Final = _vendor_rejection(request) + if rejection is not None: + return rejection + return _gemini_reply(_marker_of(request), stream=True, pause=0.5) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 12 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_each_in_its_own_form(wire.drain(), calls) + _spend_rows_by_model(model, 12) + + +def _chaos_config(wire: Wire, tmp_path: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": f"gemini/{_BACKEND}", "api_base": wire.url, "api_key": _API_KEY}, + } + ] + path: Final = tmp_path / "gemini-tool-result-part-reference-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _non_stream_spend_ids(served: Sequence[_Served]) -> frozenset[str]: + return frozenset(_response_id(item.call.marker) for item in served) + + +@dataclass(frozen=True, slots=True) +class _HeldUpstream: + release: threading.Event + held: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.held.put(_marker_of(request)) + assert self.release.wait(timeout=120), "The burst was never released" + return _echo(request) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_ref_tool_results_as_raw_text( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20, ("chat", "messages", "responses"), lambda _: False, lambda _: True) + upstream: Final = _HeldUpstream(threading.Event(), SimpleQueue()) + with wire_server(upstream.respond) as wire: + path: Final = _chaos_config(wire, tmp_path) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds() + 60, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, upstream.held.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + upstream.release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(endpoint="messages", stream=False, marker=uuid.uuid4().hex, reference=True) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + received: Final = wire.drain() + assert {(request.method, request.target) for request in received} == {("POST", _GENERATE)}, received + _assert_each_in_its_own_form(received, (*calls, follow_up)) + assert _spend_status_by_request(_non_stream_spend_ids((*served, answered))) == { + spend_id: "success" for spend_id in _non_stream_spend_ids((*served, answered)) + } + + +@pytest.mark.timeout(600) +async def test_proxy_sigterm_mid_burst_finishes_in_flight_ref_tool_results_and_a_relaunch_serves_again( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(9, ("chat", "messages", "responses"), lambda _: False, lambda _: True) + upstream: Final = _HeldUpstream(threading.Event(), SimpleQueue()) + with wire_server(upstream.respond) as wire: + first: Final = tmp_path / "first" + first.mkdir() + with owned_proxy_process(gateway, first, {}, config=_chaos_config(wire, first), workers=2) as owned: + burst: Final = asyncio.create_task( + _burst( + str(owned.gateway.client.base_url), + owned.gateway.key, + _CONFIG_MODEL, + calls, + tolerate_transport_errors=True, + ) + ) + await asyncio.to_thread(eventually, upstream.held.qsize, lambda size: size == len(calls), 120) + owned.begin_shutdown() + upstream.release.set() + served: Final = await burst + exit_code: Final = await asyncio.to_thread( + eventually, owned.process.poll, lambda code: code is not None, 120 + ) + assert exit_code is not None + assert len(served) == len(calls), (len(served), len(calls)) + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_each_in_its_own_form(wire.drain(), calls) + second: Final = tmp_path / "second" + second.mkdir() + with owned_proxy_process(gateway, second, {}, config=_chaos_config(wire, second), workers=2) as relaunched: + follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex, reference=True) + (answered,) = await _burst( + str(relaunched.gateway.client.base_url), relaunched.gateway.key, _CONFIG_MODEL, (follow_up,) + ) + _assert_answered_with_its_own_marker(answered) + _assert_each_in_its_own_form(wire.drain(), (follow_up,)) + assert _spend_status_by_request(_non_stream_spend_ids((*served, answered))) == { + spend_id: "success" for spend_id in _non_stream_spend_ids((*served, answered)) + } diff --git a/tests/integration/providers/test_gemini_tool_result_part_reference_wire.py b/tests/integration/providers/test_gemini_tool_result_part_reference_wire.py new file mode 100644 index 00000000000..9bab90be407 --- /dev/null +++ b/tests/integration/providers/test_gemini_tool_result_part_reference_wire.py @@ -0,0 +1,1216 @@ +import base64 +import itertools +import json +import os +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from queue import SimpleQueue +from typing import Final, Literal, TypeAlias +from urllib.parse import unquote + +import anthropic +import httpx +import openai +import pytest +from _pytest.mark.structures import ParameterSet +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.timing import assert_within_factor_of_control, timed +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.responses import ResponseCompletedEvent +from pydantic import JsonValue, TypeAdapter + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with + +_BACKEND: Final = "gemini-2.5-flash" +_ARGON: Final = "gemini-4-argon" +_API_KEY: Final = "synthetic-gemini-key" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_BUCKET: Final = "scripted-bucket" +_UPLOAD_PREFIX: Final = f"/upload/storage/v1/b/{_BUCKET}/o?uploadType=media&name=" +_SALT: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt") +_QUESTION: Final = "Which type does the value property of Config use?" +_FOLLOW_UP: Final = "Answer in one word." +_ANSWER: Final = "A humanScalar string." +_TOOL: Final = "inspect_schema" +_ARGS: Final[Mapping[str, JsonValue]] = {"name": "Config"} +_PARAMETERS: Final[Mapping[str, JsonValue]] = { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], +} +_FUNCTION: Final[Mapping[str, JsonValue]] = { + "name": _TOOL, + "description": "Return the JSON Schema of a named config", + "parameters": _PARAMETERS, +} +_DECLARATIONS: Final[Mapping[str, JsonValue]] = {"tools": [{"function_declarations": [dict(_FUNCTION)]}]} +_SCHEMA_TEXT: Final = ( + '{"$defs":{"humanScalar":{"type":"string"}},"type":"object","properties":{"value":{"$ref":"#/$defs/humanScalar"}}}' +) +_PLAIN_TEXT: Final = '{"type":"object","properties":{"value":{"type":"string"}}}' +_PNG: Final = "iVBORw0KGgo=" +_SKIP_SIGNATURE: Final = base64.b64encode(b"skip_thought_signature_validator").decode() +_CACHE_BUST: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +_USAGE: Final = {"promptTokenCount": 20, "candidatesTokenCount": 5, "totalTokenCount": 25} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +Provider: TypeAlias = Literal["gemini", "vertex_ai"] +Endpoint: TypeAlias = Literal["chat", "messages", "responses"] +Client: TypeAlias = Literal["sdk_sync", "sdk_async", "httpx"] +_PROVIDERS: Final[tuple[Provider, ...]] = ("gemini", "vertex_ai") +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +_CLIENTS: Final[tuple[Client, ...]] = ("sdk_sync", "sdk_async", "httpx") + + +@dataclass(frozen=True, slots=True) +class _Observed: + response_id: str + answer: str + + +def _model_path(backend: str) -> str: + return f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models/{backend}" + + +def _children(node: JsonValue) -> Sequence[JsonValue]: + if isinstance(node, dict): + return tuple(node.values()) + if isinstance(node, list): + return tuple(child for child in node if not isinstance(child, list)) + return () + + +def _references(node: JsonValue) -> Iterator[str]: + if isinstance(node, dict): + reference: Final = node.get("$ref") + if isinstance(reference, str) and reference: + yield reference + for child in _children(node): + yield from _references(child) + + +def _object(value: JsonValue) -> Mapping[str, JsonValue]: + return value if isinstance(value, dict) else {} + + +def _items(value: JsonValue) -> Sequence[JsonValue]: + return value if isinstance(value, list) else () + + +def _display_name(part: JsonValue) -> JsonValue: + blob: Final = _object(_object(part).get("inline_data") or _object(part).get("inlineData")) + return blob.get("display_name") or blob.get("displayName") + + +def _function_responses(body: Mapping[str, JsonValue]) -> Iterator[Mapping[str, JsonValue]]: + for content in _items(body.get("contents")): + for part in _items(_object(content).get("parts")): + found = _object(_object(part).get("function_response") or _object(part).get("functionResponse")) + if found: + yield found + + +def _unresolved_reference(raw: bytes) -> str | None: + for function_response in _function_responses(json.loads(raw)): + names = frozenset(str(_display_name(part)) for part in _items(function_response.get("parts"))) + for reference in _references(function_response.get("response")): + if reference not in names: + return reference + return None + + +def _rejection(reference: str) -> Reply: + message: Final = ( + f"The referenced name `{reference}` in function_response.response does not match to a display_name in the " + "function_response.parts." + ) + return Reply( + status=400, + body=json.dumps({"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}).encode(), + ) + + +def _frame(response_id: str, text: str, finished: bool) -> Mapping[str, JsonValue]: + return { + "responseId": response_id, + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": text}]}, + "index": 0, + **({"finishReason": "STOP"} if finished else {}), + } + ], + "usageMetadata": dict(_USAGE), + "modelVersion": _BACKEND, + } + + +def _reply(response_id: str, stream: bool) -> Reply: + if not stream: + return Reply(body=json.dumps(_frame(response_id, _ANSWER, finished=True)).encode()) + frames: Final = (_frame(response_id, "A humanScalar ", finished=False), _frame(response_id, "string.", True)) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in frames), + ) + + +def _vendor(response_ids: Sequence[str], stream: bool) -> Callable[[Request], Reply]: + served: Final = itertools.count() + + def respond(request: Request) -> Reply: + unresolved: Final = _unresolved_reference(request.body) + if unresolved is not None: + return _rejection(unresolved) + return _reply(response_ids[next(served)], stream) + + return respond + + +def _register(gateway: Gateway, scenario: Scenario, provider: Provider, wire: Wire, backend: str = _BACKEND) -> str: + if provider == "gemini": + return scenario.model(model=f"gemini/{backend}", api_base=wire.url, api_key=_API_KEY) + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=f"{wire.url}{_model_path(backend)}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _target(provider: Provider, stream: bool, backend: str = _BACKEND) -> str: + prefix: Final = f"/models/{backend}" if provider == "gemini" else _model_path(backend) + return f"{prefix}:streamGenerateContent?alt=sse" if stream else f"{prefix}:generateContent" + + +def _expected_auth(provider: Provider) -> tuple[str, str]: + return ("x-goog-api-key", _API_KEY) if provider == "gemini" else ("authorization", "Bearer scripted-token") + + +def _args(index: int) -> Mapping[str, JsonValue]: + return _ARGS if index == 0 else {"name": f"Config{index}"} + + +def _chat_messages(tool_contents: Sequence[JsonValue]) -> list[JsonValue]: + calls: Final = [ + {"id": f"call_{index}", "type": "function", "function": {"name": _TOOL, "arguments": json.dumps(_args(index))}} + for index in range(len(tool_contents)) + ] + results: Final = [ + {"role": "tool", "tool_call_id": f"call_{index}", "content": content} + for index, content in enumerate(tool_contents) + ] + return [ + {"role": "user", "content": _QUESTION}, + {"role": "assistant", "content": None, "tool_calls": calls}, + *results, + ] + + +def _messages_history(tool_content: JsonValue) -> list[JsonValue]: + return [ + {"role": "user", "content": _QUESTION}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_0", "name": _TOOL, "input": dict(_ARGS)}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": tool_content}]}, + ] + + +def _responses_input(tool_output: str) -> list[JsonValue]: + return [ + {"role": "user", "content": _QUESTION}, + {"type": "function_call", "call_id": "call_0", "name": _TOOL, "arguments": json.dumps(_ARGS)}, + {"type": "function_call_output", "call_id": "call_0", "output": tool_output}, + ] + + +def _extra(endpoint: Endpoint) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return {"tools": [{"type": "function", "function": dict(_FUNCTION)}], **_CACHE_BUST} + case "messages": + return { + "tools": [{"name": _TOOL, "description": _FUNCTION["description"], "input_schema": dict(_PARAMETERS)}], + **_CACHE_BUST, + } + case "responses": + return {"tools": [{"type": "function", **_FUNCTION}], **_CACHE_BUST} + + +def _http_body(endpoint: Endpoint, model: str, tool_text: str, stream: bool) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return {"model": model, "messages": _chat_messages((tool_text,)), "stream": stream, **_extra(endpoint)} + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": _messages_history(tool_text), + "stream": stream, + **_extra(endpoint), + } + case "responses": + return {"model": model, "input": _responses_input(tool_text), "stream": stream, **_extra(endpoint)} + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _openai(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_openai(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _async_anthropic(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) + + +def _single(values: frozenset[str]) -> str: + assert len(values) == 1, values + return next(iter(values)) + + +def _chat_sync(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _openai(gateway) + messages: Final = _chat_messages((text,)) + if not stream: + reply: Final = client.chat.completions.create(model=model, messages=messages, extra_body=_extra("chat")) + return _Observed(reply.id, reply.choices[0].message.content or "") + chunks: Final = list( + client.chat.completions.create(model=model, messages=messages, stream=True, extra_body=_extra("chat")) + ) + return _Observed( + _single(frozenset(chunk.id for chunk in chunks)), + "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices), + ) + + +async def _chat_async(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _async_openai(gateway) + messages: Final = _chat_messages((text,)) + if not stream: + reply: Final = await client.chat.completions.create(model=model, messages=messages, extra_body=_extra("chat")) + return _Observed(reply.id, reply.choices[0].message.content or "") + chunks: Final = [ + chunk + async for chunk in await client.chat.completions.create( + model=model, messages=messages, stream=True, extra_body=_extra("chat") + ) + ] + return _Observed( + _single(frozenset(chunk.id for chunk in chunks)), + "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices), + ) + + +def _messages_sync(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _anthropic(gateway) + history: Final = _messages_history(text) + if not stream: + reply: Final = client.messages.create( + model=model, max_tokens=64, messages=history, extra_body=_extra("messages") + ) + return _Observed(reply.id, "".join(block.text for block in reply.content if block.type == "text")) + with client.messages.stream( + model=model, max_tokens=64, messages=history, extra_body=_extra("messages") + ) as streamed: + final: Final = streamed.get_final_message() + return _Observed(final.id, "".join(block.text for block in final.content if block.type == "text")) + + +async def _messages_async(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _async_anthropic(gateway) + history: Final = _messages_history(text) + if not stream: + reply: Final = await client.messages.create( + model=model, max_tokens=64, messages=history, extra_body=_extra("messages") + ) + return _Observed(reply.id, "".join(block.text for block in reply.content if block.type == "text")) + async with client.messages.stream( + model=model, max_tokens=64, messages=history, extra_body=_extra("messages") + ) as streamed: + final: Final = await streamed.get_final_message() + return _Observed(final.id, "".join(block.text for block in final.content if block.type == "text")) + + +def _responses_sync(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _openai(gateway) + items: Final = _responses_input(text) + if not stream: + reply: Final = client.responses.create(model=model, input=items, extra_body=_extra("responses")) + return _Observed(reply.id, reply.output_text) + events: Final = list(client.responses.create(model=model, input=items, stream=True, extra_body=_extra("responses"))) + completed: Final = events[-1] + assert isinstance(completed, ResponseCompletedEvent), events + return _Observed(completed.response.id, completed.response.output_text) + + +async def _responses_async(gateway: Gateway, model: str, stream: bool, text: str) -> _Observed: + client: Final = _async_openai(gateway) + items: Final = _responses_input(text) + if not stream: + reply: Final = await client.responses.create(model=model, input=items, extra_body=_extra("responses")) + return _Observed(reply.id, reply.output_text) + events: Final = [ + event + async for event in await client.responses.create( + model=model, input=items, stream=True, extra_body=_extra("responses") + ) + ] + completed: Final = events[-1] + assert isinstance(completed, ResponseCompletedEvent), events + return _Observed(completed.response.id, completed.response.output_text) + + +def _sse_payloads(text: str) -> tuple[Mapping[str, JsonValue], ...]: + lines: Final = tuple(line for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]") + return tuple(_JSON_OBJECT.validate_json(line.removeprefix("data: ").encode()) for line in lines) + + +def _output_text(payload: Mapping[str, JsonValue]) -> str: + messages: Final = tuple(item for item in _items(payload["output"]) if _object(item).get("type") == "message") + parts: Final = tuple(itertools.chain.from_iterable(_items(_object(item)["content"]) for item in messages)) + return "".join(str(_object(part)["text"]) for part in parts if _object(part).get("type") == "output_text") + + +def _chat_answer(payload: Mapping[str, JsonValue]) -> str: + return str(_object(_object(_items(payload["choices"])[0])["message"])["content"]) + + +def _http_observed(endpoint: Endpoint, stream: bool, response: httpx.Response) -> _Observed: + assert response.status_code == 200, response.text + if not stream: + payload: Final = _JSON_OBJECT.validate_json(response.content) + match endpoint: + case "chat": + return _Observed(str(payload["id"]), _chat_answer(payload)) + case "messages": + blocks: Final = tuple(_object(block) for block in _items(payload["content"])) + return _Observed(str(payload["id"]), "".join(str(block["text"]) for block in blocks)) + case "responses": + return _Observed(str(payload["id"]), _output_text(payload)) + events: Final = _sse_payloads(response.text) + match endpoint: + case "chat": + assert response.text.rstrip().endswith("data: [DONE]"), response.text + deltas: Final = tuple(_object(_object(_items(event["choices"])[0])["delta"]) for event in events) + return _Observed( + _single(frozenset(str(event["id"]) for event in events)), + "".join(str(delta["content"]) for delta in deltas if delta.get("content")), + ) + case "messages": + assert events[-1]["type"] == "message_stop", events + start: Final = next(event for event in events if event["type"] == "message_start") + texts: Final = tuple(_object(event["delta"]) for event in events if event["type"] == "content_block_delta") + return _Observed( + str(_object(start["message"])["id"]), + "".join(str(delta["text"]) for delta in texts if delta["type"] == "text_delta"), + ) + case "responses": + assert events[-1]["type"] == "response.completed", events + completed: Final = _object(events[-1]["response"]) + return _Observed(str(completed["id"]), _output_text(completed)) + + +def _httpx(gateway: Gateway, endpoint: Endpoint, model: str, stream: bool, text: str) -> _Observed: + response: Final = gateway.request( + "POST", _path(endpoint), _http_body(endpoint, model, text, stream), headers={"anthropic-version": "2023-06-01"} + ) + return _http_observed(endpoint, stream, response) + + +async def _observe( + gateway: Gateway, endpoint: Endpoint, client: Client, model: str, stream: bool, text: str +) -> _Observed: + match (endpoint, client): + case ("chat", "sdk_sync"): + return _chat_sync(gateway, model, stream, text) + case ("chat", "sdk_async"): + return await _chat_async(gateway, model, stream, text) + case ("messages", "sdk_sync"): + return _messages_sync(gateway, model, stream, text) + case ("messages", "sdk_async"): + return await _messages_async(gateway, model, stream, text) + case ("responses", "sdk_sync"): + return _responses_sync(gateway, model, stream, text) + case ("responses", "sdk_async"): + return await _responses_async(gateway, model, stream, text) + case (_, "httpx"): + return _httpx(gateway, endpoint, model, stream, text) + raise AssertionError((endpoint, client)) + + +def _upstream_response_id(endpoint: Endpoint, caller_id: str) -> str: + if endpoint != "responses": + return caller_id + decrypted: Final = decrypt_if_encrypted_with(caller_id.removeprefix("resp_"), _SALT) + assert decrypted is not None, caller_id + issued: Final = decrypted.split(";")[0].split("response_id:")[-1] + return base64.b64decode(issued.removeprefix("resp_")).decode().split(";")[-1].removeprefix("response_id:") + + +def _assert_caller_id(endpoint: Endpoint, stream: bool, caller_id: str, response_id: str) -> None: + if endpoint == "messages" and stream: + assert caller_id.startswith("msg_"), caller_id + return + assert _upstream_response_id(endpoint, caller_id) == response_id, caller_id + + +def _call_type(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "acompletion" + case "messages": + return "anthropic_messages" + case "responses": + return "aresponses" + + +def _spend_rows(request_ids: Sequence[str], expected: int) -> tuple[Mapping[str, JsonValue], ...]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, model_group, call_type, status, prompt_tokens, completion_tokens " + 'FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array + ), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len(rows) == expected, rows + return tuple(rows) + + +def _spend_row(*request_ids: str) -> Mapping[str, JsonValue]: + row: Final = dict(_spend_rows(request_ids, 1)[0]) + return {key: value for key, value in row.items() if key != "request_id"} + + +def _user(text: str) -> Mapping[str, JsonValue]: + return {"role": "user", "parts": [{"text": text}]} + + +def _call_part(call_id: str | None = None, index: int = 0) -> Mapping[str, JsonValue]: + call: Final = {"name": _TOOL, "args": dict(_args(index)), **({"id": call_id} if call_id else {})} + return {"function_call": call, **({"thoughtSignature": _SKIP_SIGNATURE} if call_id else {})} + + +def _result_part(response: JsonValue, call_id: str | None = None, parts: JsonValue = None) -> Mapping[str, JsonValue]: + result: Final = { + "name": _TOOL, + "response": response, + **({"id": call_id} if call_id else {}), + **({"parts": parts} if parts is not None else {}), + } + return {"function_response": result} + + +def _expected_body(endpoint: Endpoint, *responses: JsonValue) -> Mapping[str, JsonValue]: + return { + "contents": [ + _user(_QUESTION), + {"role": "model", "parts": [_call_part(index=index) for index in range(len(responses))]}, + {"role": "user", "parts": [_result_part(response) for response in responses]}, + ], + **_DECLARATIONS, + **({"generationConfig": {"max_output_tokens": 64}} if endpoint == "messages" else {}), + } + + +def _only_request(wire: Wire, provider: Provider, stream: bool, backend: str = _BACKEND) -> Mapping[str, JsonValue]: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", _target(provider, stream, backend)) + ], received + header, value = _expected_auth(provider) + assert received[0].headers[header] == value, dict(received[0].headers) + return _JSON_OBJECT.validate_json(received[0].body) + + +def _matrix() -> tuple[ParameterSet, ...]: + return tuple( + pytest.param( + provider, endpoint, stream, client, id=f"{provider}-{endpoint}-{'stream' if stream else 'sync'}-{client}" + ) + for provider, endpoint, stream, client in itertools.product(_PROVIDERS, _ENDPOINTS, (False, True), _CLIENTS) + ) + + +@pytest.mark.parametrize(("provider", "endpoint", "stream", "client"), _matrix()) +async def test_a_tool_result_holding_a_string_ref_reaches_gemini_as_its_raw_text( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, client: Client +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), stream)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, provider, wire) + observed: Final = await _observe(gateway, endpoint, client, model, stream, _SCHEMA_TEXT) + assert observed.answer == _ANSWER, observed + _assert_caller_id(endpoint, stream, observed.response_id, response_id) + assert _only_request(wire, provider, stream) == _expected_body(endpoint, {"content": _SCHEMA_TEXT}) + assert _spend_row(observed.response_id, response_id) == { + "model_group": model, + "call_type": _call_type(endpoint), + "status": "success", + "prompt_tokens": 20, + "completion_tokens": 5, + } + + +_LIST_OF_LISTS_TEXT: Final = '{"rows":[[{"$ref":"#/users/1"},"alice"]]}' + + +@pytest.mark.parametrize( + ("provider", "endpoint"), + tuple(pytest.param(p, e, id=f"{p}-{e}") for p, e in itertools.product(_PROVIDERS, _ENDPOINTS)), +) +@pytest.mark.parametrize( + "text", + (pytest.param(_PLAIN_TEXT, id="no-ref"), pytest.param(_LIST_OF_LISTS_TEXT, id="ref-under-a-list-of-lists")), +) +def test_a_tool_result_without_a_part_reference_still_reaches_gemini_structured( + gateway: Gateway, provider: Provider, endpoint: Endpoint, text: str +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, provider, wire) + observed: Final = _httpx(gateway, endpoint, model, False, text) + assert observed.answer == _ANSWER, observed + _assert_caller_id(endpoint, False, observed.response_id, response_id) + assert _only_request(wire, provider, False) == _expected_body(endpoint, json.loads(text)) + assert _spend_row(observed.response_id, response_id)["status"] == "success" + + +_CHART_TEXT: Final = '{"chart":{"$ref":"#/$defs/Chart"},"$defs":{"Chart":{"type":"string"}}}' +_INLINE_PNG: Final[list[JsonValue]] = [{"inline_data": {"data": _PNG, "mime_type": "image/png"}}] + + +def _multimodal_body(endpoint: Endpoint, model: str) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return { + "model": model, + "messages": _chat_messages( + ( + [ + {"type": "text", "text": _CHART_TEXT}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_PNG}"}}, + ], + ) + ), + **_extra(endpoint), + } + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": _messages_history( + [ + {"type": "text", "text": _CHART_TEXT}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG}}, + ] + ), + **_extra(endpoint), + } + case "responses": + raise AssertionError(endpoint) + + +@pytest.mark.parametrize( + ("provider", "endpoint"), + ( + pytest.param("gemini", "chat", id="gemini-chat"), + pytest.param("vertex_ai", "chat", id="vertex_ai-chat"), + pytest.param("gemini", "messages", id="gemini-messages"), + ), +) +def test_a_multimodal_tool_result_with_a_ref_keeps_its_image_parts_beside_the_raw_text( + gateway: Gateway, provider: Provider, endpoint: Endpoint +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, provider, wire) + response: Final = gateway.request( + "POST", _path(endpoint), _multimodal_body(endpoint, model), headers={"anthropic-version": "2023-06-01"} + ) + observed: Final = _http_observed(endpoint, False, response) + assert observed == _Observed(response_id, _ANSWER), observed + expected: Final = _expected_body(endpoint, {"content": _CHART_TEXT}) + assert _only_request(wire, provider, False) == { + **expected, + "contents": [ + *_items(expected["contents"])[:2], + {"role": "user", "parts": [_result_part({"content": _CHART_TEXT}, parts=_INLINE_PNG)]}, + ], + } + assert _spend_row(response_id)["status"] == "success" + + +_CACHE_NAME: Final = "cachedContents/synthetic-cache" +_CACHED_POLICY: Final = " ".join(f"policy clause {index} applies" for index in range(600)) +_EPHEMERAL: Final[Mapping[str, JsonValue]] = {"type": "ephemeral"} + + +def _cached_chat_history(tool_text: str) -> list[JsonValue]: + return [ + {"role": "user", "content": [{"type": "text", "text": _CACHED_POLICY, "cache_control": dict(_EPHEMERAL)}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_0", "type": "function", "function": {"name": _TOOL, "arguments": json.dumps(_ARGS)}} + ], + "cache_control": dict(_EPHEMERAL), + }, + { + "role": "tool", + "tool_call_id": "call_0", + "content": [{"type": "text", "text": tool_text, "cache_control": dict(_EPHEMERAL)}], + }, + {"role": "user", "content": [{"type": "text", "text": "Acknowledged.", "cache_control": dict(_EPHEMERAL)}]}, + {"role": "user", "content": _FOLLOW_UP}, + ] + + +def _cached_messages_history(tool_text: str) -> list[JsonValue]: + return [ + {"role": "user", "content": [{"type": "text", "text": _CACHED_POLICY, "cache_control": dict(_EPHEMERAL)}]}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_0", + "name": _TOOL, + "input": dict(_ARGS), + "cache_control": dict(_EPHEMERAL), + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_0", + "content": [{"type": "text", "text": tool_text, "cache_control": dict(_EPHEMERAL)}], + }, + {"type": "text", "text": "Acknowledged.", "cache_control": dict(_EPHEMERAL)}, + ], + }, + {"role": "user", "content": _FOLLOW_UP}, + ] + + +def _cache_split(endpoint: Endpoint) -> tuple[list[JsonValue], list[JsonValue]]: + tool_turns: Final[list[JsonValue]] = [ + {"role": "model", "parts": [_call_part()]}, + {"role": "user", "parts": [_result_part({"content": _SCHEMA_TEXT})]}, + ] + if endpoint == "chat": + return [_user(_CACHED_POLICY), *tool_turns, _user("Acknowledged.")], [_user(_FOLLOW_UP)] + return [_user(_CACHED_POLICY)], [ + *tool_turns, + {"role": "user", "parts": [{"text": "Acknowledged."}, {"text": _FOLLOW_UP}]}, + ] + + +@pytest.mark.parametrize("endpoint", ("chat", "messages")) +def test_a_context_cached_tool_result_with_a_ref_is_stored_as_its_raw_text( + gateway: Gateway, endpoint: Endpoint +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b"{}") + unresolved: Final = _unresolved_reference(request.body) + if unresolved is not None: + return _rejection(unresolved) + if request.target.endswith("cachedContents"): + return Reply(body=json.dumps({"name": _CACHE_NAME, "model": f"models/{_BACKEND}"}).encode()) + return _reply(response_id, False) + + body: Final = ( + {"messages": _cached_chat_history(_SCHEMA_TEXT)} + if endpoint == "chat" + else {"max_tokens": 64, "messages": _cached_messages_history(_SCHEMA_TEXT)} + ) + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + response: Final = gateway.request( + "POST", + _path(endpoint), + {"model": model, **body, **_CACHE_BUST}, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["id"] == response_id, response.text + received: Final = wire.drain() + assert [(request.method, request.target.endswith("cachedContents")) for request in received] == [ + ("GET", True), + ("POST", True), + ("POST", False), + ], received + stored: Final = _JSON_OBJECT.validate_json(received[1].body) + assert isinstance(stored["displayName"], str) and stored["displayName"], stored + cached_turns, live_turns = _cache_split(endpoint) + assert stored == { + "contents": cached_turns, + "model": f"models/{_BACKEND}", + "displayName": stored["displayName"], + "tools": None, + }, stored + assert _JSON_OBJECT.validate_json(received[2].body) == { + "contents": live_turns, + "cachedContent": _CACHE_NAME, + **({"generationConfig": {"max_output_tokens": 64}} if endpoint == "messages" else {}), + } + assert _spend_row(response_id)["status"] == "success" + + +def _gcs_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target.startswith(_UPLOAD_PREFIX), request.target + assert request.headers["authorization"] == "Bearer scripted-token" + name: Final = unquote(request.target.removeprefix(_UPLOAD_PREFIX)) + stored: Final = { + "kind": "storage#object", + "id": f"{_BUCKET}/{name}/1759950000000000", + "name": name, + "bucket": _BUCKET, + "size": str(len(request.body)), + "timeCreated": "2026-10-15T08:00:00.000Z", + "contentType": "application/octet-stream", + } + return Reply(body=json.dumps(stored).encode()) + + +def _batch_line(url: str, model: str, tool_text: str) -> Mapping[str, JsonValue]: + body: Final = ( + {"model": model, "messages": _chat_messages((tool_text,)), "tools": _extra("chat")["tools"]} + if url == "/v1/chat/completions" + else {"model": model, "input": _responses_input(tool_text), "tools": _extra("responses")["tools"]} + ) + return {"custom_id": "schema-row", "method": "POST", "url": url, "body": body} + + +@pytest.mark.parametrize("url", ("/v1/chat/completions", "/v1/responses")) +def test_a_vertex_batch_line_with_a_ref_tool_result_is_stored_as_its_raw_text(gateway: Gateway, url: str) -> None: + with wire_server(_gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_key=None, + api_base=wire.url, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url), + gcs_bucket_name=_BUCKET, + ) + lines: Final = (_batch_line(url, model, _SCHEMA_TEXT), _batch_line(url, model, _PLAIN_TEXT)) + response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model}, + { + "file": ( + "in.jsonl", + ("\n".join(json.dumps(line) for line in lines) + "\n").encode(), + "application/jsonl", + ) + }, + ) + assert response.status_code == 200, response.text + assert response.json()["purpose"] == "batch", response.text + uploads: Final = wire.drain() + assert len(uploads) == 1, [upload.target for upload in uploads] + stored: Final = tuple( + _object(_JSON_OBJECT.validate_json(line)["request"]) + for line in uploads[0].body.splitlines() + if line.strip() + ) + assert [request["contents"] for request in stored] == [ + _expected_body("chat", {"content": _SCHEMA_TEXT})["contents"], + _expected_body("chat", json.loads(_PLAIN_TEXT))["contents"], + ] + assert [request["tools"] for request in stored] == [_DECLARATIONS["tools"]] * 2 + assert [sorted(request) for request in stored] == [["contents", "labels", "tools"]] * 2 + + +_NATIVE_CONTENTS: Final[list[JsonValue]] = [ + {"role": "user", "parts": [{"text": "Describe the chart."}]}, + {"role": "model", "parts": [{"functionCall": {"name": "render_chart", "args": {}}}]}, + { + "role": "user", + "parts": [ + { + "functionResponse": { + "name": "render_chart", + "response": {"image": {"$ref": "chart.png"}}, + "parts": [{"inlineData": {"mimeType": "image/png", "data": _PNG, "displayName": "chart.png"}}], + } + } + ], + }, +] + + +def test_a_native_generate_content_part_reference_is_forwarded_verbatim(gateway: Gateway) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + response: Final = gateway.request( + "POST", f"/v1beta/models/{model}:generateContent", {"contents": _NATIVE_CONTENTS} + ) + assert response.status_code == 200, response.text + candidate: Final = _object(_items(_JSON_OBJECT.validate_json(response.content)["candidates"])[0]) + assert _object(candidate["content"])["parts"] == [{"text": _ANSWER}], response.text + assert _only_request(wire, "gemini", False) == { + "model": _BACKEND, + "contents": _NATIVE_CONTENTS, + "tools": None, + "generationConfig": {}, + } + + +def test_a_gemini_3_tool_result_with_a_ref_keeps_its_call_id_beside_the_raw_text(gateway: Gateway) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire, _ARGON) + observed: Final = _httpx(gateway, "chat", model, False, _SCHEMA_TEXT) + assert observed == _Observed(response_id, _ANSWER), observed + assert _only_request(wire, "gemini", False, _ARGON) == { + "contents": [ + _user(_QUESTION), + {"role": "model", "parts": [_call_part("call_0")]}, + {"role": "user", "parts": [_result_part({"content": _SCHEMA_TEXT}, call_id="call_0")]}, + ], + **_DECLARATIONS, + } + assert _spend_row(response_id)["status"] == "success" + + +def _send_chat(gateway: Gateway, model: str, tool_contents: Sequence[JsonValue]) -> httpx.Response: + return gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": _chat_messages(tool_contents), **_extra("chat")} + ) + + +def _assert_answered(response: httpx.Response, response_id: str) -> None: + assert response.status_code == 200, response.text + assert _http_observed("chat", False, response) == _Observed(response_id, _ANSWER), response.text + + +_FIVE_KB_TEXT: Final = json.dumps( + {"padding": "p" * 5000, "properties": {"value": {"$ref": "#/$defs/humanScalar"}}}, separators=(",", ":") +) +_SPLIT_PARTS: Final[list[JsonValue]] = [ + {"type": "text", "text": '{"value":{"$r'}, + {"type": "text", "text": 'ef":"#/$defs/humanScalar"}}'}, +] + + +@pytest.mark.parametrize( + ("tool_contents", "expected"), + ( + pytest.param((5,), ({"content": ""},), id="content-int"), + pytest.param(("",), ({"content": ""},), id="content-empty-string"), + pytest.param( + (_SPLIT_PARTS,), + ({"content": '{"value":{"$ref":"#/$defs/humanScalar"}}'},), + id="ref-key-split-across-text-parts", + ), + pytest.param((_FIVE_KB_TEXT,), ({"content": _FIVE_KB_TEXT},), id="5kb-ref-at-the-end"), + pytest.param( + (_SCHEMA_TEXT, _PLAIN_TEXT), + ({"content": _SCHEMA_TEXT}, json.loads(_PLAIN_TEXT)), + id="two-results-one-ref", + ), + pytest.param( + (_SCHEMA_TEXT, _SCHEMA_TEXT), + ({"content": _SCHEMA_TEXT}, {"content": _SCHEMA_TEXT}), + id="two-results-both-ref", + ), + pytest.param(('{"v":{"$ref":"#/a","$ref":null}}',), ({"v": {"$ref": None}},), id="ref-twice-string-then-null"), + pytest.param( + ('{"v":{"$ref":null,"$ref":"#/a"}}',), + ({"content": '{"v":{"$ref":null,"$ref":"#/a"}}'},), + id="ref-twice-null-then-string", + ), + pytest.param(('{"v":{"$ref":"#/a"',), ({"content": '{"v":{"$ref":"#/a"'},), id="malformed-json"), + pytest.param(('[{"$ref":"#/a"}]',), ({"content": '[{"$ref":"#/a"}]'},), id="array-root"), + ), +) +def test_sad_tool_result_shapes_reach_gemini_in_the_form_it_accepts( + gateway: Gateway, tool_contents: tuple[JsonValue, ...], expected: tuple[JsonValue, ...] +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + _assert_answered(_send_chat(gateway, model, tool_contents), response_id) + assert _only_request(wire, "gemini", False) == _expected_body("chat", *expected) + assert _spend_row(response_id)["status"] == "success" + + +_UPSTREAM_ERRORS: Final = ( + pytest.param(401, "UNAUTHENTICATED", "Request had invalid authentication credentials.", id="401"), + pytest.param(404, "NOT_FOUND", f"models/{_BACKEND} is not found for API version v1beta", id="404"), + pytest.param( + 400, + "INVALID_ARGUMENT", + "The input token count (2000000) exceeds the maximum number of tokens allowed (1048576).", + id="400-oversized", + ), +) + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +@pytest.mark.parametrize(("status", "reason", "message"), _UPSTREAM_ERRORS) +def test_an_upstream_error_on_a_ref_tool_result_reaches_the_caller_and_the_next_request_succeeds( + gateway: Gateway, endpoint: Endpoint, status: int, reason: str, message: str +) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + served: Final = itertools.count() + vendor: Final = _vendor((response_id,), False) + + def respond(request: Request) -> Reply: + if next(served) == 0: + error: Final = {"error": {"code": status, "message": message, "status": reason}} + return Reply(status=status, body=json.dumps(error).encode()) + return vendor(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + failed: Final = gateway.request( + "POST", + _path(endpoint), + _http_body(endpoint, model, _SCHEMA_TEXT, False), + headers={"anthropic-version": "2023-06-01"}, + ) + assert failed.status_code == status, failed.text + assert message in failed.text, failed.text + assert _only_request(wire, "gemini", False) == _expected_body(endpoint, {"content": _SCHEMA_TEXT}) + recovered: Final = _httpx(gateway, endpoint, model, False, _SCHEMA_TEXT) + assert recovered.answer == _ANSWER, recovered + _assert_caller_id(endpoint, False, recovered.response_id, response_id) + assert _only_request(wire, "gemini", False) == _expected_body(endpoint, {"content": _SCHEMA_TEXT}) + + +def test_a_ref_tool_result_without_a_matching_call_fails_before_reaching_gemini(gateway: Gateway) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + orphan: Final = [ + *_chat_messages((_SCHEMA_TEXT,))[:2], + {"role": "tool", "tool_call_id": "call_unknown", "content": _SCHEMA_TEXT}, + ] + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": orphan, **_extra("chat")} + ) + assert response.status_code == 500, response.text + assert "Missing corresponding tool call for tool response message" in response.text, response.text + assert wire.drain() == () + _assert_answered(_send_chat(gateway, model, (_SCHEMA_TEXT,)), response_id) + + +def _escaped(key: str, index: int) -> str: + return key[:index] + f"\\u{ord(key[index]):04x}" + key[index + 1 :] + + +@pytest.mark.parametrize( + "text", + ( + pytest.param('{"v":{"$ref":null}}', id="ref-null"), + pytest.param('{"v":{"$ref":""}}', id="ref-empty-string"), + pytest.param('{"v":{"$ref":5}}', id="ref-int"), + pytest.param('{"v":{"$ref":true}}', id="ref-bool"), + pytest.param('{"v":{"$ref":{"a":1}}}', id="ref-object"), + pytest.param('{"v":{"$ref":["#/a"]}}', id="ref-list"), + pytest.param('{"note":"$ref","tags":["$ref"]}', id="ref-as-a-value"), + pytest.param(json.dumps({"text": json.dumps({"$ref": "#/a"})}), id="ref-inside-a-json-string"), + pytest.param('{"price":"\\u0024 5"}', id="escaped-dollar-in-a-value"), + pytest.param("{}", id="empty-object"), + pytest.param(_LIST_OF_LISTS_TEXT, id="ref-under-a-list-of-lists"), + pytest.param('{"rows":[[{"cell":{"$ref":"#/users/1"}}]]}', id="ref-in-an-object-under-a-list-of-lists"), + pytest.param('{"rows":[[{"cell":[{"$ref":"#/users/1"}]}]]}', id="ref-in-a-list-under-a-list-of-lists"), + pytest.param('{"grid":[[[{"$ref":"#/users/1"}]]]}', id="ref-under-three-lists"), + pytest.param('{"a":[{"b":[[{"$ref":"#/users/1"}]]}]}', id="list-of-lists-under-an-object-in-a-list"), + ), +) +def test_edge_tool_results_gemini_accepts_stay_structured(gateway: Gateway, text: str) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + _assert_answered(_send_chat(gateway, model, (text,)), response_id) + assert _only_request(wire, "gemini", False) == _expected_body("chat", json.loads(text)) + + +def _deep(levels: int, leaf: JsonValue) -> str: + return '{"n":' * levels + json.dumps(leaf, separators=(",", ":")) + "}" * levels + + +@pytest.mark.parametrize( + "text", + ( + *( + pytest.param('{"v":{"' + _escaped("$ref", index) + '":"#/a"}}', id=f"escaped-key-{index}") + for index in range(4) + ), + pytest.param('{"$ref":"#/$defs/Config","$defs":{"Config":{"type":"object"}}}', id="ref-on-the-root"), + pytest.param(_deep(200, {"$ref": "#/deep"}), id="ref-200-levels-deep"), + pytest.param('{"items":[{"a":1},{"b":[{"$ref":"#/b"}]}]}', id="ref-inside-a-list"), + pytest.param('{"rows":[{"$ref":"#/users/1"},"alice"]}', id="ref-in-a-one-level-list"), + pytest.param( + '{"rows":[[{"$ref":"#/users/1"}]],"owner":{"$ref":"#/users/2"}}', + id="list-of-lists-ref-beside-a-reachable-ref", + ), + ), +) +def test_edge_tool_results_with_a_string_ref_reach_gemini_as_raw_text(gateway: Gateway, text: str) -> None: + response_id: Final = f"gemini-reply-{uuid.uuid4().hex}" + with wire_server(_vendor((response_id,), False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + _assert_answered(_send_chat(gateway, model, (text,)), response_id) + assert _only_request(wire, "gemini", False) == _expected_body("chat", {"content": text}) + + +def test_the_same_ref_tool_result_sent_twice_is_wrapped_and_logged_once_each(gateway: Gateway) -> None: + response_ids: Final = (f"gemini-reply-{uuid.uuid4().hex}", f"gemini-reply-{uuid.uuid4().hex}") + with wire_server(_vendor(response_ids, False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + answers: Final = tuple(_send_chat(gateway, model, (_SCHEMA_TEXT,)) for _ in response_ids) + for answer, response_id in zip(answers, response_ids, strict=True): + _assert_answered(answer, response_id) + received: Final = wire.drain() + assert [_JSON_OBJECT.validate_json(request.body) for request in received] == [ + _expected_body("chat", {"content": _SCHEMA_TEXT}) + ] * 2 + rows: Final = _spend_rows(response_ids, 2) + assert sorted(str(row["request_id"]) for row in rows) == sorted(response_ids) + assert [row["status"] for row in rows] == ["success", "success"] + + +_BODY_BYTES: Final = 8 * 1024 * 1024 + + +def _escape_heavy(with_reference: bool) -> str: + count: Final = _BODY_BYTES // 6 + tail: Final = '"\\u0024ref":"#/a"' if with_reference else '"tail":"#/a"' + filler: Final = "\\u0000" * count if with_reference else "a" * (6 * count) + return '{"blob":"' + filler + '","end":{' + tail + "}}" + + +def _node_heavy(with_reference: bool) -> str: + count: Final = _BODY_BYTES // 3 + tail: Final = '{"$ref":"#/a"}' if with_reference else '{"tail":"#/a"}' + return '{"items":[' + "{}," * count + tail + "]}" + + +def _list_heavy(with_reference: bool) -> str: + count: Final = _BODY_BYTES // 3 + tail: Final = '{"$ref":"#/a"}' if with_reference else '{"tail":"#/a"}' + return '{"items":[' + "[]," * count + tail + "]}" + + +def _send_large_chat(gateway: Gateway, model: str, tool_text: str) -> httpx.Response: + return gateway.client.post( + "/v1/chat/completions", + json={"model": model, "messages": _chat_messages((tool_text,)), **_extra("chat")}, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=120, + ) + + +def _timed_liveliness(gateway: Gateway) -> tuple[float, int]: + started: Final = time.perf_counter() + probe: Final = gateway.client.get("/health/liveliness") + return time.perf_counter() - started, probe.status_code + + +@pytest.mark.timeout(600) +@pytest.mark.parametrize( + "build", + ( + pytest.param(_escape_heavy, id="escapes"), + pytest.param(_node_heavy, id="nodes"), + pytest.param(_list_heavy, id="lists"), + ), +) +def test_the_worst_case_ref_scan_stays_within_ten_times_a_same_size_control( + gateway: Gateway, build: Callable[[bool], str], record_property: Callable[[str, object], None] +) -> None: + response_ids: Final = (f"gemini-reply-{uuid.uuid4().hex}", f"gemini-reply-{uuid.uuid4().hex}") + worst_text: Final = build(True) + samples: Final[SimpleQueue[tuple[float, int]]] = SimpleQueue() + stop: Final = threading.Event() + + def poll() -> None: + while not stop.is_set(): + samples.put(_timed_liveliness(gateway)) + + with wire_server(_vendor(response_ids, False)) as wire, gateway.scenario() as scenario: + model: Final = _register(gateway, scenario, "gemini", wire) + control: Final = timed(lambda: _send_large_chat(gateway, model, build(False))) + poller: Final = threading.Thread(target=poll) + poller.start() + worst: Final = timed(lambda: _send_large_chat(gateway, model, worst_text)) + stop.set() + poller.join() + liveliness: Final = tuple(samples.get_nowait() for _ in range(samples.qsize())) + record_property("body_bytes", len(worst_text)) + record_property("control_seconds", round(control.seconds, 3)) + record_property("worst_seconds", round(worst.seconds, 3)) + record_property("liveliness_max_seconds", round(max(latency for latency, _ in liveliness), 3)) + _assert_answered(control.value, response_ids[0]) + _assert_answered(worst.value, response_ids[1]) + assert_within_factor_of_control(worst.seconds, control=control.seconds, factor=10) + assert {status for _, status in liveliness} == {200}, liveliness + assert_within_factor_of_control(max(latency for latency, _ in liveliness), control=control.seconds, factor=10) + received: Final = wire.drain() + assert len(received) == 2, [request.target for request in received] + assert _function_response_of(received[1]) == {"content": worst_text} + + +def _function_response_of(request: Request) -> JsonValue: + responses: Final = tuple(_function_responses(_JSON_OBJECT.validate_json(request.body))) + assert len(responses) == 1, len(responses) + return responses[0]["response"] diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 71b3b84bfd0..67f1108eb5e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1373,6 +1373,167 @@ def test_convert_gemini_tool_call_result_with_data_url_extra_params(): ) +def _gemini_tool_result_fixture( + content: str, +) -> tuple[ChatCompletionToolMessage, dict]: + message = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_schema", + content=content, + ) + last_message_with_tool_calls = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_schema", + "type": "function", + "index": 0, + "function": {"name": "inspect_schema", "arguments": "{}"}, + } + ], + } + return message, last_message_with_tool_calls + + +def test_convert_gemini_tool_result_wraps_output_containing_json_schema_ref(): + """ + Gemini treats {"$ref": ...} objects inside function_response.response as + references to named parts in function_response.parts, so schema-bearing + tool output must be delivered opaquely instead of structurally. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + schema_json = ( + '{"$defs":{"humanScalar":{"type":"string"}},' + '"type":"object","properties":{"value":{"$ref":"#/$defs/humanScalar"}}}' + ) + message, last_message_with_tool_calls = _gemini_tool_result_fixture(schema_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["name"] == "inspect_schema" + assert function_response["response"] == {"content": schema_json} + + +def test_convert_gemini_tool_result_wraps_deeply_nested_json_schema_ref(): + """ + A $ref key at any depth must trigger opaque delivery. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + nested_json = '{"results":[{"schemas":[{"node":{"$ref":"#/$defs/deep"}}]}]}' + message, last_message_with_tool_calls = _gemini_tool_result_fixture(nested_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["response"] == {"content": nested_json} + + +def test_convert_gemini_tool_result_wraps_non_schema_string_ref(): + """Gemini reads any non-empty string $ref as a part name, a MongoDB DBRef included.""" + dbref_json = '{"owner":{"$ref":"users","$id":{"$oid":"64b7f0c2a1"}}}' + message, last_message_with_tool_calls = _gemini_tool_result_fixture(dbref_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + assert result["function_response"]["response"] == {"content": dbref_json} + + +def test_convert_gemini_tool_result_wraps_top_level_ref(): + """A $ref on the tool result's root object is read as a part name too.""" + root_ref_json = '{"$ref":"#/$defs/Config","$defs":{"Config":{"type":"object"}}}' + message, last_message_with_tool_calls = _gemini_tool_result_fixture(root_ref_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + assert result["function_response"]["response"] == {"content": root_ref_json} + + +def _json_unicode_escape(text: str) -> str: + return "".join(f"\\u{ord(char):04x}" for char in text) + + +@pytest.mark.parametrize( + "ref_key", + ["$ref"[:index] + _json_unicode_escape("$ref"[index]) + "$ref"[index + 1 :] for index in range(4)], +) +def test_convert_gemini_tool_result_wraps_unicode_escaped_ref_key(ref_key: str): + """A $ref key spelled with JSON unicode escapes still decodes to $ref, so it must be wrapped too.""" + escaped_json = '{"v":{"' + ref_key + '":"#/a"}}' + message, last_message_with_tool_calls = _gemini_tool_result_fixture(escaped_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + assert result["function_response"]["response"] == {"content": escaped_json} + + +@pytest.mark.parametrize( + "schema_json", + [ + '{"$defs":{"Foo":{"type":"string"}},"type":"object","properties":{"bar":{"type":"string"}}}', + '{"definitions":{"Foo":{"type":"string"}},"type":"object","properties":{"bar":{"type":"string"}}}', + '{"status":"ok","count":2,"tags":["a","b"]}', + '{"v":{"$ref":null}}', + '{"v":{"$ref":""}}', + '{"v":{"$ref":5}}', + '{"v":{"$ref":true}}', + '{"v":{"$ref":{"a":1}}}', + '{"v":{"$ref":["#/a"]}}', + '{"columns":["user","name"],"rows":[[{"$ref":"#/users/1"},"alice"]]}', + '{"a":[1,[{"b":{"$ref":"x"}}]]}', + ], +) +def test_convert_gemini_tool_result_preserves_structured_passthrough_without_part_references( + schema_json: str, +): + """ + Gemini only rejects a $ref holding a non-empty string outside a list nested + directly in a list, so tool output with no $ref, a $ref holding any other + value, or a $ref under a list of lists stays structured. + Google does not document that rule. It was observed on the Gemini API and + Vertex AI (gemini-3.8-flash, gemini-4-argon) on 2026-10-15. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + message, last_message_with_tool_calls = _gemini_tool_result_fixture(schema_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["name"] == "inspect_schema" + assert function_response["response"] == json.loads(schema_json) + + +def test_convert_gemini_tool_result_malformed_json_still_wrapped_as_content(): + """Unparseable tool output keeps taking the existing content-wrapping path.""" + message, last_message_with_tool_calls = _gemini_tool_result_fixture("{not valid json") + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + assert result["function_response"]["response"] == {"content": "{not valid json"} + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 047af5d0ab2..b97724319c7 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -415,6 +415,52 @@ def test_empty_content_handling(): assert contents[0]["parts"][0]["text"] == "" +def test_transform_request_body_wraps_tool_result_with_json_schema_ref(): + """ + Regression: Gemini treats {"$ref": ...} objects inside functionResponse.response + as references to named parts in function_response.parts and rejects the request. + The shared request-body transformation must deliver schema-bearing tool output + opaquely instead of structurally. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + schema_json = ( + '{"$defs":{"humanScalar":{"type":"string"}},' + '"type":"object","properties":{"value":{"$ref":"#/$defs/humanScalar"}}}' + ) + messages = [ + {"role": "user", "content": "Inspect this schema."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_schema", + "type": "function", + "function": {"name": "inspect_schema", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_schema", "content": schema_json}, + {"role": "user", "content": "Acknowledge the result."}, + ] + + result = transform_request_body( + messages=messages, + model="gemini-2.0-flash", + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + parts = [part for content in result["contents"] for part in content.get("parts", [])] + function_responses = [ + part[key] for part in parts for key in ("function_response", "functionResponse") if key in part + ] + assert len(function_responses) == 1 + assert function_responses[0]["response"] == {"content": schema_json} + + def test_thought_signature_extraction_from_response(): """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (