mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
563dfd5075
* fix(interactions): bind Deep Research agents to their deployment and bill the agent A request that sends an agent with a model group now sends the agent to Google and is refused with a 400 unless the group's deployment is that agent. Agent interactions are logged and priced under the agent, including streamed and background jobs, with a Vertex cost row for deep-research-preview-04-2026. Background cost polling now waits 7800 seconds by default so a 120-minute job still settles. * fix(interactions): refuse extra_body that changes the agent or model A nested agent or model in extra_body was merged into the request after the agent binding check, so it could run a different agent than the one billed, inject an agent onto a model deployment, or add a second selector. Refuse those before any provider or OAuth call, while identical agent repeats and unrelated extra_body fields still pass through. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
1177 lines
46 KiB
Python
1177 lines
46 KiB
Python
import asyncio
|
|
import base64
|
|
import json
|
|
from collections.abc import Mapping
|
|
from datetime import datetime
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import responses
|
|
from pydantic import BaseModel, ConfigDict, 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"
|
|
|
|
|
|
_GEMINI_INTERACTIONS_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions"
|
|
_AGENT: Final = "antigravity-preview-09-2026"
|
|
_OTHER_AGENT: Final = "deep-research-preview-04-2026"
|
|
|
|
_EXTRA_BODY_SELECTOR_OVERRIDES: Final = (
|
|
pytest.param(
|
|
{"model": f"gemini/{_AGENT}", "agent": _AGENT}, {"agent": _OTHER_AGENT}, id="matching-agent-nested-other-agent"
|
|
),
|
|
pytest.param({"agent": _AGENT}, {"agent": None}, id="agent-nested-null-agent"),
|
|
pytest.param({"model": "gemini/gemini-4-argon"}, {"agent": _AGENT}, id="model-nested-agent"),
|
|
pytest.param({"agent": _AGENT}, {"model": "gemini-4-argon"}, id="agent-nested-model"),
|
|
pytest.param({"agent": _AGENT}, {"model": _AGENT}, id="agent-nested-model-named-like-the-agent"),
|
|
)
|
|
|
|
_EXTRA_BODY_PASS_THROUGH: Final = (
|
|
pytest.param(
|
|
{"agent": _AGENT},
|
|
{"agent": _AGENT, "store": False},
|
|
{"agent": _AGENT, "store": False},
|
|
id="same-agent-repeated",
|
|
),
|
|
pytest.param(
|
|
{"agent": _AGENT}, {"custom_flag": 1}, {"agent": _AGENT, "custom_flag": 1}, id="agent-unrelated-extra"
|
|
),
|
|
pytest.param(
|
|
{"model": "gemini/gemini-4-argon"},
|
|
{"custom_flag": 1},
|
|
{"model": "gemini-4-argon", "custom_flag": 1},
|
|
id="model-unrelated-extra",
|
|
),
|
|
pytest.param(
|
|
{"model": "gemini/gemini-4-argon"},
|
|
{"model": "gemini-4-carbon"},
|
|
{"model": "gemini-4-carbon"},
|
|
id="model-only-nested-model",
|
|
),
|
|
)
|
|
|
|
|
|
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(
|
|
{"model": "gemini/antigravity-preview-09-2026", "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="matching-model-and-agent-sends-the-agent-with-its-config",
|
|
),
|
|
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_refuses_an_agent_with_another_model_before_any_http_call(self, respx_mock: MockRouter) -> None:
|
|
with pytest.raises(litellm.BadRequestError, match="antigravity-preview-09-2026"):
|
|
interactions.create(
|
|
model="gemini/gemini-2.5-flash",
|
|
agent="antigravity-preview-09-2026",
|
|
input="Say hi in one word.",
|
|
api_key="gemini-offline",
|
|
)
|
|
|
|
assert not respx_mock.calls
|
|
|
|
@pytest.mark.parametrize(("target", "extra_body"), _EXTRA_BODY_SELECTOR_OVERRIDES)
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
def test_create_refuses_extra_body_that_changes_the_agent_or_model_before_any_http_call(
|
|
self, respx_mock: MockRouter, target: Mapping[str, str], extra_body: Mapping[str, object]
|
|
) -> None:
|
|
with pytest.raises(litellm.BadRequestError, match="extra_body"):
|
|
interactions.create(**target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline")
|
|
|
|
assert not respx_mock.calls
|
|
|
|
@pytest.mark.parametrize(("target", "extra_body", "expected_target"), _EXTRA_BODY_PASS_THROUGH)
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
def test_create_sends_extra_body_that_keeps_the_agent_and_model(
|
|
self,
|
|
respx_mock: MockRouter,
|
|
target: Mapping[str, str],
|
|
extra_body: Mapping[str, object],
|
|
expected_target: Mapping[str, object],
|
|
) -> None:
|
|
route: Final = respx_mock.post(_GEMINI_INTERACTIONS_URL).mock(
|
|
return_value=httpx.Response(200, json={"id": "interaction-extra", "status": "completed"})
|
|
)
|
|
|
|
interactions.create(**target, input="Say hi in one word.", extra_body=extra_body, 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.parametrize(("target", "extra_body"), _EXTRA_BODY_SELECTOR_OVERRIDES)
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
@pytest.mark.asyncio
|
|
async def test_acreate_refuses_extra_body_that_changes_the_agent_or_model_before_any_http_call(
|
|
self, respx_mock: MockRouter, target: Mapping[str, str], extra_body: Mapping[str, object]
|
|
) -> None:
|
|
with pytest.raises(litellm.BadRequestError, match="extra_body"):
|
|
await interactions.acreate(
|
|
**target, input="Say hi in one word.", extra_body=extra_body, api_key="gemini-offline"
|
|
)
|
|
|
|
assert not respx_mock.calls
|
|
|
|
@pytest.mark.parametrize(("target", "extra_body", "expected_target"), _EXTRA_BODY_PASS_THROUGH)
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
@pytest.mark.asyncio
|
|
async def test_acreate_sends_extra_body_that_keeps_the_agent_and_model(
|
|
self,
|
|
respx_mock: MockRouter,
|
|
target: Mapping[str, str],
|
|
extra_body: Mapping[str, object],
|
|
expected_target: Mapping[str, object],
|
|
) -> None:
|
|
route: Final = respx_mock.post(_GEMINI_INTERACTIONS_URL).mock(
|
|
return_value=httpx.Response(200, json={"id": "interaction-extra", "status": "completed"})
|
|
)
|
|
|
|
await interactions.acreate(
|
|
**target, input="Say hi in one word.", extra_body=extra_body, 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")
|
|
@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_sends_a_caller_continuation_token_to_resume_a_decode(
|
|
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-4-argon",
|
|
"status": "completed",
|
|
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hi"}]}],
|
|
"usage": {"total_input_tokens": 6, "total_output_tokens": 3},
|
|
},
|
|
)
|
|
)
|
|
response: Final = await interactions.acreate(
|
|
model="gemini/gemini-4-argon",
|
|
previous_interaction_id="interaction-hop-1",
|
|
continuation_token="dG9rLTE=",
|
|
api_key="gemini-offline",
|
|
)
|
|
# Google's spec (2026-11-05) takes continuation_token on CreateModelInteractionParams; a 2026-10-13 400 is gone
|
|
assert json.loads(route.calls.last.request.content) == {
|
|
"model": "gemini-4-argon",
|
|
"previous_interaction_id": "interaction-hop-1",
|
|
"continuation_token": "dG9rLTE=",
|
|
}
|
|
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,)
|
|
|
|
|
|
_GEMINI_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions"
|
|
_LONG_DECODE_REQUEST: Final = {
|
|
"model": "gemini/gemini-4-argon",
|
|
"input": "Prove the prime number theorem.",
|
|
"generation_config": {"max_output_tokens": 120000, "thinking_level": "high"},
|
|
}
|
|
_FIRST_HOP: Final = {
|
|
"id": "hop-1",
|
|
"object": "interaction",
|
|
"model": "gemini-4-argon",
|
|
"status": "requires_continuation",
|
|
"steps": [{"type": "model_output", "content": [{"type": "text", "text": "Lemma 1."}]}],
|
|
"usage": {
|
|
"total_input_tokens": 79,
|
|
"total_output_tokens": 32668,
|
|
"total_thought_tokens": 100,
|
|
"total_tokens": 32847,
|
|
"initial_input_tokens": 79,
|
|
},
|
|
"continuation_token": "tok-1",
|
|
}
|
|
|
|
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
def test_create_keeps_requires_continuation_status_and_token_for_the_caller(respx_mock: MockRouter) -> None:
|
|
route: Final = respx_mock.post(_GEMINI_URL).mock(return_value=httpx.Response(200, json=_FIRST_HOP))
|
|
|
|
response: Final = interactions.create(**_LONG_DECODE_REQUEST, api_key="gemini-offline")
|
|
|
|
assert route.call_count == 1
|
|
assert response.status == "requires_continuation"
|
|
assert response.continuation_token == "tok-1"
|
|
assert response.usage == _FIRST_HOP["usage"]
|
|
|
|
|
|
@pytest.mark.usefixtures("fake_provider_credentials")
|
|
def test_create_sends_a_caller_continuation_token_to_resume_a_long_decode(respx_mock: MockRouter) -> None:
|
|
route: Final = respx_mock.post(_GEMINI_URL).mock(
|
|
return_value=httpx.Response(
|
|
200, json={**_FIRST_HOP, "id": "hop-2", "status": "completed", "continuation_token": None}
|
|
)
|
|
)
|
|
|
|
response: Final = interactions.create(
|
|
model="gemini/gemini-4-argon",
|
|
previous_interaction_id="hop-1",
|
|
continuation_token="tok-1",
|
|
generation_config=_LONG_DECODE_REQUEST["generation_config"],
|
|
api_key="gemini-offline",
|
|
)
|
|
|
|
sent: Final = json.loads(route.calls.last.request.content)
|
|
assert sent == {
|
|
"model": "gemini-4-argon",
|
|
"previous_interaction_id": "hop-1",
|
|
"continuation_token": "tok-1",
|
|
"generation_config": _LONG_DECODE_REQUEST["generation_config"],
|
|
}, sent
|
|
assert response.status == "completed"
|
|
assert response.continuation_token is None
|
|
|
|
|
|
_DEEP_RESEARCH_AGENT: Final = "deep-research-preview-04-2026"
|
|
_DEEP_RESEARCH_USAGE: Final = {
|
|
"total_input_tokens": 1000,
|
|
"total_cached_tokens": 200,
|
|
"total_output_tokens": 300,
|
|
"total_thought_tokens": 100,
|
|
"total_tokens": 1400,
|
|
"input_tokens_by_modality": [{"modality": "text", "tokens": 1000}],
|
|
"output_tokens_by_modality": [{"modality": "text", "tokens": 300}],
|
|
"grounding_tool_count": [{"type": "google_search", "count": 3}],
|
|
}
|
|
|
|
|
|
def _deep_research_cost_from_the_vertex_row() -> float:
|
|
row: Final = litellm.model_cost[f"vertex_ai/{_DEEP_RESEARCH_AGENT}"]
|
|
return (
|
|
(1000 - 200) * row["input_cost_per_token"]
|
|
+ 200 * row["cache_read_input_token_cost"]
|
|
+ (300 + 100) * row["output_cost_per_token"]
|
|
+ 3 * row["search_context_cost_per_query"]["search_context_size_medium"]
|
|
)
|
|
|
|
|
|
class _LoggedModelMapValue(BaseModel):
|
|
model_config = ConfigDict(frozen=True, extra="ignore")
|
|
|
|
key: str
|
|
|
|
|
|
class _LoggedModelMap(BaseModel):
|
|
model_config = ConfigDict(frozen=True, extra="ignore")
|
|
|
|
model_map_value: _LoggedModelMapValue
|
|
|
|
|
|
class _LoggedModel(BaseModel):
|
|
model_config = ConfigDict(frozen=True, extra="ignore")
|
|
|
|
model: str
|
|
model_map_information: _LoggedModelMap
|
|
|
|
|
|
class _LoggedCost(BaseModel):
|
|
model_config = ConfigDict(frozen=True, extra="ignore")
|
|
|
|
response_cost: float
|
|
standard_logging_object: _LoggedModel
|
|
|
|
|
|
class _CostRecorder(CustomLogger):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.logged: Final = asyncio.Event()
|
|
self.cost: _LoggedCost | None = None
|
|
|
|
async def async_log_success_event(
|
|
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
|
|
) -> None:
|
|
self.cost = _LoggedCost.model_validate(kwargs)
|
|
self.logged.set()
|
|
|
|
|
|
_DEEP_RESEARCH_INTERACTION: Final = {
|
|
"id": "interaction-dr",
|
|
"object": "interaction",
|
|
"agent": _DEEP_RESEARCH_AGENT,
|
|
"status": "completed",
|
|
"usage": _DEEP_RESEARCH_USAGE,
|
|
}
|
|
_DEEP_RESEARCH_URL: Final = (
|
|
"https://aiplatform.googleapis.com/v1beta1/projects/dr-project/locations/global/interactions"
|
|
)
|
|
|
|
|
|
def _deep_research_reply(stream: bool) -> httpx.Response:
|
|
if not stream:
|
|
return httpx.Response(200, json=_DEEP_RESEARCH_INTERACTION)
|
|
completed: Final = json.dumps({"event_type": "interaction.completed", "interaction": _DEEP_RESEARCH_INTERACTION})
|
|
return httpx.Response(200, text=f"data: {completed}\n\n", headers={"content-type": "text/event-stream"})
|
|
|
|
|
|
@pytest.mark.parametrize("stream", [False, True], ids=["background-or-foreground", "streamed"])
|
|
@pytest.mark.asyncio
|
|
async def test_an_agent_only_vertex_deep_research_interaction_is_billed_under_the_agent_row(
|
|
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool
|
|
) -> None:
|
|
recorder: Final = _CostRecorder()
|
|
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
route: Final = respx_mock.post(f"{_DEEP_RESEARCH_URL}?alt=sse" if stream else _DEEP_RESEARCH_URL).mock(
|
|
return_value=_deep_research_reply(stream)
|
|
)
|
|
credentials: Final = {
|
|
"type": "authorized_user",
|
|
"client_id": "dr-client",
|
|
"client_secret": "dr-secret",
|
|
"refresh_token": "dr-refresh-token",
|
|
}
|
|
|
|
with responses.RequestsMock(assert_all_requests_are_fired=False) as google_oauth:
|
|
google_oauth.post(
|
|
"https://oauth2.googleapis.com/token",
|
|
json={"access_token": "dr-token", "expires_in": 3600, "token_type": "Bearer"},
|
|
)
|
|
response: Final = await interactions.acreate(
|
|
agent=_DEEP_RESEARCH_AGENT,
|
|
input="Compare grid-scale storage options.",
|
|
custom_llm_provider="vertex_ai",
|
|
vertex_project="dr-project",
|
|
vertex_location="global",
|
|
vertex_credentials=json.dumps(credentials),
|
|
stream=stream,
|
|
)
|
|
if stream:
|
|
async for _ in response:
|
|
pass
|
|
await asyncio.wait_for(recorder.logged.wait(), timeout=10)
|
|
|
|
assert json.loads(route.calls.last.request.content) == {
|
|
"agent": _DEEP_RESEARCH_AGENT,
|
|
"input": "Compare grid-scale storage options.",
|
|
"stream": stream,
|
|
}
|
|
assert recorder.cost is not None
|
|
assert recorder.cost.standard_logging_object.model == _DEEP_RESEARCH_AGENT
|
|
assert recorder.cost.standard_logging_object.model_map_information.model_map_value.key == (
|
|
f"vertex_ai/{_DEEP_RESEARCH_AGENT}"
|
|
)
|
|
assert recorder.cost.response_cost == pytest.approx(_deep_research_cost_from_the_vertex_row())
|