mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
* test: migrate selected llm translation contracts * test: migrate anthropic completion translation tests * test: migrate bedrock and cohere translation tests * test: tighten bedrock translation migration * test: migrate azure and openai model translation tests * test: migrate openai and small-provider translation tests * test: migrate gemini and vertex translation tests * test: migrate optional params and rerank translation tests * test: migrate huggingface, json and openrouter translation tests * test: fix openrouter integration expectations * test: migrate databricks, xai, cometapi and fireworks translation tests * test: migrate ovhcloud, prompt and small provider translation tests * test: split reasoning effort grid into unit and e2e * test: migrate router and realtime guardrail translation tests * test: migrate google interactions translation tests * test: migrate a2a, evals, skills and hosted vllm translation tests * test: remove remaining llm translation suite * test: tighten migrated llm translation unit tests and land missing rows Restores rows the first pass dropped or weakened (bedrock agentcore and openwebui thinking stream, anthropic prefix, citations and tool_choice, azure agents handler, huggingface bare api_base, optional params num_retries, alternating roles, rerank response type), replaces truthiness and call-count checks with whole-value assertions, and moves tests to the mirror file of the code they cover * test: move cohere, databricks, fireworks and gemini rows to mirror files over mocked HTTP Restores three TestCohereV2Transform methods the first pass nested out of collection, replaces direct transform calls with respx-backed completion calls, and drops the databricks and fireworks files that sat outside their mirror paths * test: close lost-intent gaps found in review of the llm translation migration Points the hosted-mock guard at the triton test's new integration path, restores exact rerank result assertions for infinity, drops the reasoning e2e fallback that passed for any multi-line answer, and adds the mixtral invoke system prompt variants, eval deletion, and async pdf url download that the legacy suite exercised * test: assert provider acceptance in the reasoning effort e2e The live half of the split grid proves each route accepts reasoning_effort. claude-fable-5-1 on the direct route answers without surfacing reasoning content, so a reasoning-presence check either fails on head or needs a fallback that passes for any answer. Wire shape stays covered by the unit grid * test: cover async and stream-end paths the legacy translation suite exercised Adds async variants for bedrock converse file and thinking content, bedrock embedding timeout and region resolution, interaction stream termination and success logging, xai reasoning token folding, and bytes chunks with a prefix before data:. Also scopes the vertex safety settings global in test_vertex_safety_settings with monkeypatch so it stops leaking into gemini request bodies on the same worker * test: type the new llm translation e2e calls and mark moved fixture dirs as packages * test: name the reasoning effort e2e models as module constants * test: assert reasoning content, register the new gemini coverage ids and give the guardrail config e2e a unique name * test: block the prompt the shared guardrail denies and drop the reasoning presence check * test: drop the duplicate litellm import left by the main sync --------- Co-authored-by: yuneng <yuneng@berri.ai>
679 lines
26 KiB
Python
679 lines
26 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.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_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,)
|