litellm/tests/unit/interactions/test_main.py
zoveki 41fcce62ff
fix(gemini)!: accept Google's current Managed Agents fields on create and interactions (#48750)
* fix(gemini): send Google's current Managed Agents fields on create and forward agent_config on interactions

* fix(gemini): accept id on managed agent get, delete, and list_versions

* fix(gemini): drop unsupported agent fields under drop_params and keep agent_config keyword-only

A model interaction drops agent_config under drop_params, agent create drops wrongly typed optional fields under drop_params and otherwise answers 400 naming the field the caller sent, and agent_config on interactions create is keyword-only so positional response settings stay in place

* fix(gemini): drop a non-object agent_config under drop_params and name the id alias the caller sent

* fix(interactions): reject agent_config on bridged providers and drop a typeless one under drop_params

Bridged providers answer 400 for agent_config without drop_params and drop it with drop_params. A Gemini agent run drops an agent_config without a string type under drop_params. The Gemini agents wire test expects Google's current upstream field names.

* fix(interactions): honor additional_drop_params for agent_config on bridged providers

* fix(agents): drop a clashing instructions alias under drop_params

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-28 12:37:59 -07:00

858 lines
33 KiB
Python

import base64
import json
from collections.abc import Mapping
from datetime import datetime
from typing import Final
import httpx
import pytest
from pydantic import JsonValue, TypeAdapter
from respx import MockRouter
import litellm
from litellm import interactions
from litellm.integrations.custom_logger import CustomLogger
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _openai_response_body(response_id: str, text: str) -> dict[str, object]:
return {
"id": response_id,
"object": "response",
"created_at": 1700000000,
"status": "completed",
"error": None,
"incomplete_details": None,
"instructions": None,
"max_output_tokens": None,
"model": "gpt-5.6",
"output": [
{
"id": f"msg-{response_id}",
"type": "message",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
],
"parallel_tool_calls": True,
"previous_response_id": None,
"reasoning": {"effort": None, "summary": None},
"store": False,
"temperature": 1.0,
"text": {"format": {"type": "text"}},
"tool_choice": "auto",
"tools": [],
"top_p": 1.0,
"truncation": "disabled",
"usage": {
"input_tokens": 6,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens": 1,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 7,
},
"user": None,
"metadata": {},
}
def _bridge_interaction_id(response_id: str) -> str:
encoded: Final = f"litellm:custom_llm_provider:openai;model_id:None;response_id:{response_id}".encode()
return f"resp_{base64.b64encode(encoded).decode()}"
def _sse_event(event_type: str, sequence_number: int, payload: dict[str, object]) -> str:
data: Final = json.dumps({"type": event_type, "sequence_number": sequence_number, **payload})
return f"event: {event_type}\ndata: {data}\n\n"
def _responses_stream_body() -> str:
completed: Final = _openai_response_body("resp-stream", "Hello")
in_progress: Final = {**completed, "status": "in_progress", "output": [], "usage": None}
text_location: Final = {"item_id": "msg-resp-stream", "output_index": 0, "content_index": 0}
empty_part: Final = {"type": "output_text", "text": "", "annotations": []}
item: Final = {"id": "msg-resp-stream", "type": "message", "role": "assistant"}
events: Final = (
("response.created", {"response": in_progress}),
("response.in_progress", {"response": in_progress}),
("response.output_item.added", {"output_index": 0, "item": {**item, "status": "in_progress", "content": []}}),
("response.content_part.added", {**text_location, "part": empty_part}),
("response.output_text.delta", {**text_location, "delta": "Hello"}),
("response.output_text.done", {**text_location, "text": "Hello"}),
("response.content_part.done", {**text_location, "part": {**empty_part, "text": "Hello"}}),
(
"response.output_item.done",
{"output_index": 0, "item": {**item, "status": "completed", "content": [{**empty_part, "text": "Hello"}]}},
),
("response.completed", {"response": completed}),
)
return "".join(_sse_event(event_type, index, payload) for index, (event_type, payload) in enumerate(events))
@pytest.fixture
def api_key():
return "test-api-key"
class TestGoogleInteractionsCreate:
@pytest.mark.usefixtures("fake_provider_credentials")
def test_missing_model_and_agent(self, api_key):
"""Test error when neither model nor agent is provided."""
with pytest.raises((ValueError, litellm.APIConnectionError)):
interactions.create(
input="Hello",
api_key=api_key,
)
@pytest.mark.parametrize(
("model", "api_key", "url", "expected_body", "response_body", "expected_response"),
[
pytest.param(
"gemini/gemini-2.5-flash",
"gemini-offline",
"https://generativelanguage.googleapis.com/v1beta/interactions",
{"model": "gemini-2.5-flash", "input": "Hello, how are you?"},
{
"id": "interaction-gemini",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"created": "2026-05-01T00:00:00Z",
"updated": "2026-05-01T00:00:00Z",
"outputs": [{"type": "text", "text": "4"}],
"usage": {"total_input_tokens": 6, "total_output_tokens": 1},
},
{
"id": "interaction-gemini",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"created": "2026-05-01T00:00:00Z",
"updated": "2026-05-01T00:00:00Z",
"outputs": [{"type": "text", "text": "4"}],
"usage": {"total_input_tokens": 6, "total_output_tokens": 1},
},
id="gemini",
),
pytest.param(
"gpt-5.6",
"sk-offline",
"https://api.openai.com/v1/responses",
{"model": "gpt-5.6", "input": "Hello, how are you?"},
_openai_response_body("resp-bridge", "4"),
{
"id": _bridge_interaction_id("resp-bridge"),
"object": "interaction",
"model": "gpt-5.6",
"status": "completed",
"created": datetime.fromtimestamp(1700000000).isoformat(),
"updated": datetime.fromtimestamp(1700000000).isoformat(),
"outputs": [{"type": "text", "text": "4"}],
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "4"}]}],
"usage": {"total_input_tokens": 6, "total_output_tokens": 1},
},
id="responses-bridge",
),
],
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_simple_string_input(
self,
respx_mock: MockRouter,
model: str,
api_key: str,
url: str,
expected_body: dict[str, object],
response_body: dict[str, object],
expected_response: dict[str, object],
) -> None:
route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body))
response: Final = interactions.create(
model=model,
input="Hello, how are you?",
api_key=api_key,
)
assert json.loads(route.calls.last.request.content) == expected_body
assert response.model_dump(exclude_none=True) == expected_response
@pytest.mark.parametrize(
("model", "api_key", "url", "expected_body", "response_body"),
[
pytest.param(
"gemini/gemini-2.5-flash",
"gemini-offline",
"https://generativelanguage.googleapis.com/v1beta/interactions",
{
"model": "gemini-2.5-flash",
"input": "Return the required token.",
"system_instruction": "Reply with exactly this token: keep-system-instruction",
},
{
"id": "interaction-system-instruction",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"outputs": [{"type": "text", "text": "keep-system-instruction"}],
},
id="gemini",
),
pytest.param(
"gpt-5.6",
"sk-offline",
"https://api.openai.com/v1/responses",
{
"model": "gpt-5.6",
"input": "Return the required token.",
"instructions": "Reply with exactly this token: keep-system-instruction",
},
_openai_response_body("resp-system-instruction", "keep-system-instruction"),
id="responses-bridge",
),
],
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_with_system_instruction(
self,
respx_mock: MockRouter,
model: str,
api_key: str,
url: str,
expected_body: dict[str, object],
response_body: dict[str, object],
) -> None:
route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body))
response: Final = interactions.create(
model=model,
input="Return the required token.",
system_instruction="Reply with exactly this token: keep-system-instruction",
api_key=api_key,
)
assert json.loads(route.calls.last.request.content) == expected_body
assert response.status == "completed"
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_with_agent_config_sends_it_to_gemini(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(
200,
json={"id": "interaction-agent", "object": "interaction", "status": "completed"},
)
)
interactions.create(
agent="antigravity-preview-09-2026",
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
input="Say hi in one word.",
environment="remote",
api_key="gemini-offline",
)
assert json.loads(route.calls.last.request.content) == {
"agent": "antigravity-preview-09-2026",
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
"input": "Say hi in one word.",
"environment": "remote",
}
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_keeps_positional_response_settings_in_place(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(200, json={"id": "interaction-model", "status": "completed"})
)
interactions.create(
"gemini/gemini-2.5-flash",
None,
"Describe a cat.",
None,
None,
None,
None,
None,
None,
None,
["TEXT"],
api_key="gemini-offline",
)
assert json.loads(route.calls.last.request.content) == {
"model": "gemini-2.5-flash",
"input": "Describe a cat.",
"response_modalities": ["TEXT"],
}
@pytest.mark.parametrize(
("target", "agent_config", "expected_target"),
[
pytest.param(
{"model": "gemini/gemini-2.5-flash"},
{"type": "antigravity", "model": "gemini-3.8-flash"},
{"model": "gemini-2.5-flash"},
id="model-drops-it",
),
pytest.param(
{"agent": "antigravity-preview-09-2026"},
{"type": "antigravity", "model": "gemini-3.8-flash"},
{
"agent": "antigravity-preview-09-2026",
"agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"},
},
id="agent-keeps-an-object",
),
pytest.param(
{"agent": "antigravity-preview-09-2026"},
"antigravity",
{"agent": "antigravity-preview-09-2026"},
id="agent-drops-a-non-object",
),
pytest.param(
{"agent": "antigravity-preview-09-2026"},
{"model": "gemini-3.8-flash"},
{"agent": "antigravity-preview-09-2026"},
id="agent-drops-an-object-without-a-type",
),
pytest.param(
{"agent": "antigravity-preview-09-2026"},
{"type": 7, "model": "gemini-3.8-flash"},
{"agent": "antigravity-preview-09-2026"},
id="agent-drops-a-non-string-type",
),
],
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_with_drop_params_drops_agent_config_gemini_cannot_take(
self,
respx_mock: MockRouter,
target: Mapping[str, str],
agent_config: object,
expected_target: Mapping[str, object],
) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(200, json={"id": "interaction-drop", "status": "completed"})
)
interactions.create(
**target,
agent_config=agent_config,
input="Say hi in one word.",
drop_params=True,
api_key="gemini-offline",
)
assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."}
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_rejects_agent_config_a_bridged_provider_cannot_take(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi"))
)
with pytest.raises(litellm.UnsupportedParamsError, match="'agent_config' is not supported for openai") as e:
interactions.create(
model="gpt-5.6",
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
input="Say hi in one word.",
api_key="sk-offline",
)
assert e.value.status_code == 400
assert not route.called
@pytest.mark.parametrize(
"drop_kwargs",
[
pytest.param({"drop_params": True}, id="drop-params"),
pytest.param({"additional_drop_params": ["agent_config"]}, id="additional-drop-params"),
],
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_with_drop_params_drops_agent_config_a_bridged_provider_cannot_take(
self, respx_mock: MockRouter, drop_kwargs: Mapping[str, object]
) -> None:
route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi"))
)
interactions.create(
model="gpt-5.6",
agent_config={"type": "antigravity", "model": "gemini-3.8-flash"},
input="Say hi in one word.",
api_key="sk-offline",
**drop_kwargs,
)
assert json.loads(route.calls.last.request.content) == {"model": "gpt-5.6", "input": "Say hi in one word."}
@pytest.mark.parametrize(
("model", "api_key", "url", "expected_body", "sse_body", "expected_chunks"),
[
pytest.param(
"gemini/gemini-2.5-flash",
"gemini-offline",
"https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse",
{"model": "gemini-2.5-flash", "input": "Stream an answer.", "stream": True},
'data: {"event_type":"step.delta","delta":{"type":"text","text":"Hello"}}\n\n',
({"event_type": "step.delta", "object": "interaction", "delta": {"type": "text", "text": "Hello"}},),
id="gemini",
),
pytest.param(
"gpt-5.6",
"sk-offline",
"https://api.openai.com/v1/responses",
{"model": "gpt-5.6", "input": "Stream an answer.", "stream": True},
_responses_stream_body(),
(
{
"event_type": "interaction.created",
"id": _bridge_interaction_id("resp-stream"),
"object": "interaction",
"model": "gpt-5.6",
"status": "in_progress",
},
{
"event_type": "step.start",
"object": "interaction",
"index": 0,
"step": {"type": "model_output", "content": []},
},
{
"event_type": "step.delta",
"object": "interaction",
"delta": {"type": "text", "text": "Hello"},
"index": 0,
},
{"event_type": "step.stop", "object": "interaction", "index": 0},
{
"event_type": "interaction.completed",
"id": _bridge_interaction_id("resp-stream"),
"object": "interaction",
"model": "gpt-5.6",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hello"}]}],
},
),
id="responses-bridge",
),
],
)
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_streaming(
self,
respx_mock: MockRouter,
model: str,
api_key: str,
url: str,
expected_body: dict[str, object],
sse_body: str,
expected_chunks: tuple[dict[str, object], ...],
) -> None:
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, text=sse_body, headers={"content-type": "text/event-stream"})
)
response_stream: Final = interactions.create(
model=model,
input="Stream an answer.",
stream=True,
api_key=api_key,
)
chunks: Final = tuple(chunk.model_dump(exclude_none=True) for chunk in response_stream)
assert json.loads(route.calls.last.request.content) == expected_body
assert chunks == expected_chunks
@pytest.mark.usefixtures("fake_provider_credentials")
def test_invalid_model_raises_not_found(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(
404,
json={"error": {"code": 404, "message": "model not found", "status": "NOT_FOUND"}},
)
)
with pytest.raises(litellm.NotFoundError):
interactions.create(
model="gemini/invalid-model-name-xyz",
input="Hello",
api_key="gemini-offline",
)
assert json.loads(route.calls.last.request.content) == {
"model": "invalid-model-name-xyz",
"input": "Hello",
}
class TestInteractionsAcreateOffline:
@pytest.fixture(autouse=True)
def _httpx_only_transport(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.asyncio
async def test_acreate_simple_gemini(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(
200,
json={
"id": "interaction-offline",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "299792458"}]}],
"usage": {"input_tokens": 6, "output_tokens": 3},
},
)
)
response: Final = await interactions.acreate(
model="gemini/gemini-2.5-flash",
input="What is the speed of light?",
api_key="gemini-offline",
)
body: Final = json.loads(route.calls.last.request.content)
assert body == {
"model": "gemini-2.5-flash",
"input": "What is the speed of light?",
}
assert response.model_dump(exclude_none=True) == {
"id": "interaction-offline",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "299792458"}]}],
"usage": {"input_tokens": 6, "output_tokens": 3},
}
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.asyncio
async def test_acreate_streaming_gemini(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse").mock(
return_value=httpx.Response(
200,
text='data: {"event_type":"step.delta","delta":{"type":"text","text":"Hello"}}\n\n',
headers={"content-type": "text/event-stream"},
)
)
response_stream: Final = await interactions.acreate(
model="gemini/gemini-2.5-flash",
input="Stream an answer.",
stream=True,
api_key="gemini-offline",
)
chunks: Final = [chunk.model_dump(exclude_none=True) async for chunk in response_stream]
assert json.loads(route.calls.last.request.content) == {
"model": "gemini-2.5-flash",
"input": "Stream an answer.",
"stream": True,
}
assert chunks == [
{"event_type": "step.delta", "object": "interaction", "delta": {"type": "text", "text": "Hello"}},
]
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.asyncio
async def test_acreate_gemini_keeps_continuation_token_out_of_the_request(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock(
return_value=httpx.Response(
200,
json={
"id": "interaction-offline",
"object": "interaction",
"model": "gemini-3.8-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hi"}]}],
"usage": {"input_tokens": 6, "output_tokens": 3},
},
)
)
response: Final = await interactions.acreate(
model="gemini/gemini-3.8-flash",
input="Prove it",
continuation_token="dG9rLTE=",
api_key="gemini-offline",
)
# The Interactions API answered 400 "Unknown parameter 'continuation_token'." on 2026-10-13
assert json.loads(route.calls.last.request.content) == {"model": "gemini-3.8-flash", "input": "Prove it"}
assert response.status == "completed"
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.asyncio
async def test_acreate_simple_litellm_responses_bridge(self, respx_mock: MockRouter) -> None:
route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(
200,
json={
"id": "resp-offline",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "299792458 m/s"}],
}
],
"usage": {"input_tokens": 6, "output_tokens": 3, "total_tokens": 9},
},
)
)
response: Final = await interactions.acreate(
model="gpt-5.6",
input="What is the speed of light?",
api_key="sk-offline",
)
body: Final = json.loads(route.calls.last.request.content)
assert body["model"] == "gpt-5.6"
serialized: Final = json.dumps(body)
assert "What is the speed of light?" in serialized
assert "response_id:resp-offline" in base64.b64decode(response.id.removeprefix("resp_")).decode()
assert response.status == "completed"
@pytest.mark.parametrize(
("model", "api_key", "url", "response_body", "expected_body"),
[
(
"gemini/gemini-2.5-flash",
"gemini-offline",
"https://generativelanguage.googleapis.com/v1beta/interactions",
{
"id": "interaction-sync",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "four"}]}],
"usage": {"input_tokens": 5, "output_tokens": 1},
},
{"model": "gemini-2.5-flash", "input": "What is 2 + 2?"},
),
(
"gpt-4o",
"sk-offline",
"https://api.openai.com/v1/responses",
{
"id": "resp-sync",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "four"}],
}
],
"usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6},
},
{"model": "gpt-4o", "input": "What is 2 + 2?"},
),
],
)
def test_create_simple_string_input_for_google_and_responses(
respx_mock: MockRouter,
model: str,
api_key: str,
url: str,
response_body: Mapping[str, JsonValue],
expected_body: Mapping[str, JsonValue],
) -> None:
route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body))
response: Final = interactions.create(
model=model,
input="What is 2 + 2?",
api_key=api_key,
)
assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == expected_body
assert response.status == "completed"
@pytest.mark.parametrize(
("model", "api_key", "url", "response_body", "expected_body"),
[
(
"gemini/gemini-2.5-flash",
"gemini-offline",
"https://generativelanguage.googleapis.com/v1beta/interactions",
{
"id": "interaction-instruction",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "ahoy"}]}],
},
{
"model": "gemini-2.5-flash",
"input": "What are you?",
"system_instruction": "Answer like a pirate.",
},
),
(
"gpt-4o",
"sk-offline",
"https://api.openai.com/v1/responses",
{
"id": "resp-instruction",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o",
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "ahoy"}],
}
],
"usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6},
},
{"model": "gpt-4o", "input": "What are you?", "instructions": "Answer like a pirate."},
),
],
)
def test_create_with_system_instruction_for_google_and_responses(
respx_mock: MockRouter,
model: str,
api_key: str,
url: str,
response_body: Mapping[str, JsonValue],
expected_body: Mapping[str, JsonValue],
) -> None:
route: Final = respx_mock.post(url).mock(return_value=httpx.Response(200, json=response_body))
response: Final = interactions.create(
model=model,
input="What are you?",
system_instruction="Answer like a pirate.",
api_key=api_key,
)
assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == expected_body
assert response.status == "completed"
def test_create_streaming_returns_interaction_events(
respx_mock: MockRouter,
) -> None:
route: Final = respx_mock.post(
"https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse"
).mock(
return_value=httpx.Response(
200,
headers={"content-type": "text/event-stream"},
content=(
b'data: {"event_type":"step.delta","delta":{"type":"text_delta","text":"1, 2, 3"}}\n\n'
b'data: {"event_type":"interaction.completed","id":"interaction-stream","status":"completed"}\n\n'
),
)
)
chunks: Final = tuple(
interactions.create(
model="gemini/gemini-2.5-flash",
input="Count from 1 to 3.",
stream=True,
api_key="gemini-offline",
)
)
assert tuple((chunk.event_type, chunk.delta) for chunk in chunks) == (
("step.delta", {"type": "text_delta", "text": "1, 2, 3"}),
("interaction.completed", None),
)
assert _JSON_OBJECT.validate_json(route.calls.last.request.content) == {
"model": "gemini-2.5-flash",
"input": "Count from 1 to 3.",
"stream": True,
}
_GEMINI_STREAM_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse"
_GEMINI_COMPLETED_INTERACTION: Final = {
"id": "interaction-stream",
"object": "interaction",
"model": "gemini-2.5-flash",
"status": "completed",
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hello"}]}],
"usage": {"input_tokens": 4, "output_tokens": 1},
}
_GEMINI_STREAM_WITH_DONE: Final = (
'data: {"event_type":"step.delta","delta":{"type":"text","text":"Hello"}}\n\n'
"data: {not-json\n\n"
f"data: {json.dumps({'event_type': 'interaction.completed', **_GEMINI_COMPLETED_INTERACTION})}\n\n"
"data: [DONE]\n\n"
'data: {"event_type":"step.delta","delta":{"type":"text","text":"after done"}}\n\n'
)
class _CompletedInteractionRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.logged: tuple[dict[str, object], ...] = ()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
self.logged = (*self.logged, response_obj.model_dump(exclude_none=True))
async def _stream_gemini_interaction(sync_mode: bool) -> list[dict[str, object]]:
if sync_mode:
return [
chunk.model_dump(exclude_none=True)
for chunk in interactions.create(
model="gemini/gemini-2.5-flash", input="Stream an answer.", stream=True, api_key="gemini-offline"
)
]
response_stream: Final = await interactions.acreate(
model="gemini/gemini-2.5-flash", input="Stream an answer.", stream=True, api_key="gemini-offline"
)
return [chunk.model_dump(exclude_none=True) async for chunk in response_stream]
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.usefixtures("fake_provider_credentials")
@pytest.mark.asyncio
async def test_streaming_skips_malformed_lines_and_stops_at_done(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post(_GEMINI_STREAM_URL).mock(
return_value=httpx.Response(200, text=_GEMINI_STREAM_WITH_DONE, headers={"content-type": "text/event-stream"})
)
chunks: Final = await _stream_gemini_interaction(sync_mode)
assert chunks == [
{"event_type": "step.delta", "object": "interaction", "delta": {"type": "text", "text": "Hello"}},
{"event_type": "interaction.completed", **_GEMINI_COMPLETED_INTERACTION},
]
@pytest.mark.usefixtures("fake_provider_credentials")
def test_create_streaming_hands_completed_interaction_to_success_callback(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
recorder: Final = _CompletedInteractionRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
respx_mock.post(_GEMINI_STREAM_URL).mock(
return_value=httpx.Response(200, text=_GEMINI_STREAM_WITH_DONE, headers={"content-type": "text/event-stream"})
)
tuple(
interactions.create(
model="gemini/gemini-2.5-flash", input="Stream an answer.", stream=True, api_key="gemini-offline"
)
)
assert recorder.logged == (_GEMINI_COMPLETED_INTERACTION,)