mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover google_genai endpoint contracts (#44680)
* test(integration): cover google_genai endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten google_genai contract assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover non-admin agent api_base guard and exactly-once spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): expand Google GenAI route contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover adapter responseJsonSchema, tool-loop history, /models alias and anthropic stream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin raw SSE id/retry framing and agent list/versions key guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): strengthen Gemini endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
903f04dd46
commit
32f0d680d2
5 changed files with 1563 additions and 0 deletions
98
tests/integration/providers/test_gemini_agents_wire.py
Normal file
98
tests/integration/providers/test_gemini_agents_wire.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
import json
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
_ENV_GEMINI_KEY: Final = "synthetic-shared-env-gemini-key"
|
||||
_AGENT_NAME: Final = "integration-slides-agent"
|
||||
|
||||
_AGENT_REPLY: Final = {"id": _AGENT_NAME, "base_agent": "waverunner"}
|
||||
|
||||
|
||||
def _template(params: dict) -> str:
|
||||
return urllib.parse.quote(json.dumps(params))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("caller", ["admin", "virtual_key"], ids=["admin", "virtual_key"])
|
||||
@pytest.mark.parametrize(
|
||||
"method,path",
|
||||
[
|
||||
("POST", "/v1beta/agents"),
|
||||
("GET", "/v1beta/agents"),
|
||||
("GET", f"/v1beta/agents/{_AGENT_NAME}"),
|
||||
("GET", f"/v1beta/agents/{_AGENT_NAME}/versions"),
|
||||
("DELETE", f"/v1beta/agents/{_AGENT_NAME}"),
|
||||
],
|
||||
)
|
||||
def test_agents_rejects_api_base_without_caller_key(
|
||||
gateway: Gateway, tmp_path: Path, method: str, path: str, caller: str
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(body=b"{}")
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {"GEMINI_API_KEY": _ENV_GEMINI_KEY}
|
||||
) as owned, owned.gateway.scenario() as scenario:
|
||||
virtual_key: Final = None if caller == "admin" else scenario.key()
|
||||
if method == "POST":
|
||||
response: Final = owned.gateway.request(
|
||||
"POST",
|
||||
path,
|
||||
{
|
||||
"name": _AGENT_NAME,
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "make slides",
|
||||
"litellm_params_template": {"api_base": wire.url},
|
||||
},
|
||||
key=virtual_key,
|
||||
)
|
||||
else:
|
||||
response = owned.gateway.request(
|
||||
method,
|
||||
f"{path}?litellm_params_template={_template({'api_base': wire.url})}",
|
||||
key=virtual_key,
|
||||
)
|
||||
if caller == "admin":
|
||||
assert response.status_code == 500, response.text
|
||||
assert response.json()["error"]["type"] == "internal_server_error", response.text
|
||||
assert "api_base" in response.text and "api_key" in response.text, response.text
|
||||
else:
|
||||
assert response.status_code == 401, response.text
|
||||
assert "caller-supplied" in response.json()["detail"], response.text
|
||||
assert wire.drain() == (), response.text
|
||||
|
||||
|
||||
def test_agents_forwards_caller_supplied_key_to_chosen_api_base(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/agents"
|
||||
assert request.headers["x-goog-api-key"] == "synthetic-caller-key"
|
||||
assert _ENV_GEMINI_KEY not in request.headers.values()
|
||||
assert json.loads(request.body) == {
|
||||
"name": _AGENT_NAME,
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "make slides",
|
||||
}
|
||||
return Reply(body=json.dumps(_AGENT_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
with owned_proxy_process(gateway, tmp_path, {"GEMINI_API_KEY": _ENV_GEMINI_KEY}) as owned:
|
||||
response: Final = owned.gateway.request(
|
||||
"POST",
|
||||
"/v1beta/agents",
|
||||
{
|
||||
"name": _AGENT_NAME,
|
||||
"base_agent": "waverunner",
|
||||
"instructions": "make slides",
|
||||
"litellm_params_template": {"api_base": wire.url, "api_key": "synthetic-caller-key"},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert json.loads(response.content) == {**_AGENT_REPLY, "name": _AGENT_NAME}, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/agents")]
|
||||
279
tests/integration/providers/test_gemini_interactions_wire.py
Normal file
279
tests/integration/providers/test_gemini_interactions_wire.py
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel
|
||||
|
||||
_GEMINI_MODEL: Final = "gemini/gemini-2.5-flash"
|
||||
_GEMINI_DEPLOYMENT_KEY: Final = "synthetic-gemini-deployment-key"
|
||||
|
||||
_MODEL_BODY: Final = {
|
||||
"input": "plan a trip",
|
||||
"system_instruction": "be brief",
|
||||
"tools": [{"type": "function", "name": "search", "parameters": {"type": "object"}}],
|
||||
"generation_config": {"temperature": 0.3, "image_config": {"aspect_ratio": "16:9"}},
|
||||
"response_mime_type": "application/json",
|
||||
"response_format": {"type": "object", "properties": {"plan": {"type": "string"}}},
|
||||
"previous_interaction_id": "interaction-prev-123",
|
||||
"store": True,
|
||||
"background": True,
|
||||
}
|
||||
|
||||
_INTERACTION_REPLY: Final = {
|
||||
"id": "interaction-wire-1",
|
||||
"object": "interaction",
|
||||
"model": "gemini-2.5-flash",
|
||||
"status": "completed",
|
||||
"created": "2026-05-20T00:00:00Z",
|
||||
"steps": [{"type": "content", "content": {"type": "text", "text": "itinerary"}}],
|
||||
"usage": {"total_tokens": 42},
|
||||
}
|
||||
_KNOWN_DONE_TERMINATOR: Final = b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
class _Interaction(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
model: str | None = None
|
||||
agent: str | None = None
|
||||
status: str
|
||||
created: str
|
||||
steps: list[dict]
|
||||
usage: dict
|
||||
|
||||
|
||||
def _sse_event(event_type: str, payload: dict) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _proxy_stream_event(payload: dict, model: str) -> bytes:
|
||||
event: Final = {
|
||||
"event_type": payload["event_type"],
|
||||
"id": payload["id"],
|
||||
**({"object": payload["object"]} if "object" in payload else {}),
|
||||
"model": model,
|
||||
**{key: value for key, value in payload.items() if key not in {"event_type", "id", "object"}},
|
||||
}
|
||||
return f"data: {json.dumps(event, separators=(',', ':'))}\n\n".encode()
|
||||
|
||||
|
||||
_STREAM_EVENTS: Final = (
|
||||
(
|
||||
"interaction.start",
|
||||
{
|
||||
"event_type": "interaction.start",
|
||||
"id": "interaction-wire-2",
|
||||
"object": "interaction",
|
||||
"status": "in_progress",
|
||||
},
|
||||
),
|
||||
(
|
||||
"content.delta",
|
||||
{
|
||||
"event_type": "content.delta",
|
||||
"id": "interaction-wire-2",
|
||||
"delta": {"type": "text", "text": "itin"},
|
||||
},
|
||||
),
|
||||
(
|
||||
"interaction.complete",
|
||||
{
|
||||
"event_type": "interaction.complete",
|
||||
"id": "interaction-wire-2",
|
||||
"object": "interaction",
|
||||
"status": "completed",
|
||||
"usage": {"total_tokens": 42},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/interactions", "/interactions"])
|
||||
@pytest.mark.parametrize(
|
||||
"request_input",
|
||||
[
|
||||
"plan a trip",
|
||||
[
|
||||
{"role": "user", "content": "plan a trip"},
|
||||
{"role": "model", "content": "where to"},
|
||||
{"role": "user", "content": "lisbon"},
|
||||
],
|
||||
],
|
||||
ids=["string", "turns"],
|
||||
)
|
||||
def test_interaction_forwards_native_body_to_gemini(gateway: Gateway, prefix: str, request_input: object) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/interactions"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
assert request.headers["api-revision"] == "2026-05-20"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"input": request_input,
|
||||
"tools": _MODEL_BODY["tools"],
|
||||
"system_instruction": "be brief",
|
||||
"store": True,
|
||||
"background": True,
|
||||
"previous_interaction_id": "interaction-prev-123",
|
||||
"generation_config": {"temperature": 0.3},
|
||||
"response_format": [
|
||||
{
|
||||
"type": "text",
|
||||
"mime_type": "application/json",
|
||||
"schema": {"type": "object", "properties": {"plan": {"type": "string"}}},
|
||||
},
|
||||
{"type": "image", "aspect_ratio": "16:9"},
|
||||
],
|
||||
}
|
||||
return Reply(body=json.dumps(_INTERACTION_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST", prefix, {"model": model, **_MODEL_BODY, "input": request_input}, key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = _Interaction.model_validate_json(response.content)
|
||||
assert payload.model == model, response.text
|
||||
assert response.json() == {
|
||||
**_INTERACTION_REPLY,
|
||||
"model": model,
|
||||
"agent": None,
|
||||
"updated": None,
|
||||
"outputs": None,
|
||||
}, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/interactions", "/interactions"])
|
||||
def test_interaction_authenticates_google_key_header(gateway: Gateway, prefix: str) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/interactions"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
assert all(key not in value for value in request.headers.values()), dict(request.headers)
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"input": "plan a trip",
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_INTERACTION_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
prefix,
|
||||
json={"model": model, "input": "plan a trip"},
|
||||
headers={"x-goog-api-key": key},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {
|
||||
**_INTERACTION_REPLY,
|
||||
"model": model,
|
||||
"agent": None,
|
||||
"updated": None,
|
||||
"outputs": None,
|
||||
}, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/interactions", "/interactions"])
|
||||
def test_agent_interaction_uses_env_gemini_credentials(gateway: Gateway, tmp_path: Path, prefix: str) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/interactions"
|
||||
assert request.headers["x-goog-api-key"] == "synthetic-env-gemini-key"
|
||||
assert request.headers["api-revision"] == "2026-05-20"
|
||||
body: Final = json.loads(request.body)
|
||||
assert body == {
|
||||
"agent": "deep-research-pro-preview-12-2025",
|
||||
"input": [{"role": "user", "content": "research trails"}],
|
||||
}, body
|
||||
return Reply(body=json.dumps({**_INTERACTION_REPLY, "agent": "deep-research-pro-preview-12-2025"}).encode())
|
||||
|
||||
with wire_server(respond) as wire:
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "synthetic-env-gemini-key"},
|
||||
) as owned:
|
||||
response: Final = owned.gateway.request(
|
||||
"POST",
|
||||
prefix,
|
||||
{
|
||||
"agent": "deep-research-pro-preview-12-2025",
|
||||
"input": [{"role": "user", "content": "research trails"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
_Interaction.model_validate_json(response.content)
|
||||
assert response.json() == {
|
||||
**_INTERACTION_REPLY,
|
||||
"agent": "deep-research-pro-preview-12-2025",
|
||||
"updated": None,
|
||||
"outputs": None,
|
||||
}, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/interactions", "/interactions"])
|
||||
def test_stream_interaction_relays_gemini_events(gateway: Gateway, prefix: str) -> None:
|
||||
frames: Final = tuple(_sse_event(event_type, payload) for event_type, payload in _STREAM_EVENTS)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/interactions?alt=sse"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
body: Final = json.loads(request.body)
|
||||
assert body == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"input": "plan a trip",
|
||||
"stream": True,
|
||||
}, body
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST", prefix, {"model": model, "input": "plan a trip", "stream": True}, key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
expected: Final = b"".join(_proxy_stream_event(payload, model) for _event_type, payload in _STREAM_EVENTS)
|
||||
assert response.content.removesuffix(_KNOWN_DONE_TERMINATOR) == expected, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions?alt=sse")], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/interactions", "/interactions"])
|
||||
def test_stream_interaction_ends_without_openai_done_terminator(gateway: Gateway, prefix: str) -> None:
|
||||
pytest.skip("BUG: /interactions stream=true appends an OpenAI data: [DONE] frame to the Gemini event stream")
|
||||
frames: Final = tuple(_sse_event(event_type, payload) for event_type, payload in _STREAM_EVENTS)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1beta/interactions?alt=sse"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"input": "plan a trip",
|
||||
"stream": True,
|
||||
}
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST", prefix, {"model": model, "input": "plan a trip", "stream": True}, key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
expected: Final = b"".join(_proxy_stream_event(payload, model) for _event_type, payload in _STREAM_EVENTS)
|
||||
assert response.content == expected, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions?alt=sse")]
|
||||
780
tests/integration/providers/test_google_genai_adapter_wire.py
Normal file
780
tests/integration/providers/test_google_genai_adapter_wire.py
Normal file
|
|
@ -0,0 +1,780 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
_SYSTEM_INSTRUCTION: Final = {"parts": [{"text": "use the tools"}]}
|
||||
_RESPONSE_SCHEMA: Final = {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
}
|
||||
_GENERATION_CONFIG: Final = {
|
||||
"maxOutputTokens": 50,
|
||||
"temperature": 0.1,
|
||||
"stopSequences": ["END"],
|
||||
"responseMimeType": "application/json",
|
||||
"responseSchema": _RESPONSE_SCHEMA,
|
||||
}
|
||||
_PREFIXES: Final = ("/v1beta/models", "/models")
|
||||
_SCHEMA_KEYS: Final = ("responseSchema", "responseJsonSchema")
|
||||
_IMAGE_DATA: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
_TOOLS: Final = [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "get_time",
|
||||
"parametersJsonSchema": {
|
||||
"type": "object",
|
||||
"properties": {"zone": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
_REQUEST: Final = {
|
||||
"contents": [],
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"generationConfig": _GENERATION_CONFIG,
|
||||
"tools": _TOOLS,
|
||||
"toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
|
||||
"cache": {"no-cache": True},
|
||||
}
|
||||
|
||||
_EXPECTED_TOOLS: Final = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_time",
|
||||
"parameters": {"type": "object", "properties": {"zone": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
_EXPECTED_ANTHROPIC_TOOLS: Final = [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
"type": "custom",
|
||||
},
|
||||
{
|
||||
"name": "get_time",
|
||||
"input_schema": {"type": "object", "properties": {"zone": {"type": "string"}}},
|
||||
"type": "custom",
|
||||
},
|
||||
]
|
||||
|
||||
_OPENAI_REPLY: Final = {
|
||||
"id": "chatcmpl-wire1",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "checking now",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_weather1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"SF"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 20, "completion_tokens": 9, "total_tokens": 29},
|
||||
}
|
||||
|
||||
_ANTHROPIC_REPLY: Final = {
|
||||
"id": "msg_wire1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-haiku-4-5",
|
||||
"content": [
|
||||
{"type": "text", "text": "checking now"},
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "SF"}},
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
"usage": {"input_tokens": 20, "output_tokens": 9},
|
||||
}
|
||||
|
||||
_EXPECTED_CANDIDATES: Final = [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{"text": "checking now"},
|
||||
{"functionCall": {"name": "get_weather", "args": {"city": "SF"}}},
|
||||
],
|
||||
"role": "model",
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
"index": 0,
|
||||
"safetyRatings": [],
|
||||
}
|
||||
]
|
||||
_EXPECTED_USAGE: Final = {"promptTokenCount": 20, "candidatesTokenCount": 9, "totalTokenCount": 29}
|
||||
|
||||
|
||||
def _openai_stream_frames() -> tuple[bytes, ...]:
|
||||
def frame(chunk: dict) -> bytes:
|
||||
return f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
|
||||
base: Final = {
|
||||
"id": "chatcmpl-wire2",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
return (
|
||||
frame(
|
||||
{
|
||||
**base,
|
||||
"choices": [
|
||||
{"index": 0, "delta": {"role": "assistant", "content": "checking "}, "finish_reason": None}
|
||||
],
|
||||
}
|
||||
),
|
||||
frame(
|
||||
{
|
||||
**base,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_weather1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"ci'},
|
||||
}
|
||||
]
|
||||
},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
frame(
|
||||
{
|
||||
**base,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"tool_calls": [{"index": 0, "function": {"arguments": 'ty":"SF"}'}}]},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
frame(
|
||||
{
|
||||
**base,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
|
||||
}
|
||||
),
|
||||
frame({**base, "choices": [], "usage": {"prompt_tokens": 20, "completion_tokens": 9, "total_tokens": 29}}),
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
|
||||
|
||||
def _expected_gemini_stream_frames(usage: Mapping[str, int] | None) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": [{"text": "checking "}], "role": "model"},
|
||||
"finishReason": None,
|
||||
"index": 0,
|
||||
"safetyRatings": [],
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [{"functionCall": {"name": "get_weather", "args": {"city": "SF"}}}],
|
||||
"role": "model",
|
||||
},
|
||||
"finishReason": None,
|
||||
"index": 0,
|
||||
"safetyRatings": [],
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": [], "role": "model"},
|
||||
"finishReason": "STOP",
|
||||
"index": 0,
|
||||
"safetyRatings": [],
|
||||
}
|
||||
],
|
||||
**({"usageMetadata": dict(usage)} if usage is not None else {}),
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _contents(marker: str) -> list[dict]:
|
||||
return [
|
||||
{"role": "user", "parts": [{"text": f"what is the weather in SF {marker}"}]},
|
||||
{"role": "model", "parts": [{"text": "let me check"}]},
|
||||
{"role": "user", "parts": [{"text": "and the time"}]},
|
||||
]
|
||||
|
||||
|
||||
def _request(marker: str, schema_key: str = "responseSchema") -> dict:
|
||||
generation_config: Final = {
|
||||
**{key: value for key, value in _GENERATION_CONFIG.items() if key != "responseSchema"},
|
||||
schema_key: _RESPONSE_SCHEMA,
|
||||
}
|
||||
return {**_REQUEST, "contents": _contents(marker), "generationConfig": generation_config}
|
||||
|
||||
|
||||
def _tool_loop_contents(marker: str, image_part: dict) -> list[dict]:
|
||||
return [
|
||||
{"role": "user", "parts": [{"text": f"what is in this photo {marker}"}, image_part]},
|
||||
{
|
||||
"role": "model",
|
||||
"parts": [{"text": "checking"}, {"functionCall": {"name": "get_weather", "args": {"city": "SF"}}}],
|
||||
},
|
||||
{"role": "user", "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "sunny"}}}]},
|
||||
]
|
||||
|
||||
|
||||
def _expected_anthropic_tool_loop_messages(marker: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"what is in this photo {marker}"},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _IMAGE_DATA}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "checking"},
|
||||
{"type": "tool_use", "id": "call_get_weather", "name": "get_weather", "input": {"city": "SF"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "call_get_weather", "content": '{"forecast": "sunny"}'}],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _expected_openai_tool_loop_messages(marker: str) -> list[dict]:
|
||||
return [
|
||||
{"role": "system", "content": "use the tools"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"what is in this photo {marker}"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_IMAGE_DATA}"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "checking",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_get_weather",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "SF"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_get_weather", "content": '{"forecast": "sunny"}'},
|
||||
]
|
||||
|
||||
|
||||
def _expected_openai_messages(marker: str) -> list[dict]:
|
||||
return [
|
||||
{"role": "system", "content": "use the tools"},
|
||||
{"role": "user", "content": f"what is the weather in SF {marker}"},
|
||||
{"role": "assistant", "content": "let me check"},
|
||||
{"role": "user", "content": "and the time"},
|
||||
]
|
||||
|
||||
|
||||
def _expected_anthropic_messages(marker: str) -> list[dict]:
|
||||
return [
|
||||
{"role": "user", "content": [{"type": "text", "text": f"what is the weather in SF {marker}"}]},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "let me check"}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "and the time"}]},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("schema_key", _SCHEMA_KEYS)
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_generate_content_adapts_openai_backend_request_and_response(
|
||||
gateway: Gateway, prefix: str, schema_key: str
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/chat/completions"
|
||||
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": _expected_openai_messages(marker),
|
||||
"max_tokens": 50,
|
||||
"temperature": 0.1,
|
||||
"stop": ["END"],
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "response", "schema": _RESPONSE_SCHEMA},
|
||||
},
|
||||
"tools": _EXPECTED_TOOLS,
|
||||
"tool_choice": "required",
|
||||
}
|
||||
return Reply(body=json.dumps(_OPENAI_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST", f"{prefix}/{model}:generateContent", _request(marker, schema_key), key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert body["usageMetadata"] == _EXPECTED_USAGE, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/chat/completions")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("schema_key", _SCHEMA_KEYS)
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_generate_content_adapts_anthropic_backend_request_and_response(
|
||||
gateway: Gateway, prefix: str, schema_key: str
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/messages"
|
||||
assert request.headers["x-api-key"] == "synthetic-anthropic-key"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "claude-haiku-4-5",
|
||||
"system": [{"type": "text", "text": "use the tools"}],
|
||||
"messages": _expected_anthropic_messages(marker),
|
||||
"max_tokens": 50,
|
||||
"temperature": 0.1,
|
||||
"stop_sequences": ["END"],
|
||||
"output_format": {
|
||||
"type": "json_schema",
|
||||
"schema": {**_RESPONSE_SCHEMA, "additionalProperties": False},
|
||||
},
|
||||
"tools": _EXPECTED_ANTHROPIC_TOOLS,
|
||||
"tool_choice": {"type": "any"},
|
||||
}
|
||||
return Reply(body=json.dumps(_ANTHROPIC_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="anthropic/claude-haiku-4-5", api_base=wire.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST", f"{prefix}/{model}:generateContent", _request(marker, schema_key), key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert body["usageMetadata"] == _EXPECTED_USAGE, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/messages")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_stream_generate_content_adapts_openai_sse_to_gemini_frames(gateway: Gateway, prefix: str) -> None:
|
||||
frames: Final = _openai_stream_frames()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/chat/completions"
|
||||
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": _expected_openai_messages(marker),
|
||||
"stream": True,
|
||||
"tools": _EXPECTED_TOOLS,
|
||||
"tool_choice": "required",
|
||||
}
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
f"{prefix}/{model}:streamGenerateContent?alt=sse",
|
||||
json={
|
||||
"contents": _contents(marker),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"tools": _TOOLS,
|
||||
"toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
|
||||
},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
text: Final = response.text
|
||||
assert "data: [DONE]" not in text, text
|
||||
frames_out: Final = [part for part in text.split("\n\n") if part]
|
||||
assert all(part.startswith("data: ") for part in frames_out), text
|
||||
parsed: Final = [json.loads(part.removeprefix("data: ")) for part in frames_out]
|
||||
assert "usageMetadata" in parsed[-1], text
|
||||
assert set(parsed[-1]["usageMetadata"]) == {
|
||||
"promptTokenCount",
|
||||
"candidatesTokenCount",
|
||||
"totalTokenCount",
|
||||
}, text
|
||||
assert [
|
||||
{key: value for key, value in frame.items() if key != "usageMetadata"} for frame in parsed
|
||||
] == _expected_gemini_stream_frames(None), text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/chat/completions")]
|
||||
|
||||
|
||||
def test_generate_content_adapter_accepts_snake_case_generation_config(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: top-level snake_case generation_config is dropped on the completion-adapter path, no max_tokens, temperature, stop or response_format"
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/chat/completions"
|
||||
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": _expected_openai_messages(marker),
|
||||
"max_tokens": 50,
|
||||
"temperature": 0.1,
|
||||
"stop": ["END"],
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "response", "schema": _RESPONSE_SCHEMA},
|
||||
},
|
||||
"tools": _EXPECTED_TOOLS,
|
||||
"tool_choice": "required",
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_OPENAI_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
request_body: Final = {
|
||||
"contents": _contents(marker),
|
||||
"system_instruction": _SYSTEM_INSTRUCTION,
|
||||
"generation_config": _GENERATION_CONFIG,
|
||||
"tools": _TOOLS,
|
||||
"tool_config": {"functionCallingConfig": {"mode": "ANY"}},
|
||||
"cache": {"no-cache": True},
|
||||
}
|
||||
response: Final = gateway.request("POST", f"/v1beta/models/{model}:generateContent", request_body, key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert body["usageMetadata"] == _EXPECTED_USAGE, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/chat/completions")], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
@pytest.mark.parametrize("backend", ["openai", "anthropic"])
|
||||
def test_stream_generate_content_adapter_reports_upstream_usage(gateway: Gateway, prefix: str, backend: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: native Google stream served by an openai or anthropic deployment ends with usageMetadata 0/0/0 although the upstream sent usage 20/9/29"
|
||||
)
|
||||
frames: Final = _openai_stream_frames() if backend == "openai" else _anthropic_stream_frames()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == ("/v1/chat/completions" if backend == "openai" else "/v1/messages")
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = (
|
||||
scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key")
|
||||
if backend == "openai"
|
||||
else scenario.model(
|
||||
model="anthropic/claude-haiku-4-5", api_base=wire.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
f"{prefix}/{model}:streamGenerateContent?alt=sse",
|
||||
json={
|
||||
"contents": _contents(marker),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
**({"generationConfig": {"maxOutputTokens": 50}} if backend == "anthropic" else {}),
|
||||
"tools": _TOOLS,
|
||||
"toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
|
||||
},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert "data: [DONE]" not in response.text, response.text
|
||||
frames_out: Final = [part for part in response.text.split("\n\n") if part]
|
||||
parsed: Final = [json.loads(part.removeprefix("data: ")) for part in frames_out]
|
||||
assert parsed == _expected_gemini_stream_frames(
|
||||
{"promptTokenCount": 20, "candidatesTokenCount": 9, "totalTokenCount": 29}
|
||||
), response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [
|
||||
("POST", "/v1/chat/completions" if backend == "openai" else "/v1/messages")
|
||||
], response.text
|
||||
|
||||
|
||||
def _anthropic_stream_frames() -> tuple[bytes, ...]:
|
||||
def frame(event: str, data: dict) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
|
||||
|
||||
return (
|
||||
frame(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_wire2",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-haiku-4-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 20, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
),
|
||||
frame(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
frame(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "checking "}},
|
||||
),
|
||||
frame("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
frame(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 1,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}},
|
||||
},
|
||||
),
|
||||
frame(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": '{"ci'}},
|
||||
),
|
||||
frame(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 1,
|
||||
"delta": {"type": "input_json_delta", "partial_json": 'ty":"SF"}'},
|
||||
},
|
||||
),
|
||||
frame("content_block_stop", {"type": "content_block_stop", "index": 1}),
|
||||
frame(
|
||||
"message_delta",
|
||||
{"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 9}},
|
||||
),
|
||||
frame("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_stream_generate_content_adapts_anthropic_sse_to_gemini_frames(gateway: Gateway, prefix: str) -> None:
|
||||
frames: Final = _anthropic_stream_frames()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/messages"
|
||||
assert request.headers["x-api-key"] == "synthetic-anthropic-key"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "claude-haiku-4-5",
|
||||
"system": [{"type": "text", "text": "use the tools"}],
|
||||
"messages": _expected_anthropic_messages(marker),
|
||||
"max_tokens": 50,
|
||||
"tools": _EXPECTED_ANTHROPIC_TOOLS,
|
||||
"tool_choice": {"type": "any"},
|
||||
"stream": True,
|
||||
}, request.body.decode()
|
||||
return Reply(content_type="text/event-stream", chunks=frames)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="anthropic/claude-haiku-4-5", api_base=wire.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
f"{prefix}/{model}:streamGenerateContent?alt=sse",
|
||||
json={
|
||||
"contents": _contents(marker),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"generationConfig": {"maxOutputTokens": 50},
|
||||
"tools": _TOOLS,
|
||||
"toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
|
||||
},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
text: Final = response.text
|
||||
assert "data: [DONE]" not in text, text
|
||||
frames_out: Final = [part for part in text.split("\n\n") if part]
|
||||
assert all(part.startswith("data: ") for part in frames_out), text
|
||||
parsed: Final = [json.loads(part.removeprefix("data: ")) for part in frames_out]
|
||||
assert set(parsed[-1]["usageMetadata"]) == {
|
||||
"promptTokenCount",
|
||||
"candidatesTokenCount",
|
||||
"totalTokenCount",
|
||||
}, text
|
||||
assert [
|
||||
{key: value for key, value in frame.items() if key != "usageMetadata"} for frame in parsed
|
||||
] == _expected_gemini_stream_frames(None), text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/messages")]
|
||||
|
||||
|
||||
def _tool_loop_request(marker: str, image_part: dict) -> dict:
|
||||
return {
|
||||
"contents": _tool_loop_contents(marker, image_part),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"tools": _TOOLS,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_generate_content_adapts_openai_tool_loop_history_and_inline_image(gateway: Gateway, prefix: str) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/chat/completions"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": _expected_openai_tool_loop_messages(marker),
|
||||
"tools": _EXPECTED_TOOLS,
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_OPENAI_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
image_part: Final = {"inline_data": {"mime_type": "image/png", "data": _IMAGE_DATA}}
|
||||
response: Final = gateway.request(
|
||||
"POST", f"{prefix}/{model}:generateContent", _tool_loop_request(marker, image_part), key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert body["usageMetadata"] == _EXPECTED_USAGE, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/chat/completions")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _PREFIXES)
|
||||
def test_generate_content_adapts_anthropic_tool_loop_history_and_inline_image(gateway: Gateway, prefix: str) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/messages"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "claude-haiku-4-5",
|
||||
"system": [{"type": "text", "text": "use the tools"}],
|
||||
"messages": _expected_anthropic_tool_loop_messages(marker),
|
||||
"max_tokens": 64000,
|
||||
"tools": _EXPECTED_ANTHROPIC_TOOLS,
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_ANTHROPIC_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="anthropic/claude-haiku-4-5", api_base=wire.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
image_part: Final = {"inline_data": {"mime_type": "image/png", "data": _IMAGE_DATA}}
|
||||
response: Final = gateway.request(
|
||||
"POST", f"{prefix}/{model}:generateContent", _tool_loop_request(marker, image_part), key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert body["usageMetadata"] == _EXPECTED_USAGE, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/messages")]
|
||||
|
||||
|
||||
def test_generate_content_adapter_keeps_camel_case_inline_data(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: a camelCase inlineData part (the google-genai SDK wire spelling) is silently dropped on the completion-adapter path, the image never reaches the provider"
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/chat/completions"
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": _expected_openai_tool_loop_messages(marker),
|
||||
"tools": _EXPECTED_TOOLS,
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_OPENAI_REPLY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
image_part: Final = {"inlineData": {"mimeType": "image/png", "data": _IMAGE_DATA}}
|
||||
response: Final = gateway.request(
|
||||
"POST", f"/v1beta/models/{model}:generateContent", _tool_loop_request(marker, image_part), key=key
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
assert body["candidates"] == _EXPECTED_CANDIDATES, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1/chat/completions")], response.text
|
||||
318
tests/integration/providers/test_google_genai_native_wire.py
Normal file
318
tests/integration/providers/test_google_genai_native_wire.py
Normal file
|
|
@ -0,0 +1,318 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel
|
||||
|
||||
_GEMINI_MODEL: Final = "gemini/gemini-2.5-flash"
|
||||
_GEMINI_DEPLOYMENT_KEY: Final = "synthetic-gemini-deployment-key"
|
||||
|
||||
_CONTENTS: Final = (
|
||||
{"role": "user", "parts": [{"text": "name a color"}]},
|
||||
{"role": "model", "parts": [{"text": "blue"}]},
|
||||
{"role": "user", "parts": [{"text": "name another"}]},
|
||||
)
|
||||
_SYSTEM_INSTRUCTION: Final = {"parts": [{"text": "answer with one word"}]}
|
||||
_GENERATION_CONFIG_REQUEST: Final = {
|
||||
"maxOutputTokens": 64,
|
||||
"temperature": 0.2,
|
||||
"top_p": 0.9,
|
||||
"response_mime_type": "application/json",
|
||||
"responseSchema": {
|
||||
"type": "object",
|
||||
"properties": {"color": {"type": "string"}},
|
||||
"required": ["color"],
|
||||
},
|
||||
}
|
||||
_TOOLS: Final = (
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "read the weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
_TOOL_CONFIG: Final = {"functionCallingConfig": {"mode": "AUTO"}}
|
||||
_SAFETY_SETTINGS: Final = ({"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},)
|
||||
|
||||
_REPLY_BODY: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": [{"text": "green"}], "role": "model"},
|
||||
"finishReason": "STOP",
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 11, "candidatesTokenCount": 7, "totalTokenCount": 18},
|
||||
}
|
||||
|
||||
_STREAM_FRAMES: Final = (
|
||||
b'data: {"candidates":[{"content":{"parts":[{"text":"gre"}],"role":"model"},"index":0}]}\n\n',
|
||||
b'id: gemini-frame-2\ndata: {"candidates":[{"content":{"parts":[{"text":"en"}],"role":"model"},"index":0}]}\n\n',
|
||||
b"retry: 3000\n\n",
|
||||
b'data: {"candidates":[{"content":{"parts":[],"role":"model"},"finishReason":"STOP","index":0}],'
|
||||
b'"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":7,"totalTokenCount":18}}\n\n',
|
||||
)
|
||||
|
||||
|
||||
class _Candidate(BaseModel):
|
||||
content: dict
|
||||
finishReason: str
|
||||
index: int
|
||||
|
||||
|
||||
class _GenerateContentResponse(BaseModel):
|
||||
candidates: tuple[_Candidate, ...]
|
||||
usageMetadata: dict
|
||||
|
||||
|
||||
def _assert_no_virtual_key(request: Request, key: str) -> None:
|
||||
assert all(key not in value for value in request.headers.values()), dict(request.headers)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/models", "/models"])
|
||||
def test_generate_content_rewrites_native_request_for_gemini_backend(gateway: Gateway, prefix: str) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/models/gemini-2.5-flash:generateContent"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
_assert_no_virtual_key(request, key)
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": list(_CONTENTS),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"tools": list(_TOOLS),
|
||||
"toolConfig": _TOOL_CONFIG,
|
||||
"safetySettings": list(_SAFETY_SETTINGS),
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 64,
|
||||
"temperature": 0.2,
|
||||
"topP": 0.9,
|
||||
"responseMimeType": "application/json",
|
||||
"responseJsonSchema": {
|
||||
"type": "object",
|
||||
"properties": {"color": {"type": "string"}},
|
||||
"required": ["color"],
|
||||
},
|
||||
},
|
||||
}
|
||||
return Reply(body=json.dumps(_REPLY_BODY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"{prefix}/{model}:generateContent",
|
||||
{
|
||||
"contents": list(_CONTENTS),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"generationConfig": _GENERATION_CONFIG_REQUEST,
|
||||
"tools": list(_TOOLS),
|
||||
"toolConfig": _TOOL_CONFIG,
|
||||
"safetySettings": list(_SAFETY_SETTINGS),
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = _GenerateContentResponse.model_validate_json(response.content)
|
||||
assert [candidate.model_dump() for candidate in payload.candidates] == _REPLY_BODY["candidates"], response.text
|
||||
assert payload.usageMetadata == _REPLY_BODY["usageMetadata"], response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/models/gemini-2.5-flash:generateContent")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix,query,chunking",
|
||||
[
|
||||
("/v1beta/models", "?alt=sse", "whole"),
|
||||
("/v1beta/models", "", "whole"),
|
||||
("/models", "?alt=sse", "whole"),
|
||||
("/models", "", "whole"),
|
||||
("/v1beta/models", "?alt=sse", "split"),
|
||||
("/v1beta/models", "", "split"),
|
||||
("/models", "?alt=sse", "split"),
|
||||
("/models", "", "split"),
|
||||
("/v1beta/models", "?alt=sse", "packed"),
|
||||
("/v1beta/models", "", "packed"),
|
||||
("/models", "?alt=sse", "packed"),
|
||||
("/models", "", "packed"),
|
||||
],
|
||||
)
|
||||
def test_stream_generate_content_relays_gemini_sse_frames(
|
||||
gateway: Gateway, prefix: str, query: str, chunking: str
|
||||
) -> None:
|
||||
stream_bytes: Final = b"".join(_STREAM_FRAMES)
|
||||
json_boundary: Final = stream_bytes.index(b'"gre"') + 2
|
||||
delimiter_boundary: Final = stream_bytes.index(b"\n\n") + 1
|
||||
id_boundary: Final = stream_bytes.index(b"id: gemini-frame-2") + len(b"id:")
|
||||
chunks: Final = {
|
||||
"whole": _STREAM_FRAMES,
|
||||
"split": (
|
||||
stream_bytes[:json_boundary],
|
||||
stream_bytes[json_boundary:delimiter_boundary],
|
||||
stream_bytes[delimiter_boundary:id_boundary],
|
||||
stream_bytes[id_boundary:],
|
||||
),
|
||||
"packed": (stream_bytes,),
|
||||
}[chunking]
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
_assert_no_virtual_key(request, key)
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": list(_CONTENTS),
|
||||
"tools": None,
|
||||
"generationConfig": {"maxOutputTokens": 64},
|
||||
}
|
||||
return Reply(content_type="text/event-stream", chunks=chunks)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"{prefix}/{model}:streamGenerateContent{query}",
|
||||
{"contents": list(_CONTENTS), "generationConfig": {"maxOutputTokens": 64}},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
assert response.content == stream_bytes, response.text
|
||||
assert b"[DONE]" not in response.content, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [
|
||||
("POST", "/models/gemini-2.5-flash:streamGenerateContent?alt=sse")
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/models", "/models"])
|
||||
def test_generate_content_accepts_snake_case_top_level_fields(gateway: Gateway, prefix: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: top-level snake_case generation_config, tool_config and safety_settings are dropped before reaching Gemini"
|
||||
)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/models/gemini-2.5-flash:generateContent"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
_assert_no_virtual_key(request, key)
|
||||
assert json.loads(request.body) == {
|
||||
"model": "gemini-2.5-flash",
|
||||
"contents": list(_CONTENTS),
|
||||
"systemInstruction": _SYSTEM_INSTRUCTION,
|
||||
"tools": list(_TOOLS),
|
||||
"toolConfig": _TOOL_CONFIG,
|
||||
"safetySettings": list(_SAFETY_SETTINGS),
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 64,
|
||||
"temperature": 0.2,
|
||||
"topP": 0.9,
|
||||
"responseMimeType": "application/json",
|
||||
"responseJsonSchema": {
|
||||
"type": "object",
|
||||
"properties": {"color": {"type": "string"}},
|
||||
"required": ["color"],
|
||||
},
|
||||
},
|
||||
}, request.body.decode()
|
||||
return Reply(body=json.dumps(_REPLY_BODY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"{prefix}/{model}:generateContent",
|
||||
{
|
||||
"contents": list(_CONTENTS),
|
||||
"system_instruction": _SYSTEM_INSTRUCTION,
|
||||
"generation_config": _GENERATION_CONFIG_REQUEST,
|
||||
"tools": list(_TOOLS),
|
||||
"tool_config": _TOOL_CONFIG,
|
||||
"safety_settings": list(_SAFETY_SETTINGS),
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload: Final = _GenerateContentResponse.model_validate_json(response.content)
|
||||
assert [candidate.model_dump() for candidate in payload.candidates] == _REPLY_BODY["candidates"], response.text
|
||||
assert payload.usageMetadata == _REPLY_BODY["usageMetadata"], response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/models/gemini-2.5-flash:generateContent")], (
|
||||
response.text
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/models", "/models"])
|
||||
@pytest.mark.parametrize("spelling", ["header", "query"])
|
||||
@pytest.mark.parametrize(
|
||||
"route,path",
|
||||
[
|
||||
(":generateContent", "/models/gemini-2.5-flash:generateContent"),
|
||||
(":streamGenerateContent?alt=sse", "/models/gemini-2.5-flash:streamGenerateContent?alt=sse"),
|
||||
],
|
||||
ids=["generateContent", "streamGenerateContent"],
|
||||
)
|
||||
def test_generate_content_authenticates_google_key_spellings(
|
||||
gateway: Gateway, prefix: str, spelling: str, route: str, path: str
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == path
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
_assert_no_virtual_key(request, key)
|
||||
if route == ":generateContent":
|
||||
return Reply(body=json.dumps(_REPLY_BODY).encode())
|
||||
return Reply(content_type="text/event-stream", chunks=_STREAM_FRAMES)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
key: Final = scenario.key(models=[model])
|
||||
response: Final = gateway.client.post(
|
||||
f"{prefix}/{model}{route}",
|
||||
json={"contents": list(_CONTENTS)},
|
||||
headers={"x-goog-api-key": key} if spelling == "header" else {},
|
||||
params={"key": key} if spelling == "query" else {},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [("POST", path)], response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["/v1beta/models", "/models"])
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[":generateContent", ":streamGenerateContent?alt=sse"],
|
||||
ids=["generateContent", "streamGenerateContent"],
|
||||
)
|
||||
@pytest.mark.parametrize("spelling", ["header", "query"])
|
||||
def test_generate_content_rejects_key_without_model_access(
|
||||
gateway: Gateway, prefix: str, route: str, spelling: str
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(body=json.dumps(_REPLY_BODY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_GEMINI_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY)
|
||||
other_model: Final = scenario.model(
|
||||
model="gemini/gemini-2.5-flash-lite", api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY
|
||||
)
|
||||
other_only_key: Final = scenario.key(models=[other_model])
|
||||
response: Final = gateway.client.post(
|
||||
f"{prefix}/{model}{route}",
|
||||
json={"contents": list(_CONTENTS)},
|
||||
headers={"x-goog-api-key": other_only_key} if spelling == "header" else {},
|
||||
params={"key": other_only_key} if spelling == "query" else {},
|
||||
)
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json()["error"]["type"] == "key_model_access_denied", response.text
|
||||
assert wire.drain() == (), response.text
|
||||
88
tests/integration/spend/test_google_genai_native_spend.py
Normal file
88
tests/integration/spend/test_google_genai_native_spend.py
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
import json
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
_GEMINI_MODEL: Final = "gemini/gemini-2.5-flash"
|
||||
_GEMINI_DEPLOYMENT_KEY: Final = "synthetic-gemini-deployment-key"
|
||||
_CONTENTS: Final = ({"role": "user", "parts": [{"text": "spend probe"}]},)
|
||||
_USAGE: Final = {"promptTokenCount": 11, "candidatesTokenCount": 7, "totalTokenCount": 18}
|
||||
|
||||
_REPLY_BODY: Final = {
|
||||
"candidates": [{"content": {"parts": [{"text": "ok"}], "role": "model"}, "finishReason": "STOP", "index": 0}],
|
||||
"usageMetadata": _USAGE,
|
||||
}
|
||||
|
||||
_STREAM_FRAMES: Final = (
|
||||
b'data: {"candidates":[{"content":{"parts":[{"text":"ok"}],"role":"model"},"index":0}]}\n\n',
|
||||
(
|
||||
b'data: {"candidates":[{"content":{"parts":[],"role":"model"},"finishReason":"STOP","index":0}],'
|
||||
b'"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":7,"totalTokenCount":18}}\n\n'
|
||||
),
|
||||
)
|
||||
|
||||
_CALL_TYPE: Final = {False: "agenerate_content", True: "agenerate_content_stream"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
def test_generate_content_bills_from_usage_metadata(gateway: Gateway, stream: bool) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.headers["x-goog-api-key"] == _GEMINI_DEPLOYMENT_KEY
|
||||
if stream:
|
||||
assert request.target == "/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
|
||||
return Reply(content_type="text/event-stream", chunks=_STREAM_FRAMES)
|
||||
assert request.target == "/models/gemini-2.5-flash:generateContent"
|
||||
return Reply(body=json.dumps(_REPLY_BODY).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=_GEMINI_MODEL,
|
||||
api_base=wire.url,
|
||||
api_key=_GEMINI_DEPLOYMENT_KEY,
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
)
|
||||
key: Final = scenario.key(models=[model], max_budget=10)
|
||||
query: Final = "?alt=sse" if stream else ""
|
||||
action: Final = "streamGenerateContent" if stream else "generateContent"
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1beta/models/{model}:{action}{query}",
|
||||
{"contents": list(_CONTENTS)},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT call_type, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
||||
(digest,),
|
||||
),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert rows[0]["call_type"] == _CALL_TYPE[stream], rows
|
||||
assert rows[0]["prompt_tokens"] == 11, rows
|
||||
assert rows[0]["completion_tokens"] == 7, rows
|
||||
assert float(rows[0]["spend"]) == pytest.approx(0.025), rows
|
||||
info: Final = eventually(
|
||||
lambda: object_value(gateway.get("/key/info", {"key": digest})["info"]),
|
||||
lambda values: float(values.get("spend", 0)) > 0,
|
||||
seconds=70,
|
||||
)
|
||||
assert float(info["spend"]) == pytest.approx(0.025), info
|
||||
reread: Final = read_rows(
|
||||
'SELECT call_type, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
||||
(digest,),
|
||||
)
|
||||
assert reread == rows, reread
|
||||
if not stream:
|
||||
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.025), response.headers
|
||||
assert [(r.method, r.target) for r in wire.drain()] == [
|
||||
("POST", f"/models/gemini-2.5-flash:{action}{'?alt=sse' if stream else ''}")
|
||||
]
|
||||
Loading…
Add table
Reference in a new issue