mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover pass_through endpoint contracts (#44690)
* test(integration): cover pass_through endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): run configured endpoint contracts on one proxy worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): prove pass-through streams, SDK spellings, endpoint renames and bedrock stream bodies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover pass-through wire behavior Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover Vertex and Gemini passthrough authentication Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Prove pass-through runtime method, credential-less Anthropic, stream barriers and Vertex strip layers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(pass_through): gate unquoted Bedrock stream targets and assert configured endpoint part types and methods after rename Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert pass-through query pairs so duplicate params fail 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
feb92d3255
commit
cfa98e2f40
6 changed files with 2156 additions and 4 deletions
|
|
@ -1,16 +1,20 @@
|
|||
import gzip
|
||||
import json
|
||||
import threading
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import AsyncOpenAI, NotFoundError, OpenAI
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -265,9 +269,7 @@ def test_gemini_streaming_passthrough_request_tags_reach_spend_log(gateway: Gate
|
|||
raw_tags: Final = row["request_tags"]
|
||||
request_tags: Final = json.loads(raw_tags) if isinstance(raw_tags, str) else raw_tags
|
||||
policy_tags: Final = [
|
||||
tag
|
||||
for tag in request_tags
|
||||
if not (isinstance(tag, str) and tag.startswith("User-Agent: "))
|
||||
tag for tag in request_tags if not (isinstance(tag, str) and tag.startswith("User-Agent: "))
|
||||
]
|
||||
assert policy_tags == ["gemini-js-sdk", "pass-through-endpoint"], row
|
||||
assert row["call_type"] == "pass_through_endpoint", row
|
||||
|
|
@ -650,3 +652,359 @@ def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tm
|
|||
)
|
||||
)
|
||||
assert len(budget_rows) == 1, budget_rows
|
||||
|
||||
|
||||
_GEMINI_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"candidates": [{"content": {"parts": [{"text": "scripted gemini"}], "role": "model"}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5},
|
||||
}
|
||||
_GEMINI_SSE: Final = (f"data: {json.dumps(_GEMINI_REPLY)}\r\n\r\n".encode(),)
|
||||
_GEMINI_BODY: Final[dict[str, JsonValue]] = {
|
||||
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "be terse"}]},
|
||||
"generationConfig": {"temperature": 0.2, "maxOutputTokens": 32, "thinkingConfig": {"thinkingBudget": 0}},
|
||||
}
|
||||
_GEMINI_SDK_BODY: Final[dict[str, JsonValue]] = {
|
||||
**_GEMINI_BODY,
|
||||
"systemInstruction": {"parts": [{"text": "be terse"}], "role": "user"},
|
||||
}
|
||||
|
||||
|
||||
def _upstream_query(request: Request) -> tuple[str, list[tuple[str, str]]]:
|
||||
split: Final = urlsplit(request.target)
|
||||
return split.path, parse_qsl(split.query, keep_blank_values=True)
|
||||
|
||||
|
||||
def test_gemini_passthrough_both_auth_spellings_swap_the_virtual_key_for_the_proxy_key(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
if ":streamGenerateContent" in request.target:
|
||||
return Reply(chunks=_GEMINI_SSE, content_type="text/event-stream")
|
||||
return Reply(body=json.dumps(_GEMINI_REPLY).encode())
|
||||
|
||||
path: Final = tmp_path / "gemini-auth-spellings.yaml"
|
||||
model_path: Final = "/v1beta/models/gemini-2.5-flash"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
client: Final = genai.Client(
|
||||
api_key=key,
|
||||
http_options=types.HttpOptions(
|
||||
base_url=f"{str(candidate.client.base_url).rstrip('/')}/gemini",
|
||||
client_args={"event_hooks": {"request": [sent.append]}, "timeout": 30, "trust_env": False},
|
||||
),
|
||||
)
|
||||
generated: Final = client.models.generate_content(
|
||||
model="gemini-2.5-flash",
|
||||
contents="hi",
|
||||
config=types.GenerateContentConfig(
|
||||
system_instruction="be terse",
|
||||
temperature=0.2,
|
||||
max_output_tokens=32,
|
||||
thinking_config=types.ThinkingConfig(thinking_budget=0),
|
||||
),
|
||||
)
|
||||
streamed_text: Final = "".join(
|
||||
chunk.text or ""
|
||||
for chunk in client.models.generate_content_stream(
|
||||
model="gemini-2.5-flash",
|
||||
contents="hi",
|
||||
config=types.GenerateContentConfig(
|
||||
system_instruction="be terse",
|
||||
temperature=0.2,
|
||||
max_output_tokens=32,
|
||||
thinking_config=types.ThinkingConfig(thinking_budget=0),
|
||||
),
|
||||
)
|
||||
)
|
||||
assert generated.text == "scripted gemini", generated
|
||||
assert streamed_text == "scripted gemini", streamed_text
|
||||
assert len(sent) == 2, sent
|
||||
assert all(request.headers.get("x-goog-api-key") == key for request in sent), sent
|
||||
assert parse_qsl(sent[1].url.query.decode(), keep_blank_values=True) == [("alt", "sse")], sent[1].url
|
||||
query_spelling: Final = candidate.client.post(
|
||||
f"/gemini{model_path}:generateContent", params={"key": key}, json=_GEMINI_BODY
|
||||
)
|
||||
assert query_spelling.status_code == 200, query_spelling.text
|
||||
assert query_spelling.json() == _GEMINI_REPLY, query_spelling.text
|
||||
with candidate.client.stream(
|
||||
"POST",
|
||||
f"/gemini{model_path}:streamGenerateContent",
|
||||
params={"key": key, "alt": "sse"},
|
||||
json=_GEMINI_BODY,
|
||||
) as streamed_response:
|
||||
streamed: Final = streamed_response.read()
|
||||
assert streamed_response.status_code == 200, streamed
|
||||
assert streamed == b"".join(_GEMINI_SSE), streamed
|
||||
received: Final = wire.drain()
|
||||
assert [(request.method, *_upstream_query(request)) for request in received] == [
|
||||
("POST", f"{model_path}:generateContent", [("key", "scripted")]),
|
||||
("POST", f"{model_path}:streamGenerateContent", [("alt", "sse"), ("key", "scripted")]),
|
||||
("POST", f"{model_path}:generateContent", [("key", "scripted")]),
|
||||
("POST", f"{model_path}:streamGenerateContent", [("key", "scripted"), ("alt", "sse")]),
|
||||
], received
|
||||
for upstream, sdk_request in zip(received[:2], sent, strict=True):
|
||||
sdk_body: Final = json.loads(sdk_request.content)
|
||||
assert json.loads(upstream.body) == sdk_body, upstream.body
|
||||
assert sdk_body == _GEMINI_SDK_BODY, sdk_request.content
|
||||
for request in received[2:]:
|
||||
assert json.loads(request.body) == _GEMINI_BODY, request.body
|
||||
for request in received:
|
||||
assert {name: value for name, value in request.headers.items() if key in value} == {}, (
|
||||
request.headers
|
||||
)
|
||||
assert "x-goog-api-key" not in request.headers, request.headers
|
||||
assert key not in request.target, request.target
|
||||
|
||||
|
||||
_CHAT_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"id": "chatcmpl-scripted",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted chat"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
|
||||
}
|
||||
_RESPONSE_OBJECT: Final[dict[str, JsonValue]] = {
|
||||
"id": "resp_scripted",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_scripted",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "scripted response", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"output_tokens": 3,
|
||||
"total_tokens": 8,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
},
|
||||
}
|
||||
_RESPONSE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**_RESPONSE_OBJECT, "status": "in_progress", "output": []},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_scripted",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "scripted response",
|
||||
"logprobs": [],
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": _RESPONSE_OBJECT},
|
||||
)
|
||||
_RESPONSE_SSE: Final = tuple(
|
||||
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _RESPONSE_EVENTS
|
||||
)
|
||||
_ASSISTANTS_PAGE: Final[dict[str, JsonValue]] = {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "asst_scripted",
|
||||
"object": "assistant",
|
||||
"created_at": 1,
|
||||
"name": "scripted",
|
||||
"description": None,
|
||||
"model": "gpt-4o-mini",
|
||||
"instructions": None,
|
||||
"tools": [],
|
||||
"metadata": {},
|
||||
}
|
||||
],
|
||||
"first_id": "asst_scripted",
|
||||
"last_id": "asst_scripted",
|
||||
"has_more": False,
|
||||
}
|
||||
|
||||
|
||||
_RESPONSE_STREAM_GATE: Final = threading.Event()
|
||||
_RESPONSE_STREAM_GATE.set()
|
||||
|
||||
|
||||
def _openai_peer(request: Request) -> Reply:
|
||||
if request.target.startswith("/v1/assistants"):
|
||||
return Reply(body=json.dumps(_ASSISTANTS_PAGE).encode())
|
||||
if request.target == "/v1/chat/completions":
|
||||
return Reply(body=json.dumps(_CHAT_REPLY).encode())
|
||||
if json.loads(request.body).get("stream") is True:
|
||||
return Reply(chunks=_RESPONSE_SSE, content_type="text/event-stream", gate_after_first=_RESPONSE_STREAM_GATE)
|
||||
return Reply(body=json.dumps(_RESPONSE_OBJECT).encode())
|
||||
|
||||
|
||||
def _openai_sdk(candidate: Gateway, key: str, prefix: str, sent: list[httpx.Request]) -> OpenAI:
|
||||
return OpenAI(
|
||||
api_key=key,
|
||||
base_url=f"{str(candidate.client.base_url).rstrip('/')}{prefix}/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}),
|
||||
)
|
||||
|
||||
|
||||
def _assert_openai_upstream_auth(request: Request, key: str) -> None:
|
||||
assert request.headers["authorization"] == "Bearer scripted", request.headers
|
||||
assert {name: value for name, value in request.headers.items() if key in value} == {}, request.headers
|
||||
assert key.encode() not in request.body, request.body
|
||||
|
||||
|
||||
def _assert_openai_sdk_request(
|
||||
wire: Wire,
|
||||
sent: list[httpx.Request],
|
||||
key: str,
|
||||
method: str,
|
||||
target: str,
|
||||
expected_body: dict[str, JsonValue] | None,
|
||||
) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in received] == [(method, target)], received
|
||||
assert len(sent) == 1, sent
|
||||
client_request: Final = sent[0]
|
||||
upstream: Final = received[0]
|
||||
_assert_openai_upstream_auth(upstream, key)
|
||||
if expected_body is None:
|
||||
assert client_request.content == b"", client_request.content
|
||||
assert upstream.body == b"", upstream.body
|
||||
else:
|
||||
assert json.loads(client_request.content) == expected_body, client_request.content
|
||||
assert json.loads(upstream.body) == expected_body, upstream.body
|
||||
return upstream
|
||||
|
||||
|
||||
_OPENAI_CHAT_BODY: Final[dict[str, JsonValue]] = {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"user": "end-user-1",
|
||||
"store": True,
|
||||
"temperature": 0.2,
|
||||
}
|
||||
_OPENAI_RESPONSES_BODY: Final[dict[str, JsonValue]] = {
|
||||
"model": "gpt-4o-mini",
|
||||
"input": "hi",
|
||||
"store": False,
|
||||
"user": "end-user-1",
|
||||
"temperature": 0.2,
|
||||
}
|
||||
_OPENAI_RESPONSES_STREAM_BODY: Final[dict[str, JsonValue]] = {
|
||||
**_OPENAI_RESPONSES_BODY,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("base_path", "surface"),
|
||||
[
|
||||
pytest.param("/openai", "chat", id="openai-chat"),
|
||||
pytest.param("/openai_passthrough", "chat", id="openai-passthrough-chat"),
|
||||
pytest.param("/openai_passthrough", "responses", id="openai-passthrough-responses"),
|
||||
pytest.param("/openai_passthrough", "responses-stream", id="openai-passthrough-responses-stream"),
|
||||
pytest.param("/openai", "assistants", id="openai-assistants"),
|
||||
pytest.param("/openai_passthrough", "assistants", id="openai-passthrough-assistants"),
|
||||
],
|
||||
)
|
||||
def test_openai_passthrough_sdk_success_on_both_aliases_forwards_the_native_request(
|
||||
gateway: Gateway, tmp_path: Path, base_path: str, surface: str
|
||||
) -> None:
|
||||
path: Final = tmp_path / "openai-success.yaml"
|
||||
with wire_server(_openai_peer) as wire:
|
||||
_openai_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
with _openai_sdk(candidate, key, base_path, sent) as sdk:
|
||||
if surface == "chat":
|
||||
chat: Final = sdk.chat.completions.create(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
user="end-user-1",
|
||||
store=True,
|
||||
temperature=0.2,
|
||||
)
|
||||
assert chat.model_dump(exclude_unset=True) == _CHAT_REPLY, chat
|
||||
_assert_openai_sdk_request(wire, sent, key, "POST", "/v1/chat/completions", _OPENAI_CHAT_BODY)
|
||||
elif surface == "responses":
|
||||
response: Final = sdk.responses.create(
|
||||
model="gpt-4o-mini",
|
||||
input="hi",
|
||||
store=False,
|
||||
user="end-user-1",
|
||||
temperature=0.2,
|
||||
)
|
||||
assert response.output_text == "scripted response", response
|
||||
assert response.id == "resp_scripted", response
|
||||
_assert_openai_sdk_request(wire, sent, key, "POST", "/v1/responses", _OPENAI_RESPONSES_BODY)
|
||||
elif surface == "responses-stream":
|
||||
_RESPONSE_STREAM_GATE.clear()
|
||||
try:
|
||||
stream: Final = iter(
|
||||
sdk.responses.create(
|
||||
model="gpt-4o-mini",
|
||||
input="hi",
|
||||
store=False,
|
||||
user="end-user-1",
|
||||
temperature=0.2,
|
||||
stream=True,
|
||||
)
|
||||
)
|
||||
first_event: Final = next(stream)
|
||||
finally:
|
||||
_RESPONSE_STREAM_GATE.set()
|
||||
assert first_event.type == _RESPONSE_EVENTS[0]["type"], first_event
|
||||
events: Final = (first_event, *stream)
|
||||
streamed_text: Final = "".join(
|
||||
event.delta for event in events if event.type == "response.output_text.delta"
|
||||
)
|
||||
final: Final = next(event.response for event in events if event.type == "response.completed")
|
||||
assert [event.type for event in events] == [event["type"] for event in _RESPONSE_EVENTS], events
|
||||
assert streamed_text == "scripted response", events
|
||||
assert final.output_text == "scripted response", final
|
||||
_assert_openai_sdk_request(
|
||||
wire, sent, key, "POST", "/v1/responses", _OPENAI_RESPONSES_STREAM_BODY
|
||||
)
|
||||
else:
|
||||
assistants: Final = sdk.beta.assistants.list(limit=2, order="desc")
|
||||
assert [assistant.id for assistant in assistants.data] == ["asst_scripted"], assistants
|
||||
upstream: Final = _assert_openai_sdk_request(
|
||||
wire,
|
||||
sent,
|
||||
key,
|
||||
"GET",
|
||||
"/v1/assistants?limit=2&order=desc",
|
||||
None,
|
||||
)
|
||||
assert upstream.headers["openai-beta"] == "assistants=v2", upstream.headers
|
||||
|
||||
|
||||
def test_openai_passthrough_responses_metadata_reaches_upstream(gateway: Gateway, tmp_path: Path) -> None:
|
||||
pytest.skip("BUG: /openai_passthrough/v1/responses drops the native metadata field before forwarding upstream")
|
||||
path: Final = tmp_path / "openai-metadata.yaml"
|
||||
with wire_server(_openai_peer) as wire:
|
||||
_openai_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.scenario() as scenario:
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
with _openai_sdk(candidate, scenario.key(), "/openai_passthrough", sent) as sdk:
|
||||
sdk.responses.create(model="gpt-4o-mini", input="hi", metadata={"session": "scripted"})
|
||||
received: Final = wire.drain()
|
||||
assert [json.loads(request.body) for request in received] == [json.loads(sent[0].content)], received
|
||||
|
|
|
|||
332
tests/integration/providers/test_anthropic_passthrough_wire.py
Normal file
332
tests/integration/providers/test_anthropic_passthrough_wire.py
Normal file
|
|
@ -0,0 +1,332 @@
|
|||
import json
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from itertools import accumulate, dropwhile
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_API_KEY: Final = "scripted-anthropic-key"
|
||||
_OAUTH_TOKEN: Final = "sk-ant-oat01-scripted-proxy-token"
|
||||
_MODEL: Final = "claude-sonnet-4-5"
|
||||
_TEXT: Final = "scripted passthrough answer"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_SYSTEM: Final[list[JsonValue]] = [{"type": "text", "text": "You are a terse assistant"}]
|
||||
_TOOLS: Final[list[JsonValue]] = [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Look up the weather",
|
||||
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
}
|
||||
]
|
||||
_THINKING: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024}
|
||||
_MESSAGE: Final[dict[str, JsonValue]] = {
|
||||
"id": "msg_scripted",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": _MODEL,
|
||||
"content": [{"type": "text", "text": _TEXT}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 7, "output_tokens": 3},
|
||||
}
|
||||
_SSE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{"type": "message_start", "message": {**_MESSAGE, "content": [], "stop_reason": None}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _TEXT}},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 3},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
_SSE_FRAMES: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _SSE_EVENTS)
|
||||
_CLAUDE_CODE_STREAM_GATE: Final = threading.Event()
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
gate: Final = _CLAUDE_CODE_STREAM_GATE if request.headers.get("x-app") == "cli" else None
|
||||
return Reply(chunks=_SSE_FRAMES, content_type="text/event-stream", gate_after_first=gate)
|
||||
return Reply(body=json.dumps(_MESSAGE).encode())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig() -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as gateway:
|
||||
yield gateway
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wire() -> Iterator[Wire]:
|
||||
with wire_server(_respond) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained(wire: Wire) -> None:
|
||||
wire.drain()
|
||||
|
||||
|
||||
def _config(directory: Path) -> Path:
|
||||
path: Final = directory / "anthropic-passthrough.yaml"
|
||||
path.write_text(
|
||||
"model_list: []\n"
|
||||
"general_settings:\n"
|
||||
" master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
" database_url: os.environ/DATABASE_URL\n"
|
||||
" store_model_in_db: true\n"
|
||||
"router_settings:\n"
|
||||
" disable_cooldowns: true\n"
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def api_key_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-api-key")
|
||||
overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _API_KEY}
|
||||
with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def oauth_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-oauth")
|
||||
overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _OAUTH_TOKEN}
|
||||
with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def credentialless_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-credentialless")
|
||||
overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": "", "ANTHROPIC_AUTH_TOKEN": ""}
|
||||
with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _sdk(proxy: Gateway, key: str, sent: list[httpx.Request]) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(
|
||||
api_key=key,
|
||||
base_url=f"{str(proxy.client.base_url).rstrip('/')}/anthropic",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}),
|
||||
)
|
||||
|
||||
|
||||
def _only_upstream_request(wire: Wire) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in received] == [("POST", "/v1/messages")], received
|
||||
return received[0]
|
||||
|
||||
|
||||
def _assert_virtual_key_absent(upstream: Request, key: str) -> None:
|
||||
leaked: Final = {name: value for name, value in upstream.headers.items() if key in value}
|
||||
assert leaked == {}, upstream.headers
|
||||
assert key.encode() not in upstream.body, upstream.body
|
||||
|
||||
|
||||
def _assert_sdk_call_reached_upstream_intact(upstream: Request, sent: httpx.Request, key: str) -> None:
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent.content), upstream.body
|
||||
assert upstream.headers["x-api-key"] == _API_KEY, upstream.headers
|
||||
assert "authorization" not in upstream.headers, upstream.headers
|
||||
assert upstream.headers["anthropic-version"] == sent.headers["anthropic-version"], upstream.headers
|
||||
assert upstream.headers["anthropic-beta"] == "interleaved-thinking-2025-05-14", upstream.headers
|
||||
_assert_virtual_key_absent(upstream, key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"])
|
||||
def test_anthropic_sdk_call_reaches_upstream_with_full_body_and_proxy_key(
|
||||
api_key_proxy: Gateway, wire: Wire, stream: bool
|
||||
) -> None:
|
||||
with api_key_proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
with _sdk(api_key_proxy, key, sent) as sdk:
|
||||
arguments: Final = {
|
||||
"model": _MODEL,
|
||||
"max_tokens": 2048,
|
||||
"system": _SYSTEM,
|
||||
"tools": _TOOLS,
|
||||
"thinking": _THINKING,
|
||||
"messages": [{"role": "user", "content": "weather in Paris?"}],
|
||||
"extra_headers": {"anthropic-beta": "interleaved-thinking-2025-05-14"},
|
||||
}
|
||||
if stream:
|
||||
with sdk.messages.stream(**arguments) as events:
|
||||
message = events.get_final_message()
|
||||
else:
|
||||
message = sdk.messages.create(**arguments)
|
||||
assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump(
|
||||
exclude_none=True
|
||||
), message
|
||||
upstream: Final = _only_upstream_request(wire)
|
||||
assert len(sent) == 1, sent
|
||||
_assert_sdk_call_reached_upstream_intact(upstream, sent[0], key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("auth_style", ["api-key", "oauth"])
|
||||
def test_anthropic_proxy_credentials_override_caller_credentials(
|
||||
api_key_proxy: Gateway, oauth_proxy: Gateway, wire: Wire, auth_style: str
|
||||
) -> None:
|
||||
proxy: Final = api_key_proxy if auth_style == "api-key" else oauth_proxy
|
||||
caller_api_key: Final = f"caller-anthropic-key-{auth_style}"
|
||||
caller_authorization: Final = f"Bearer caller-anthropic-token-{auth_style}"
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": _MODEL,
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
with (
|
||||
proxy.scenario() as scenario,
|
||||
httpx.Client(base_url=str(proxy.client.base_url), timeout=30, trust_env=False) as client,
|
||||
):
|
||||
key: Final = scenario.key()
|
||||
response: Final = client.post(
|
||||
"/anthropic/v1/messages",
|
||||
json=body,
|
||||
headers={
|
||||
"x-litellm-api-key": key,
|
||||
"x-api-key": caller_api_key,
|
||||
"Authorization": caller_authorization,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
message: Final = anthropic.types.Message.model_validate_json(response.content)
|
||||
assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump(
|
||||
exclude_none=True
|
||||
), response.text
|
||||
upstream: Final = _only_upstream_request(wire)
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
assert {
|
||||
name: value
|
||||
for name, value in upstream.headers.items()
|
||||
if caller_api_key in value or caller_authorization in value
|
||||
} == {}, upstream.headers
|
||||
if auth_style == "api-key":
|
||||
assert upstream.headers["x-api-key"] == _API_KEY, upstream.headers
|
||||
assert "authorization" not in upstream.headers, upstream.headers
|
||||
else:
|
||||
assert upstream.headers["authorization"] == f"Bearer {_OAUTH_TOKEN}", upstream.headers
|
||||
assert "x-api-key" not in upstream.headers, upstream.headers
|
||||
_assert_virtual_key_absent(upstream, key)
|
||||
|
||||
|
||||
def test_claude_code_bearer_stream_merges_caller_beta_with_the_credential_beta(
|
||||
oauth_proxy: Gateway, wire: Wire
|
||||
) -> None:
|
||||
_CLAUDE_CODE_STREAM_GATE.clear()
|
||||
with oauth_proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": _MODEL,
|
||||
"max_tokens": 2048,
|
||||
"stream": True,
|
||||
"system": _SYSTEM,
|
||||
"tools": _TOOLS,
|
||||
"thinking": _THINKING,
|
||||
"messages": [{"role": "user", "content": "weather in Paris?"}],
|
||||
}
|
||||
with oauth_proxy.client.stream(
|
||||
"POST",
|
||||
"/anthropic/v1/messages",
|
||||
json=body,
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14",
|
||||
"x-app": "cli",
|
||||
},
|
||||
) as response:
|
||||
raw: Final = iter(response.iter_raw())
|
||||
try:
|
||||
assert response.status_code == 200, response.status_code
|
||||
assert response.headers.get("content-type", "").startswith("text/event-stream"), response.headers
|
||||
first_frame: Final = next(
|
||||
dropwhile(
|
||||
lambda buffered: b"\n\n" not in buffered,
|
||||
accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""),
|
||||
)
|
||||
)
|
||||
assert first_frame == _SSE_FRAMES[0], first_frame
|
||||
finally:
|
||||
_CLAUDE_CODE_STREAM_GATE.set()
|
||||
streamed: Final = first_frame + b"".join(raw)
|
||||
assert streamed == b"".join(_SSE_FRAMES), streamed
|
||||
_CLAUDE_CODE_STREAM_GATE.clear()
|
||||
upstream: Final = _only_upstream_request(wire)
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
assert upstream.headers["authorization"] == f"Bearer {_OAUTH_TOKEN}", upstream.headers
|
||||
assert "x-api-key" not in upstream.headers, upstream.headers
|
||||
assert upstream.headers["anthropic-version"] == "2023-06-01", upstream.headers
|
||||
assert upstream.headers["anthropic-beta"] == (
|
||||
"claude-code-20250219,interleaved-thinking-2025-05-14,oauth-2025-04-20"
|
||||
), upstream.headers
|
||||
assert upstream.headers["x-app"] == "cli", upstream.headers
|
||||
_assert_virtual_key_absent(upstream, key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("credential_header", "credential"),
|
||||
[("authorization", "Bearer caller-anthropic-oauth-token"), ("x-api-key", "caller-anthropic-api-key")],
|
||||
ids=["caller-oauth-token", "caller-api-key"],
|
||||
)
|
||||
def test_credentialless_proxy_forwards_the_caller_anthropic_credential_and_never_the_virtual_key(
|
||||
credentialless_proxy: Gateway, wire: Wire, credential_header: str, credential: str
|
||||
) -> None:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": _MODEL,
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
with (
|
||||
credentialless_proxy.scenario() as scenario,
|
||||
httpx.Client(base_url=str(credentialless_proxy.client.base_url), timeout=30, trust_env=False) as client,
|
||||
):
|
||||
key: Final = scenario.key()
|
||||
response: Final = client.post(
|
||||
"/anthropic/v1/messages", json=body, headers={"x-litellm-api-key": key, credential_header: credential}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
message: Final = anthropic.types.Message.model_validate_json(response.content)
|
||||
assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump(
|
||||
exclude_none=True
|
||||
), response.text
|
||||
upstream: Final = _only_upstream_request(wire)
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
assert upstream.headers[credential_header] == credential, upstream.headers
|
||||
assert {"authorization", "x-api-key"} & set(upstream.headers) == {credential_header}, upstream.headers
|
||||
_assert_virtual_key_absent(upstream, key)
|
||||
|
||||
refused: Final = client.post("/anthropic/v1/messages", json=body, headers={"x-litellm-api-key": key})
|
||||
assert refused.status_code == 401, refused.text
|
||||
assert "The LiteLLM virtual key is not forwarded to Anthropic" in refused.text, refused.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_anthropic_sdk_metadata_user_id_reaches_upstream(api_key_proxy: Gateway, wire: Wire) -> None:
|
||||
pytest.skip("BUG: /anthropic/v1/messages drops the native metadata.user_id field before forwarding upstream")
|
||||
with api_key_proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
with _sdk(api_key_proxy, key, sent) as sdk:
|
||||
sdk.messages.create(
|
||||
model=_MODEL,
|
||||
max_tokens=64,
|
||||
metadata={"user_id": "user_scripted_session_1"},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
upstream: Final = _only_upstream_request(wire)
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent[0].content), upstream.body
|
||||
|
|
@ -1,10 +1,22 @@
|
|||
import base64
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
import yaml
|
||||
from botocore.awsrequest import AWSPreparedRequest
|
||||
from botocore.config import Config
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.upstream import _aws_event_frame
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
|
||||
|
|
@ -40,3 +52,299 @@ def test_bedrock_passthrough_converse_stream_response_carries_event_stream_conte
|
|||
assert len(wire.drain()) == 1, response.text
|
||||
assert response.headers.get("content-type") == _EVENT_STREAM, dict(response.headers)
|
||||
assert response.content == _STREAM_BYTES, response.text
|
||||
|
||||
|
||||
_ACCOUNT: Final = "123456789012"
|
||||
_CONVERSE_BODY: Final[dict[str, JsonValue]] = {
|
||||
"messages": [{"role": "user", "content": [{"text": "full body"}]}],
|
||||
"system": [{"text": "be terse"}],
|
||||
"inferenceConfig": {"maxTokens": 64, "temperature": 0.2, "stopSequences": ["END"]},
|
||||
"toolConfig": {
|
||||
"tools": [
|
||||
{
|
||||
"toolSpec": {
|
||||
"name": "lookup",
|
||||
"description": "look a value up",
|
||||
"inputSchema": {"json": {"type": "object", "properties": {"q": {"type": "string"}}}},
|
||||
}
|
||||
}
|
||||
],
|
||||
"toolChoice": {"auto": {}},
|
||||
},
|
||||
}
|
||||
_INVOKE_BODY: Final[dict[str, JsonValue]] = {
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 64,
|
||||
"system": "be terse",
|
||||
"messages": [{"role": "user", "content": "full body"}],
|
||||
}
|
||||
_CONVERSE_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock converse"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 5, "outputTokens": 2, "totalTokens": 7},
|
||||
"metrics": {"latencyMs": 12},
|
||||
}
|
||||
_INVOKE_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"id": "msg_bedrock",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude",
|
||||
"content": [{"type": "text", "text": "bedrock invoke"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 5, "output_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
_CONVERSE_STREAM_FRAMES: Final = tuple(_aws_event_frame(kind, payload, "sc", "u") for kind, payload in _EVENTS)
|
||||
_INVOKE_STREAM_FRAMES: Final = (
|
||||
_aws_event_frame(
|
||||
"chunk",
|
||||
{"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()},
|
||||
"sc",
|
||||
"u",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_GATED_STREAM_TARGETS: Final = (f"/model/{_MODEL_ID}/converse-stream", "-direct-")
|
||||
_BEDROCK_STREAM_GATE: Final = threading.Event()
|
||||
_BEDROCK_STREAM_GATE.set()
|
||||
|
||||
|
||||
def _full_body_peer(request: Request) -> Reply:
|
||||
if request.target.endswith("/converse-stream"):
|
||||
gated: Final = any(marker in unquote(request.target) for marker in _GATED_STREAM_TARGETS)
|
||||
return Reply(
|
||||
chunks=_CONVERSE_STREAM_FRAMES,
|
||||
content_type=_EVENT_STREAM,
|
||||
gate_after_first=_BEDROCK_STREAM_GATE if gated else None,
|
||||
)
|
||||
if request.target.endswith("/invoke-with-response-stream"):
|
||||
return Reply(chunks=_INVOKE_STREAM_FRAMES, content_type=_EVENT_STREAM)
|
||||
if request.target.endswith("/converse"):
|
||||
return Reply(body=json.dumps(_CONVERSE_REPLY).encode())
|
||||
return Reply(body=json.dumps(_INVOKE_REPLY).encode())
|
||||
|
||||
|
||||
def _boto3_runtime(gateway: Gateway, key: str): # noqa: ANN202 - boto3 clients are untyped
|
||||
client: Final = boto3.client(
|
||||
"bedrock-runtime",
|
||||
endpoint_url=f"{str(gateway.client.base_url).rstrip('/')}/bedrock",
|
||||
region_name="us-east-1",
|
||||
aws_access_key_id="AKIACALLERSIDE",
|
||||
aws_secret_access_key="caller-side-secret",
|
||||
config=Config(retries={"max_attempts": 1}),
|
||||
)
|
||||
|
||||
def bearer(request: AWSPreparedRequest, **_: object) -> None:
|
||||
request.headers["Authorization"] = f"Bearer {key}"
|
||||
|
||||
client.meta.events.register("before-send.bedrock-runtime.*", bearer)
|
||||
return client
|
||||
|
||||
|
||||
def _events_read_before_the_upstream_finishes(open_stream: Callable[[], Iterable[object]]) -> tuple[object, ...]:
|
||||
_BEDROCK_STREAM_GATE.clear()
|
||||
events: Final = iter(open_stream())
|
||||
try:
|
||||
first: Final = next(events)
|
||||
finally:
|
||||
_BEDROCK_STREAM_GATE.set()
|
||||
return (first, *events)
|
||||
|
||||
|
||||
def _assert_signed_with_the_deployment_key(request: Request, key: str) -> None:
|
||||
authorization: Final = request.headers["authorization"]
|
||||
assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIASCRIPTEDPROVIDER/"), request.headers
|
||||
assert "/us-east-1/bedrock/aws4_request" in authorization, request.headers
|
||||
assert {name: value for name, value in request.headers.items() if key in value} == {}, request.headers
|
||||
assert key not in request.target, request.target
|
||||
|
||||
|
||||
def test_boto3_converse_invoke_and_stream_methods_forward_full_bodies_for_deployments_ids_and_profile_arns(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
suffix: Final = uuid.uuid4().hex[:12]
|
||||
raw_id: Final = f"anthropic.claude-scripted-{suffix}-v1:0"
|
||||
profile_arn: Final = f"arn:aws:bedrock:us-east-1:{_ACCOUNT}:application-inference-profile/{suffix}"
|
||||
router_name: Final = f"bedrock-router-{suffix}"
|
||||
base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
with wire_server(_full_body_peer) as wire:
|
||||
deployment_params: Final = {
|
||||
"api_base": wire.url,
|
||||
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
|
||||
"aws_secret_access_key": "scripted-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
"use_in_pass_through": True,
|
||||
}
|
||||
config: Final = {
|
||||
**base_config,
|
||||
"model_list": [
|
||||
*base_config.get("model_list", []),
|
||||
{
|
||||
"model_name": router_name,
|
||||
"litellm_params": {"model": f"bedrock/{_MODEL_ID}", **deployment_params},
|
||||
},
|
||||
{
|
||||
"model_name": raw_id,
|
||||
"litellm_params": {"model": f"bedrock/{raw_id}", **deployment_params},
|
||||
},
|
||||
{
|
||||
"model_name": profile_arn,
|
||||
"litellm_params": {"model": f"bedrock/{profile_arn}", **deployment_params},
|
||||
},
|
||||
],
|
||||
}
|
||||
config_path: Final = tmp_path / "bedrock-passthrough.yaml"
|
||||
config_path.write_text(yaml.safe_dump(config))
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {}, config=config_path, workers=1) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key()
|
||||
runtime: Final = _boto3_runtime(candidate, key)
|
||||
router_converse: Final = runtime.converse(modelId=router_name, **_CONVERSE_BODY)
|
||||
profile_converse: Final = runtime.converse(modelId=profile_arn, **_CONVERSE_BODY)
|
||||
raw_invoke: Final = runtime.invoke_model(modelId=raw_id, body=json.dumps(_INVOKE_BODY))
|
||||
profile_invoke: Final = runtime.invoke_model(modelId=profile_arn, body=json.dumps(_INVOKE_BODY))
|
||||
raw_converse_stream: Final = runtime.converse_stream(modelId=raw_id, **_CONVERSE_BODY)
|
||||
profile_converse_stream: Final = runtime.converse_stream(modelId=profile_arn, **_CONVERSE_BODY)
|
||||
raw_invoke_stream: Final = runtime.invoke_model_with_response_stream(
|
||||
modelId=raw_id, body=json.dumps(_INVOKE_BODY)
|
||||
)
|
||||
profile_invoke_stream: Final = runtime.invoke_model_with_response_stream(
|
||||
modelId=profile_arn, body=json.dumps(_INVOKE_BODY)
|
||||
)
|
||||
router_converse_stream: Final = _events_read_before_the_upstream_finishes(
|
||||
lambda: runtime.converse_stream(modelId=router_name, **_CONVERSE_BODY)["stream"]
|
||||
)
|
||||
router_invoke_stream: Final = runtime.invoke_model_with_response_stream(
|
||||
modelId=router_name, body=json.dumps(_INVOKE_BODY)
|
||||
)
|
||||
assert tuple(
|
||||
{name: reply[name] for name in _CONVERSE_REPLY} for reply in (router_converse, profile_converse)
|
||||
) == (_CONVERSE_REPLY, _CONVERSE_REPLY), (router_converse, profile_converse)
|
||||
assert tuple(json.loads(response["body"].read()) for response in (raw_invoke, profile_invoke)) == (
|
||||
_INVOKE_REPLY,
|
||||
_INVOKE_REPLY,
|
||||
), (raw_invoke, profile_invoke)
|
||||
expected_converse_events: Final = [{kind: payload} for kind, payload in _EVENTS]
|
||||
assert [list(response["stream"]) for response in (raw_converse_stream, profile_converse_stream)] == [
|
||||
expected_converse_events,
|
||||
expected_converse_events,
|
||||
], (raw_converse_stream, profile_converse_stream)
|
||||
assert [
|
||||
[json.loads(event["chunk"]["bytes"]) for event in response["body"]]
|
||||
for response in (raw_invoke_stream, profile_invoke_stream)
|
||||
] == [[_INVOKE_REPLY], [_INVOKE_REPLY]], (raw_invoke_stream, profile_invoke_stream)
|
||||
assert list(router_converse_stream) == expected_converse_events, router_converse_stream
|
||||
assert [json.loads(event["chunk"]["bytes"]) for event in router_invoke_stream["body"]] == [_INVOKE_REPLY]
|
||||
|
||||
received: Final = wire.drain()
|
||||
encoded_arn: Final = quote(profile_arn, safe=":")
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"/model/{_MODEL_ID}/converse"),
|
||||
("POST", f"/model/{encoded_arn}/converse"),
|
||||
("POST", f"/model/{raw_id}/invoke"),
|
||||
("POST", f"/model/{encoded_arn}/invoke"),
|
||||
("POST", f"/model/{raw_id}/converse-stream"),
|
||||
("POST", f"/model/{encoded_arn}/converse-stream"),
|
||||
("POST", f"/model/{raw_id}/invoke-with-response-stream"),
|
||||
("POST", f"/model/{encoded_arn}/invoke-with-response-stream"),
|
||||
("POST", f"/model/{_MODEL_ID}/converse-stream"),
|
||||
("POST", f"/model/{_MODEL_ID}/invoke-with-response-stream"),
|
||||
], received
|
||||
assert tuple(json.loads(request.body) for request in received) == (
|
||||
_CONVERSE_BODY,
|
||||
_CONVERSE_BODY,
|
||||
_INVOKE_BODY,
|
||||
_INVOKE_BODY,
|
||||
_CONVERSE_BODY,
|
||||
_CONVERSE_BODY,
|
||||
_INVOKE_BODY,
|
||||
_INVOKE_BODY,
|
||||
_CONVERSE_BODY,
|
||||
_INVOKE_BODY,
|
||||
), tuple(request.body for request in received)
|
||||
for request in received:
|
||||
_assert_signed_with_the_deployment_key(request, key)
|
||||
|
||||
|
||||
def test_boto3_direct_raw_model_ids_forward_full_bodies(gateway: Gateway, tmp_path: Path) -> None:
|
||||
direct_raw_id: Final = f"anthropic.claude-direct-{uuid.uuid4().hex[:12]}-v1:0"
|
||||
with wire_server(_full_body_peer) as wire:
|
||||
config_path: Final = tmp_path / "bedrock-direct.yaml"
|
||||
config_path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [{"model_name": "bedrock/*", "litellm_params": {"model": "bedrock/*"}}],
|
||||
"general_settings": {"supported_db_objects": []},
|
||||
}
|
||||
)
|
||||
)
|
||||
aws_environment: Final = {
|
||||
"AWS_ACCESS_KEY_ID": "AKIASCRIPTEDPROVIDER",
|
||||
"AWS_SECRET_ACCESS_KEY": "scripted-secret",
|
||||
"AWS_REGION_NAME": "us-east-1",
|
||||
"AWS_BEDROCK_RUNTIME_ENDPOINT": wire.url,
|
||||
}
|
||||
with (
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
aws_environment,
|
||||
config=config_path,
|
||||
remove_environment=(
|
||||
"AWS_ACCESS_KEY_ID",
|
||||
"AWS_SECRET_ACCESS_KEY",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"AWS_SECURITY_TOKEN",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"AWS_PROFILE",
|
||||
"AWS_DEFAULT_PROFILE",
|
||||
"AWS_SHARED_CREDENTIALS_FILE",
|
||||
"AWS_CONFIG_FILE",
|
||||
"AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
"AWS_ROLE_ARN",
|
||||
"AWS_CONTAINER_CREDENTIALS_FULL_URI",
|
||||
"AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
|
||||
"AWS_REGION",
|
||||
"AWS_DEFAULT_REGION",
|
||||
),
|
||||
workers=1,
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key()
|
||||
runtime: Final = _boto3_runtime(candidate, key)
|
||||
direct_converse: Final = runtime.converse(modelId=direct_raw_id, **_CONVERSE_BODY)
|
||||
direct_invoke: Final = runtime.invoke_model(modelId=direct_raw_id, body=json.dumps(_INVOKE_BODY))
|
||||
direct_converse_stream: Final = _events_read_before_the_upstream_finishes(
|
||||
lambda: runtime.converse_stream(modelId=direct_raw_id, **_CONVERSE_BODY)["stream"]
|
||||
)
|
||||
direct_invoke_stream: Final = runtime.invoke_model_with_response_stream(
|
||||
modelId=direct_raw_id, body=json.dumps(_INVOKE_BODY)
|
||||
)
|
||||
assert {name: direct_converse[name] for name in _CONVERSE_REPLY} == _CONVERSE_REPLY, direct_converse
|
||||
assert json.loads(direct_invoke["body"].read()) == _INVOKE_REPLY, direct_invoke
|
||||
assert list(direct_converse_stream) == [{kind: payload} for kind, payload in _EVENTS], (
|
||||
direct_converse_stream
|
||||
)
|
||||
assert [json.loads(event["chunk"]["bytes"]) for event in direct_invoke_stream["body"]] == [_INVOKE_REPLY]
|
||||
|
||||
received: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"/model/{direct_raw_id}/converse"),
|
||||
("POST", f"/model/{direct_raw_id}/invoke"),
|
||||
("POST", f"/model/{direct_raw_id}/converse-stream"),
|
||||
("POST", f"/model/{direct_raw_id}/invoke-with-response-stream"),
|
||||
], received
|
||||
assert tuple(json.loads(request.body) for request in received) == (
|
||||
_CONVERSE_BODY,
|
||||
_INVOKE_BODY,
|
||||
_CONVERSE_BODY,
|
||||
_INVOKE_BODY,
|
||||
), tuple(request.body for request in received)
|
||||
for request in received:
|
||||
_assert_signed_with_the_deployment_key(request, key)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,276 @@
|
|||
import io
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from email.parser import BytesParser
|
||||
from email.policy import HTTP
|
||||
from itertools import accumulate, dropwhile
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_PNG_A: Final = b"\x89PNG\r\n\x1a\nscripted-image-a"
|
||||
_PNG_B: Final = b"\x89PNG\r\n\x1a\nscripted-image-b"
|
||||
_MASK: Final = b"\x89PNG\r\n\x1a\nscripted-mask"
|
||||
_AUDIO: Final = b"RIFF\x24\x00\x00\x00WAVEscripted-audio"
|
||||
_JSONL: Final = b'{"custom_id":"1","method":"POST","url":"/v1/chat/completions","body":{}}\n'
|
||||
_EDITED: Final = "c2NyaXB0ZWQtZWRpdA=="
|
||||
_TRANSCRIPT: Final = "scripted transcript"
|
||||
_IMAGE_STREAM_FRAMES: Final = (
|
||||
b'event: image_generation.partial_image\ndata: {"type":"image_generation.partial_image"}\n\n',
|
||||
b'event: image_generation.completed\ndata: {"type":"image_generation.completed"}\n\n',
|
||||
)
|
||||
_IMAGE_STREAM_GATE: Final = threading.Event()
|
||||
_FILE_OBJECT: Final = {
|
||||
"id": "file-scripted",
|
||||
"object": "file",
|
||||
"bytes": len(_JSONL),
|
||||
"created_at": 1,
|
||||
"filename": "batch.jsonl",
|
||||
"purpose": "batch",
|
||||
"status": "processed",
|
||||
}
|
||||
|
||||
Part = tuple[str, str | None, str | None, bytes]
|
||||
|
||||
|
||||
class _Image(BaseModel):
|
||||
b64_json: str
|
||||
|
||||
|
||||
class _ImageResponse(BaseModel):
|
||||
created: int
|
||||
data: tuple[_Image, ...]
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
if request.target == "/v1/audio/transcriptions":
|
||||
return Reply(body=_TRANSCRIPT.encode(), content_type="text/plain")
|
||||
if request.target == "/v1/files":
|
||||
return Reply(body=json.dumps(_FILE_OBJECT).encode())
|
||||
if request.target == "/v1/images/edits":
|
||||
parts: Final = _parts(request.headers["content-type"], request.body)
|
||||
if ("stream", None, None, b"true") in parts:
|
||||
return Reply(
|
||||
chunks=_IMAGE_STREAM_FRAMES,
|
||||
content_type="text/event-stream",
|
||||
gate_after_first=_IMAGE_STREAM_GATE,
|
||||
)
|
||||
return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": _EDITED}]}).encode())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wire() -> Iterator[Wire]:
|
||||
with wire_server(_respond) as served:
|
||||
yield served
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained(wire: Wire) -> None:
|
||||
wire.drain()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy(wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("openai-passthrough-multipart")
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["environment_variables"] = {"OPENAI_API_BASE": wire.url, "OPENAI_API_KEY": "scripted"}
|
||||
path: Final = directory / "openai-multipart.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with gateway_from_environment() as rig, owned_proxy(rig, directory, {}, config=path, workers=2) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _parts(content_type: str, body: bytes) -> tuple[Part, ...]:
|
||||
envelope: Final = f"content-type: {content_type}\r\n\r\n".encode() + body
|
||||
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
|
||||
assert parsed.is_multipart(), content_type
|
||||
return tuple(
|
||||
(
|
||||
str(part.get_param("name", header="content-disposition")),
|
||||
part.get_filename(),
|
||||
part.get_content_type() if part.get_filename() is not None else None,
|
||||
part.get_payload(decode=True),
|
||||
)
|
||||
for part in parsed.iter_parts()
|
||||
)
|
||||
|
||||
|
||||
def _only_upstream(wire: Wire, target: str) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in received] == [("POST", target)], received
|
||||
upstream: Final = received[0]
|
||||
assert upstream.headers["authorization"] == "Bearer scripted", upstream.headers
|
||||
return upstream
|
||||
|
||||
|
||||
def _sdk(proxy: Gateway, key: str, prefix: str) -> OpenAI:
|
||||
return OpenAI(
|
||||
api_key=key,
|
||||
base_url=f"{str(proxy.client.base_url).rstrip('/')}{prefix}/v1",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=30, trust_env=False),
|
||||
)
|
||||
|
||||
|
||||
def _image_edit_files(
|
||||
stream: bool, image_field: str, mask_field: str
|
||||
) -> list[tuple[str, tuple[str | None, str | bytes] | tuple[str, bytes, str]]]:
|
||||
return [
|
||||
("model", (None, "gpt-image-1")),
|
||||
("prompt", (None, "make them match")),
|
||||
("n", (None, "2")),
|
||||
("stream", (None, "true" if stream else "false")),
|
||||
(image_field, ("a.png", _PNG_A, "image/png")),
|
||||
(image_field, ("b.png", _PNG_B, "image/png")),
|
||||
(mask_field, ("m.png", _MASK, "image/png")),
|
||||
]
|
||||
|
||||
|
||||
def _expected_image_edit_parts(stream: bool, image_field: str, mask_field: str) -> tuple[Part, ...]:
|
||||
return (
|
||||
("model", None, None, b"gpt-image-1"),
|
||||
("prompt", None, None, b"make them match"),
|
||||
("n", None, None, b"2"),
|
||||
("stream", None, None, b"true" if stream else b"false"),
|
||||
(image_field, "a.png", "image/png", _PNG_A),
|
||||
(image_field, "b.png", "image/png", _PNG_B),
|
||||
(mask_field, "m.png", "image/png", _MASK),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_field", "mask_field"),
|
||||
[("image[]", "mask"), ("image", "mask"), ("image[]", "mask[]")],
|
||||
ids=["bracketed-images", "plain-repeated-images", "bracketed-mask"],
|
||||
)
|
||||
def test_curl_image_edit_forwards_repeated_image_parts_mask_and_string_typed_fields(
|
||||
proxy: Gateway, wire: Wire, image_field: str, mask_field: str
|
||||
) -> None:
|
||||
with proxy.scenario() as scenario:
|
||||
response: Final = proxy.client.post(
|
||||
"/openai/v1/images/edits",
|
||||
headers={"Authorization": f"Bearer {scenario.key()}"},
|
||||
files=_image_edit_files(False, image_field, mask_field),
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _ImageResponse.model_validate_json(response.content).data == (_Image(b64_json=_EDITED),), response.text
|
||||
upstream: Final = _only_upstream(wire, "/v1/images/edits")
|
||||
assert _parts(upstream.headers["content-type"], upstream.body) == _expected_image_edit_parts(
|
||||
False, image_field, mask_field
|
||||
), upstream.body
|
||||
|
||||
|
||||
def test_curl_streaming_image_edit_relays_each_sse_frame_as_the_upstream_sends_it(proxy: Gateway, wire: Wire) -> None:
|
||||
pytest.skip("BUG: multipart image edits with stream=true do not relay the SSE response incrementally")
|
||||
_IMAGE_STREAM_GATE.clear()
|
||||
with (
|
||||
proxy.scenario() as scenario,
|
||||
proxy.client.stream(
|
||||
"POST",
|
||||
"/openai/v1/images/edits",
|
||||
headers={"Authorization": f"Bearer {scenario.key()}"},
|
||||
files=_image_edit_files(True, "image[]", "mask"),
|
||||
) as response,
|
||||
):
|
||||
raw: Final = iter(response.iter_raw())
|
||||
try:
|
||||
assert response.status_code == 200, response.status_code
|
||||
assert response.headers.get("content-type") == "text/event-stream", response.headers
|
||||
first_frame: Final = next(
|
||||
dropwhile(
|
||||
lambda buffered: b"\n\n" not in buffered,
|
||||
accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""),
|
||||
)
|
||||
)
|
||||
assert first_frame == _IMAGE_STREAM_FRAMES[0], first_frame
|
||||
finally:
|
||||
_IMAGE_STREAM_GATE.set()
|
||||
streamed: Final = first_frame + b"".join(raw)
|
||||
assert streamed == b"".join(_IMAGE_STREAM_FRAMES), streamed
|
||||
upstream: Final = _only_upstream(wire, "/v1/images/edits")
|
||||
assert _parts(upstream.headers["content-type"], upstream.body) == _expected_image_edit_parts(
|
||||
True, "image[]", "mask"
|
||||
)
|
||||
|
||||
|
||||
def test_sdk_image_edit_with_two_images_reaches_upstream_as_two_file_parts(proxy: Gateway, wire: Wire) -> None:
|
||||
with proxy.scenario() as scenario, _sdk(proxy, scenario.key(), "/openai") as sdk:
|
||||
edited: Final = sdk.images.edit(
|
||||
model="gpt-image-1",
|
||||
image=[("a.png", io.BytesIO(_PNG_A), "image/png"), ("b.png", io.BytesIO(_PNG_B), "image/png")],
|
||||
mask=("m.png", io.BytesIO(_MASK), "image/png"),
|
||||
prompt="make them match",
|
||||
n=2,
|
||||
)
|
||||
assert [image.b64_json for image in edited.data or []] == [_EDITED], edited
|
||||
upstream: Final = _only_upstream(wire, "/v1/images/edits")
|
||||
assert _parts(upstream.headers["content-type"], upstream.body) == (
|
||||
("prompt", None, None, b"make them match"),
|
||||
("model", None, None, b"gpt-image-1"),
|
||||
("n", None, None, b"2"),
|
||||
("image[]", "a.png", "image/png", _PNG_A),
|
||||
("image[]", "b.png", "image/png", _PNG_B),
|
||||
("mask", "m.png", "image/png", _MASK),
|
||||
), upstream.body
|
||||
|
||||
|
||||
def test_sdk_transcription_and_file_upload_forward_the_file_and_its_text_fields(proxy: Gateway, wire: Wire) -> None:
|
||||
with proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
with _sdk(proxy, key, "/openai") as sdk:
|
||||
transcript: Final = sdk.audio.transcriptions.create(
|
||||
model="whisper-1", file=("clip.wav", io.BytesIO(_AUDIO), "audio/wav"), response_format="text"
|
||||
)
|
||||
assert transcript == _TRANSCRIPT, transcript
|
||||
transcription: Final = _only_upstream(wire, "/v1/audio/transcriptions")
|
||||
assert _parts(transcription.headers["content-type"], transcription.body) == (
|
||||
("model", None, None, b"whisper-1"),
|
||||
("response_format", None, None, b"text"),
|
||||
("file", "clip.wav", "audio/wav", _AUDIO),
|
||||
), transcription.body
|
||||
with _sdk(proxy, key, "/openai_passthrough") as sdk:
|
||||
uploaded: Final = sdk.files.create(
|
||||
file=("batch.jsonl", io.BytesIO(_JSONL), "application/jsonl"), purpose="batch"
|
||||
)
|
||||
assert uploaded.model_dump(exclude_unset=True) == _FILE_OBJECT, uploaded
|
||||
upload: Final = _only_upstream(wire, "/v1/files")
|
||||
assert _parts(upload.headers["content-type"], upload.body) == (
|
||||
("purpose", None, None, b"batch"),
|
||||
("file", "batch.jsonl", "application/jsonl", _JSONL),
|
||||
), upload.body
|
||||
|
||||
|
||||
def test_file_less_multipart_form_stays_multipart_with_every_repeated_field(proxy: Gateway, wire: Wire) -> None:
|
||||
with proxy.scenario() as scenario:
|
||||
response: Final = proxy.client.post(
|
||||
"/openai/v1/images/edits",
|
||||
headers={"Authorization": f"Bearer {scenario.key()}"},
|
||||
files=[
|
||||
("model", (None, "gpt-image-1")),
|
||||
("prompt", (None, "no file")),
|
||||
("image_url[]", (None, "https://example.com/a.png")),
|
||||
("image_url[]", (None, "https://example.com/b.png")),
|
||||
],
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
parsed: Final = _ImageResponse.model_validate_json(response.content)
|
||||
assert parsed == _ImageResponse(created=1, data=(_Image(b64_json=_EDITED),)), response.text
|
||||
upstream: Final = _only_upstream(wire, "/v1/images/edits")
|
||||
assert upstream.headers["content-type"].startswith("multipart/form-data; boundary="), upstream.headers
|
||||
assert _parts(upstream.headers["content-type"], upstream.body) == (
|
||||
("model", None, None, b"gpt-image-1"),
|
||||
("prompt", None, None, b"no file"),
|
||||
("image_url[]", None, None, b"https://example.com/a.png"),
|
||||
("image_url[]", None, None, b"https://example.com/b.png"),
|
||||
), upstream.body
|
||||
516
tests/integration/providers/test_vertex_passthrough_wire.py
Normal file
516
tests/integration/providers/test_vertex_passthrough_wire.py
Normal file
|
|
@ -0,0 +1,516 @@
|
|||
import contextlib
|
||||
import datetime
|
||||
import json
|
||||
import socket
|
||||
import socketserver
|
||||
import ssl
|
||||
import threading
|
||||
from collections.abc import Generator, Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, dropwhile
|
||||
from pathlib import Path
|
||||
from queue import Empty, SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
from cryptography.x509.oid import NameOID
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.oauth2.credentials import Credentials
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_PROJECT: Final = "scripted-project"
|
||||
_TOKEN: Final = "scripted-vertex-token"
|
||||
_CALLER_GOOGLE_TOKEN: Final = "caller-google-token"
|
||||
_GEMINI_ROUTER_NAME: Final = "vertex-gemini-router"
|
||||
_REGIONAL: Final = "us-central1-aiplatform.googleapis.com:443"
|
||||
_GLOBAL: Final = "aiplatform.googleapis.com:443"
|
||||
_BYO: Final = "us-east5-aiplatform.googleapis.com:443"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_GEMINI_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"candidates": [{"content": {"parts": [{"text": "scripted vertex"}], "role": "model"}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5},
|
||||
}
|
||||
_GEMINI_SSE: Final = (f"data: {json.dumps(_GEMINI_REPLY)}\r\n\r\n".encode(),)
|
||||
_GEMINI_BODY: Final[dict[str, JsonValue]] = {
|
||||
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "be terse"}]},
|
||||
"generationConfig": {"temperature": 0.2, "maxOutputTokens": 32},
|
||||
}
|
||||
_GEMINI_SDK_BODY: Final[dict[str, JsonValue]] = {
|
||||
**_GEMINI_BODY,
|
||||
"systemInstruction": {"parts": [{"text": "be terse"}], "role": "user"},
|
||||
}
|
||||
_CLAUDE_REPLY: Final[dict[str, JsonValue]] = {
|
||||
"id": "msg_scripted",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "scripted claude"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 2},
|
||||
}
|
||||
_CLAUDE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{"type": "message_start", "message": {**_CLAUDE_REPLY, "content": [], "stop_reason": None}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "scripted claude"}},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
_CLAUDE_SSE: Final = tuple(
|
||||
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _CLAUDE_EVENTS
|
||||
)
|
||||
|
||||
|
||||
_VERTEX_STREAM_GATE: Final = threading.Event()
|
||||
_VERTEX_STREAM_GATE.set()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Tunnel:
|
||||
url: str
|
||||
authorities: SimpleQueue[str]
|
||||
|
||||
def drain(self) -> list[str]:
|
||||
drained: Final[list[str]] = []
|
||||
with contextlib.suppress(Empty):
|
||||
while True:
|
||||
drained.append(self.authorities.get_nowait())
|
||||
return drained
|
||||
|
||||
|
||||
def _tls_context(directory: Path) -> ssl.SSLContext:
|
||||
key: Final = ec.generate_private_key(ec.SECP256R1())
|
||||
now: Final = datetime.datetime.now(datetime.UTC)
|
||||
name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "aiplatform.googleapis.com")])
|
||||
certificate: Final = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(name)
|
||||
.issuer_name(name)
|
||||
.public_key(key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(now - datetime.timedelta(days=1))
|
||||
.not_valid_after(now + datetime.timedelta(days=1))
|
||||
.sign(key, hashes.SHA256())
|
||||
)
|
||||
certificate_file: Final = directory / "vertex.pem"
|
||||
key_file: Final = directory / "vertex.key"
|
||||
certificate_file.write_bytes(certificate.public_bytes(serialization.Encoding.PEM))
|
||||
key_file.write_bytes(
|
||||
key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
|
||||
)
|
||||
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(certificate_file, key_file)
|
||||
return context
|
||||
|
||||
|
||||
def _pipe(source: socket.socket, sink: socket.socket) -> None:
|
||||
with contextlib.suppress(OSError):
|
||||
for chunk in iter(lambda: source.recv(65536), b""):
|
||||
sink.sendall(chunk)
|
||||
with contextlib.suppress(OSError):
|
||||
sink.shutdown(socket.SHUT_WR)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _vertex_tunnel(destination: Wire) -> Generator[_Tunnel, None, None]:
|
||||
authorities: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
destination_port: Final = int(destination.url.rsplit(":", 1)[1])
|
||||
|
||||
class Tunnel(socketserver.StreamRequestHandler):
|
||||
rbufsize = 0
|
||||
request: socket.socket
|
||||
|
||||
def handle(self) -> None:
|
||||
authority: Final = self.rfile.readline().decode().split()[1]
|
||||
while self.rfile.readline() not in (b"\r\n", b""):
|
||||
pass
|
||||
authorities.put(authority)
|
||||
if authority not in (_REGIONAL, _GLOBAL, _BYO):
|
||||
self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n")
|
||||
return
|
||||
self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n")
|
||||
self.request.settimeout(10)
|
||||
with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream:
|
||||
outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream))
|
||||
outbound.start()
|
||||
_pipe(upstream, self.request)
|
||||
outbound.join(timeout=12)
|
||||
|
||||
with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Tunnel) as server:
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
|
||||
thread.start()
|
||||
try:
|
||||
yield _Tunnel(f"http://127.0.0.1:{server.server_address[1]}", authorities)
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join(timeout=6)
|
||||
|
||||
|
||||
def _vertex_peer(request: Request) -> Reply:
|
||||
if "enerateContent" in request.target:
|
||||
if ":streamGenerateContent" in request.target:
|
||||
return Reply(chunks=_GEMINI_SSE, content_type="text/event-stream")
|
||||
return Reply(body=json.dumps(_GEMINI_REPLY).encode())
|
||||
if ":streamRawPredict" in request.target or request.target.endswith(":rawPredict?alt=sse"):
|
||||
return Reply(chunks=_CLAUDE_SSE, content_type="text/event-stream", gate_after_first=_VERTEX_STREAM_GATE)
|
||||
return Reply(body=json.dumps(_CLAUDE_REPLY).encode())
|
||||
|
||||
|
||||
def _token_peer(request: Request) -> Reply:
|
||||
assert request.target == "/_oauth/token", request.target
|
||||
return Reply(body=json.dumps({"access_token": _TOKEN, "expires_in": 3600, "token_type": "Bearer"}).encode())
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
proxy: Gateway
|
||||
vertex: Wire
|
||||
tokens: Wire
|
||||
tunnel: _Tunnel
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("vertex-passthrough")
|
||||
with (
|
||||
wire_server(_vertex_peer, tls=_tls_context(directory)) as vertex,
|
||||
wire_server(_token_peer) as tokens,
|
||||
_vertex_tunnel(vertex) as tunnel,
|
||||
):
|
||||
credentials: Final = service_account_json(_PROJECT, tokens.url)
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["model_list"] = [
|
||||
*config.get("model_list", []),
|
||||
{
|
||||
"model_name": _GEMINI_ROUTER_NAME,
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-flash",
|
||||
"vertex_project": _PROJECT,
|
||||
"vertex_location": "us-central1",
|
||||
"vertex_credentials": credentials,
|
||||
"use_in_pass_through": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-sonnet-4-5",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/claude-sonnet-4-5",
|
||||
"vertex_project": _PROJECT,
|
||||
"vertex_location": "global",
|
||||
"vertex_credentials": credentials,
|
||||
"use_in_pass_through": True,
|
||||
},
|
||||
},
|
||||
]
|
||||
path: Final = directory / "vertex-passthrough.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
environment: Final = {
|
||||
"HTTPS_PROXY": tunnel.url,
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
"AIOHTTP_TRUST_ENV": "True",
|
||||
"SSL_VERIFY": "False",
|
||||
}
|
||||
with (
|
||||
gateway_from_environment() as base,
|
||||
owned_proxy(base, directory, environment, config=path, workers=2) as proxy,
|
||||
):
|
||||
yield _Rig(proxy, vertex, tokens, tunnel)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _drained(rig: _Rig) -> None:
|
||||
rig.vertex.drain()
|
||||
rig.tunnel.drain()
|
||||
|
||||
|
||||
def _assert_minted_token_and_no_virtual_key(upstream: Request, key: str) -> None:
|
||||
assert upstream.headers["authorization"] == f"Bearer {_TOKEN}", upstream.headers
|
||||
assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers
|
||||
assert key not in upstream.target, upstream.target
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"])
|
||||
def test_google_genai_vertex_calls_reach_the_router_deployment_with_a_minted_token(rig: _Rig, alias: str) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
client: Final = genai.Client(
|
||||
vertexai=True,
|
||||
project=_PROJECT,
|
||||
location="us-central1",
|
||||
credentials=Credentials(token=key),
|
||||
http_options=types.HttpOptions(
|
||||
base_url=f"{str(rig.proxy.client.base_url).rstrip('/')}{alias}",
|
||||
client_args={"event_hooks": {"request": [sent.append]}, "timeout": 30, "trust_env": False},
|
||||
),
|
||||
)
|
||||
generated: Final = client.models.generate_content(
|
||||
model=_GEMINI_ROUTER_NAME,
|
||||
contents="hi",
|
||||
config=types.GenerateContentConfig(
|
||||
system_instruction="be terse",
|
||||
temperature=0.2,
|
||||
max_output_tokens=32,
|
||||
),
|
||||
)
|
||||
streamed: Final = "".join(
|
||||
chunk.text or ""
|
||||
for chunk in client.models.generate_content_stream(
|
||||
model=_GEMINI_ROUTER_NAME,
|
||||
contents="hi",
|
||||
config=types.GenerateContentConfig(
|
||||
system_instruction="be terse",
|
||||
temperature=0.2,
|
||||
max_output_tokens=32,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert generated.text == "scripted vertex", generated
|
||||
assert streamed == "scripted vertex", streamed
|
||||
received: Final = rig.vertex.drain()
|
||||
model_path: Final = f"/v1beta1/projects/{_PROJECT}/locations/us-central1/publishers/google/models"
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"{model_path}/gemini-2.5-flash:generateContent"),
|
||||
("POST", f"{model_path}/gemini-2.5-flash:streamGenerateContent?alt=sse"),
|
||||
], received
|
||||
assert rig.tunnel.drain() == [_REGIONAL, _REGIONAL]
|
||||
assert len(sent) == len(received), sent
|
||||
for upstream, sdk_request in zip(received, sent, strict=True):
|
||||
sdk_body: Final = _JSON_OBJECT.validate_json(sdk_request.content)
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == sdk_body, upstream.body
|
||||
assert sdk_body == _GEMINI_SDK_BODY, upstream.body
|
||||
_assert_minted_token_and_no_virtual_key(upstream, key)
|
||||
|
||||
|
||||
def test_vertex_abbreviated_beta_route_uses_router_project_and_location(rig: _Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
response: Final = rig.proxy.client.post(
|
||||
f"/vertex-ai/v1beta1/publishers/google/models/{_GEMINI_ROUTER_NAME}:generateContent",
|
||||
json=_GEMINI_BODY,
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _GEMINI_REPLY, response.text
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
(
|
||||
"POST",
|
||||
f"/v1beta1/projects/{_PROJECT}/locations/us-central1/publishers/google/models/"
|
||||
"gemini-2.5-flash:generateContent",
|
||||
)
|
||||
], received
|
||||
assert rig.tunnel.drain() == [_REGIONAL]
|
||||
assert _JSON_OBJECT.validate_json(received[0].body) == _GEMINI_BODY, received[0].body
|
||||
_assert_minted_token_and_no_virtual_key(received[0], key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"])
|
||||
def test_vertex_rest_spellings_mint_token_and_add_sse_query_upstream(rig: _Rig, alias: str) -> None:
|
||||
model_path: Final = f"/v1/projects/{_PROJECT}/locations/us-central1/publishers/google/models"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
headers: Final = {"Authorization": f"Bearer {key}"}
|
||||
generated: Final = rig.proxy.client.post(
|
||||
f"{alias}{model_path}/{_GEMINI_ROUTER_NAME}:generateContent", json=_GEMINI_BODY, headers=headers
|
||||
)
|
||||
assert generated.status_code == 200, generated.text
|
||||
assert generated.json() == _GEMINI_REPLY, generated.text
|
||||
with rig.proxy.client.stream(
|
||||
"POST",
|
||||
f"{alias}{model_path}/{_GEMINI_ROUTER_NAME}:streamGenerateContent",
|
||||
json=_GEMINI_BODY,
|
||||
headers=headers,
|
||||
) as streamed_response:
|
||||
streamed: Final = streamed_response.read()
|
||||
assert streamed_response.status_code == 200, streamed
|
||||
assert streamed == b"".join(_GEMINI_SSE), streamed
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"{model_path}/gemini-2.5-flash:generateContent"),
|
||||
("POST", f"{model_path}/gemini-2.5-flash:streamGenerateContent?alt=sse"),
|
||||
], received
|
||||
assert rig.tunnel.drain() == [_REGIONAL, _REGIONAL]
|
||||
for upstream in received:
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == _GEMINI_BODY, upstream.body
|
||||
_assert_minted_token_and_no_virtual_key(upstream, key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"])
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["rawPredict", "streamRawPredict"])
|
||||
def test_anthropic_vertex_sdk_partner_call_on_the_global_location_reaches_upstream_intact(
|
||||
rig: _Rig, alias: str, stream: bool
|
||||
) -> None:
|
||||
sent: Final[list[httpx.Request]] = []
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
client: Final = anthropic.AnthropicVertex(
|
||||
base_url=f"{str(rig.proxy.client.base_url).rstrip('/')}{alias}/v1",
|
||||
region="global",
|
||||
project_id=_PROJECT,
|
||||
access_token=key,
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}),
|
||||
)
|
||||
arguments: Final[dict[str, JsonValue]] = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 64,
|
||||
"system": "be terse",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
if stream:
|
||||
with client.messages.stream(**arguments) as events: # type: ignore[arg-type]
|
||||
text = events.get_final_text()
|
||||
else:
|
||||
text = "".join(
|
||||
block.text
|
||||
for block in client.messages.create(**arguments).content # type: ignore[arg-type]
|
||||
if block.type == "text"
|
||||
)
|
||||
client.close()
|
||||
assert text == "scripted claude"
|
||||
action: Final = "streamRawPredict" if stream else "rawPredict"
|
||||
model_path: Final = f"/v1/projects/{_PROJECT}/locations/global/publishers/anthropic/models/claude-sonnet-4-5"
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"{model_path}:{action}" + ("?alt=sse" if stream else ""))
|
||||
], received
|
||||
assert rig.tunnel.drain() == [_GLOBAL]
|
||||
upstream: Final = received[0]
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent[0].content), upstream.body
|
||||
_assert_minted_token_and_no_virtual_key(upstream, key)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"])
|
||||
@pytest.mark.parametrize("action", ["rawPredict", "streamRawPredict"])
|
||||
def test_vertex_partner_stream_body_adds_the_sse_query_and_relays_each_frame_before_the_upstream_finishes(
|
||||
rig: _Rig, alias: str, action: str
|
||||
) -> None:
|
||||
model_path: Final = f"/v1/projects/{_PROJECT}/locations/global/publishers/anthropic/models/claude-sonnet-4-5"
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"max_tokens": 64,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
_VERTEX_STREAM_GATE.clear()
|
||||
with rig.proxy.client.stream(
|
||||
"POST", f"{alias}{model_path}:{action}", json=body, headers={"Authorization": f"Bearer {key}"}
|
||||
) as response:
|
||||
raw: Final = iter(response.iter_raw())
|
||||
try:
|
||||
assert response.status_code == 200, response.status_code
|
||||
first_frame: Final = next(
|
||||
dropwhile(
|
||||
lambda buffered: b"\n\n" not in buffered,
|
||||
accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""),
|
||||
),
|
||||
b"",
|
||||
)
|
||||
assert first_frame == _CLAUDE_SSE[0], first_frame
|
||||
finally:
|
||||
_VERTEX_STREAM_GATE.set()
|
||||
streamed: Final = first_frame + b"".join(raw)
|
||||
assert streamed == b"".join(_CLAUDE_SSE), streamed
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [
|
||||
("POST", f"{model_path}:{action}?alt=sse")
|
||||
], received
|
||||
assert rig.tunnel.drain() == [_GLOBAL]
|
||||
assert _JSON_OBJECT.validate_json(received[0].body) == body, received[0].body
|
||||
_assert_minted_token_and_no_virtual_key(received[0], key)
|
||||
|
||||
|
||||
def test_credential_less_vertex_project_forwards_the_caller_google_token_and_strips_the_virtual_key(
|
||||
rig: _Rig,
|
||||
) -> None:
|
||||
model_path: Final = "/v1/projects/byo-project/locations/us-east5/publishers/anthropic/models/claude-byo"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
response: Final = rig.proxy.client.post(
|
||||
f"/vertex_ai{model_path}:rawPredict",
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {key}", "x-goog-api-key": _CALLER_GOOGLE_TOKEN},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _CLAUDE_REPLY, response.text
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [("POST", f"{model_path}:rawPredict")], (
|
||||
received
|
||||
)
|
||||
assert rig.tunnel.drain() == [_BYO]
|
||||
upstream: Final = received[0]
|
||||
assert upstream.headers["x-goog-api-key"] == _CALLER_GOOGLE_TOKEN, upstream.headers
|
||||
assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers
|
||||
assert "authorization" not in upstream.headers, upstream.headers
|
||||
assert "x-litellm-api-key" not in upstream.headers, upstream.headers
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"])
|
||||
def test_credential_less_vertex_project_forwards_the_caller_bearer_token_when_the_virtual_key_rides_in_x_litellm_api_key(
|
||||
rig: _Rig, alias: str
|
||||
) -> None:
|
||||
model_path: Final = "/v1/projects/byo-project/locations/us-east5/publishers/anthropic/models/claude-byo"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
second_key: Final = scenario.key()
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
response: Final = rig.proxy.client.post(
|
||||
f"{alias}{model_path}:rawPredict",
|
||||
json=body,
|
||||
headers={
|
||||
"x-litellm-api-key": f"Bearer {key}",
|
||||
"Authorization": f"Bearer {_CALLER_GOOGLE_TOKEN}",
|
||||
"api-key": second_key,
|
||||
"x-caller-echo": key,
|
||||
"x-caller-trace": "trace-1",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _CLAUDE_REPLY, response.text
|
||||
received: Final = rig.vertex.drain()
|
||||
assert [(request.method, request.target) for request in received] == [("POST", f"{model_path}:rawPredict")], (
|
||||
received
|
||||
)
|
||||
assert rig.tunnel.drain() == [_BYO]
|
||||
upstream: Final = received[0]
|
||||
assert upstream.headers["authorization"] == f"Bearer {_CALLER_GOOGLE_TOKEN}", upstream.headers
|
||||
assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers
|
||||
assert {name: value for name, value in upstream.headers.items() if second_key in value} == {}, upstream.headers
|
||||
assert "x-litellm-api-key" not in upstream.headers, upstream.headers
|
||||
assert "x-goog-api-key" not in upstream.headers, upstream.headers
|
||||
assert upstream.headers["x-caller-trace"] == "trace-1", upstream.headers
|
||||
assert key not in upstream.target, upstream.target
|
||||
assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body
|
||||
|
|
@ -1,9 +1,12 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from email.parser import BytesParser
|
||||
from email.policy import HTTP
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value
|
||||
|
|
@ -419,3 +422,362 @@ def test_anthropic_passthrough_spend_row_carries_key_team_project_tags_and_spend
|
|||
row
|
||||
)
|
||||
assert _spend_logs_metadata(row) == {"cost_center": marker, "team_field": marker}, row
|
||||
|
||||
|
||||
_CONTRACT_TOKEN_ENV: Final = "INTEGRATION_PASSTHROUGH_CONTRACT_TOKEN"
|
||||
_CONTRACT_TOKEN: Final = "scripted-configured-endpoint-token"
|
||||
_ENDPOINTS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
|
||||
|
||||
def _received_reply(request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"received": request.target}).encode())
|
||||
|
||||
|
||||
def _create_endpoint(gateway: Gateway, scenario: Scenario, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
response: Final = gateway.request("POST", "/config/pass_through_endpoint", dict(body))
|
||||
assert response.status_code == 200, response.text
|
||||
created: Final = _ENDPOINTS.validate_python(response.json()["endpoints"])
|
||||
assert len(created) == 1, response.text
|
||||
endpoint_id: Final = str(created[0]["id"])
|
||||
scenario.cleanups.callback(
|
||||
lambda: gateway.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id})
|
||||
)
|
||||
return created[0]
|
||||
|
||||
|
||||
def _readback(gateway: Gateway, endpoint_id: str) -> list[dict[str, JsonValue]]:
|
||||
response: Final = gateway.request("GET", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id})
|
||||
assert response.status_code == 200, response.text
|
||||
return _ENDPOINTS.validate_python(response.json()["endpoints"])
|
||||
|
||||
|
||||
def _stored_endpoint_ids() -> list[JsonValue]:
|
||||
rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("general_settings",))
|
||||
assert len(rows) == 1, rows
|
||||
settings: Final = object_value(rows[0]["param_value"])
|
||||
stored: Final = _ENDPOINTS.validate_python(settings.get("pass_through_endpoints") or [])
|
||||
return [endpoint["id"] for endpoint in stored]
|
||||
|
||||
|
||||
def _stored_endpoint(endpoint_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("general_settings",))
|
||||
assert len(rows) == 1, rows
|
||||
settings: Final = object_value(rows[0]["param_value"])
|
||||
stored: Final = _ENDPOINTS.validate_python(settings.get("pass_through_endpoints") or [])
|
||||
matches: Final = [endpoint for endpoint in stored if endpoint.get("id") == endpoint_id]
|
||||
assert len(matches) == 1, matches
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _query(request: Request) -> tuple[str, tuple[tuple[str, str], ...]]:
|
||||
path, _, query = request.target.partition("?")
|
||||
return path, tuple(parse_qsl(query, keep_blank_values=True))
|
||||
|
||||
|
||||
def _multipart_parts(content_type: str, body: bytes) -> list[tuple[str, str | None, str | None, bytes]]:
|
||||
parsed: Final = BytesParser(policy=HTTP).parsebytes(f"content-type: {content_type}\r\n\r\n".encode() + body)
|
||||
assert parsed.is_multipart(), content_type
|
||||
return [
|
||||
(
|
||||
str(part.get_param("name", header="content-disposition")),
|
||||
part.get_filename(),
|
||||
part.get("content-type"),
|
||||
part.get_payload(decode=True),
|
||||
)
|
||||
for part in parsed.iter_parts()
|
||||
]
|
||||
|
||||
|
||||
def test_configured_endpoint_forwards_json_multipart_and_query_and_an_update_keeps_omitted_fields(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
path: Final = f"/integration-contract-{marker}"
|
||||
subpath_root: Final = f"/integration-contract-subpath-{marker}"
|
||||
with (
|
||||
wire_server(_received_reply) as wire,
|
||||
owned_proxy(gateway, tmp_path, {_CONTRACT_TOKEN_ENV: _CONTRACT_TOKEN}) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
created: Final = _create_endpoint(
|
||||
candidate,
|
||||
scenario,
|
||||
{
|
||||
"path": path,
|
||||
"target": f"{wire.url}/base",
|
||||
"headers": {"Authorization": f"Bearer os.environ/{_CONTRACT_TOKEN_ENV}", "x-static": "static-v1"},
|
||||
"methods": ["POST"],
|
||||
"default_query_params": {"api-version": "2024-01-01"},
|
||||
"auth": False,
|
||||
},
|
||||
)
|
||||
_create_endpoint(
|
||||
candidate,
|
||||
scenario,
|
||||
{"path": subpath_root, "target": f"{wire.url}/nested", "include_subpath": True, "auth": True},
|
||||
)
|
||||
endpoint_id: Final = str(created["id"])
|
||||
body: Final[dict[str, JsonValue]] = {"input": marker, "options": {"trace": marker, "stream": False}, "n": 2}
|
||||
|
||||
sent_json: Final = candidate.client.post(path, params={"trace": "1"}, json=body)
|
||||
assert sent_json.status_code == 200, sent_json.text
|
||||
assert sent_json.json() == {"received": "/base?api-version=2024-01-01&trace=1"}, sent_json.text
|
||||
overridden: Final = candidate.client.post(path, params={"api-version": "2025-02-02"}, json=body)
|
||||
assert overridden.status_code == 200, overridden.text
|
||||
form: Final = candidate.client.post(
|
||||
path,
|
||||
files=[
|
||||
("tag", (None, "first")),
|
||||
("tag", (None, "second")),
|
||||
("file", ("one.bin", b"\x00one", "application/octet-stream")),
|
||||
("file", ("two.bin", b"\x00two", "application/octet-stream")),
|
||||
],
|
||||
)
|
||||
assert form.status_code == 200, form.text
|
||||
refused_verb: Final = candidate.client.get(path)
|
||||
assert refused_verb.status_code == 405, refused_verb.text
|
||||
nested: Final = candidate.request("POST", f"{subpath_root}/sub/route", body, params={"q": "1"})
|
||||
assert nested.status_code == 200, nested.text
|
||||
before: Final = wire.drain()
|
||||
assert [(request.method, *_query(request)) for request in before] == [
|
||||
("POST", "/base", (("api-version", "2024-01-01"), ("trace", "1"))),
|
||||
("POST", "/base", (("api-version", "2025-02-02"),)),
|
||||
("POST", "/base", (("api-version", "2024-01-01"),)),
|
||||
("POST", "/nested/sub/route", (("q", "1"),)),
|
||||
], before
|
||||
for request in before[:3]:
|
||||
assert request.headers["authorization"] == f"Bearer {_CONTRACT_TOKEN}", request.headers
|
||||
assert request.headers["x-static"] == "static-v1", request.headers
|
||||
assert [json.loads(request.body) for request in (before[0], before[1], before[3])] == [body] * 3, before
|
||||
assert _multipart_parts(before[2].headers["content-type"], before[2].body) == [
|
||||
("tag", None, None, b"first"),
|
||||
("tag", None, None, b"second"),
|
||||
("file", "one.bin", "application/octet-stream", b"\x00one"),
|
||||
("file", "two.bin", "application/octet-stream", b"\x00two"),
|
||||
], before[2].body
|
||||
|
||||
retargeted: Final = candidate.request(
|
||||
"POST", f"/config/pass_through_endpoint/{endpoint_id}", {"path": path, "target": f"{wire.url}/v2"}
|
||||
)
|
||||
assert retargeted.status_code == 200, retargeted.text
|
||||
expected: Final = {**created, "target": f"{wire.url}/v2"}
|
||||
assert _ENDPOINTS.validate_python(retargeted.json()["endpoints"]) == [expected], retargeted.text
|
||||
retargeted_readback: Final = _readback(candidate, endpoint_id)
|
||||
assert retargeted_readback == [expected]
|
||||
assert _stored_endpoint(endpoint_id) == {
|
||||
key: value for key, value in retargeted_readback[0].items() if key != "is_from_config"
|
||||
}
|
||||
after_retarget: Final = candidate.client.post(path, json=body)
|
||||
assert after_retarget.status_code == 200, after_retarget.text
|
||||
assert after_retarget.json() == {"received": "/v2?api-version=2024-01-01"}, after_retarget.text
|
||||
retargeted_request: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in retargeted_request] == [
|
||||
("POST", "/v2?api-version=2024-01-01")
|
||||
], retargeted_request
|
||||
assert retargeted_request[0].headers["authorization"] == f"Bearer {_CONTRACT_TOKEN}", retargeted_request
|
||||
|
||||
locked: Final = candidate.request(
|
||||
"POST",
|
||||
f"/config/pass_through_endpoint/{endpoint_id}",
|
||||
{"path": path, "target": f"{wire.url}/v2", "auth": True, "headers": {"x-static": "static-v2"}},
|
||||
)
|
||||
assert locked.status_code == 200, locked.text
|
||||
locked_expected: Final = {**expected, "auth": True, "headers": {"x-static": "static-v2"}}
|
||||
locked_readback: Final = _readback(candidate, endpoint_id)
|
||||
assert locked_readback == [locked_expected]
|
||||
assert _stored_endpoint(endpoint_id) == {
|
||||
key: value for key, value in locked_readback[0].items() if key != "is_from_config"
|
||||
}
|
||||
anonymous: Final = candidate.client.post(path, json=body)
|
||||
assert anonymous.status_code == 401, anonymous.text
|
||||
assert wire.drain() == ()
|
||||
keyed: Final = candidate.request("POST", path, body)
|
||||
assert keyed.status_code == 200, keyed.text
|
||||
locked_request: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in locked_request] == [
|
||||
("POST", "/v2?api-version=2024-01-01")
|
||||
], locked_request
|
||||
assert locked_request[0].headers["x-static"] == "static-v2", locked_request[0].headers
|
||||
assert "authorization" not in locked_request[0].headers, locked_request[0].headers
|
||||
refused_after_update: Final = candidate.request("GET", path)
|
||||
assert refused_after_update.status_code == 405, refused_after_update.text
|
||||
keyed_form: Final = candidate.client.post(
|
||||
path,
|
||||
headers={"Authorization": f"Bearer {candidate.key}"},
|
||||
files=[
|
||||
("tag", (None, "first")),
|
||||
("tag", (None, "second")),
|
||||
("file", ("one.bin", b"\x00one", "application/octet-stream")),
|
||||
("file", ("two.bin", b"\x00two", "application/octet-stream")),
|
||||
],
|
||||
)
|
||||
assert keyed_form.status_code == 200, keyed_form.text
|
||||
locked_form: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in locked_form] == [
|
||||
("POST", "/v2?api-version=2024-01-01")
|
||||
], locked_form
|
||||
assert _multipart_parts(locked_form[0].headers["content-type"], locked_form[0].body) == [
|
||||
("tag", None, None, b"first"),
|
||||
("tag", None, None, b"second"),
|
||||
("file", "one.bin", "application/octet-stream", b"\x00one"),
|
||||
("file", "two.bin", "application/octet-stream", b"\x00two"),
|
||||
], locked_form[0].body
|
||||
|
||||
|
||||
def test_renaming_a_subpath_endpoint_moves_both_routes_keeps_stored_fields_and_never_forwards_caller_headers(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
path: Final = f"/integration-rename-{marker}"
|
||||
renamed_path: Final = f"/integration-renamed-{marker}"
|
||||
body: Final[dict[str, JsonValue]] = {"input": marker}
|
||||
with (
|
||||
wire_server(_received_reply) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}, workers=1) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
created: Final = _create_endpoint(
|
||||
candidate,
|
||||
scenario,
|
||||
{
|
||||
"path": path,
|
||||
"target": f"{wire.url}/base",
|
||||
"include_subpath": True,
|
||||
"auth": True,
|
||||
"headers": {"x-static": "v1"},
|
||||
"methods": ["POST"],
|
||||
},
|
||||
)
|
||||
endpoint_id: Final = str(created["id"])
|
||||
expected: Final = {
|
||||
"id": endpoint_id,
|
||||
"path": path,
|
||||
"target": f"{wire.url}/base",
|
||||
"headers": {"x-static": "v1"},
|
||||
"default_query_params": {},
|
||||
"include_subpath": True,
|
||||
"cost_per_request": 0.0,
|
||||
"timeout": None,
|
||||
"auth": True,
|
||||
"guardrails": None,
|
||||
"is_from_config": False,
|
||||
"methods": ["POST"],
|
||||
}
|
||||
assert created == expected, created
|
||||
assert _readback(candidate, endpoint_id) == [expected]
|
||||
key: Final = scenario.key(allowed_passthrough_routes=[path, renamed_path])
|
||||
caller_headers: Final = {"x-caller": f"caller-{marker}"}
|
||||
first: Final = candidate.request("POST", path, body, key=key, headers=caller_headers)
|
||||
assert first.status_code == 200, first.text
|
||||
assert first.json() == {"received": "/base"}, first.text
|
||||
second: Final = candidate.request("POST", f"{path}/sub", body, key=key, headers=caller_headers)
|
||||
assert second.status_code == 200, second.text
|
||||
assert second.json() == {"received": "/base/sub"}, second.text
|
||||
initial_requests: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in initial_requests] == [
|
||||
("POST", "/base"),
|
||||
("POST", "/base/sub"),
|
||||
], initial_requests
|
||||
|
||||
updated: Final = candidate.request(
|
||||
"POST",
|
||||
f"/config/pass_through_endpoint/{endpoint_id}",
|
||||
{"path": renamed_path, "target": f"{wire.url}/base"},
|
||||
)
|
||||
assert updated.status_code == 200, updated.text
|
||||
updated_expected: Final = {**expected, "path": renamed_path}
|
||||
assert _ENDPOINTS.validate_python(updated.json()["endpoints"]) == [updated_expected], updated.text
|
||||
updated_readback: Final = _readback(candidate, endpoint_id)
|
||||
assert updated_readback == [updated_expected]
|
||||
assert _stored_endpoint(endpoint_id) == {
|
||||
key: value for key, value in updated_readback[0].items() if key != "is_from_config"
|
||||
}
|
||||
|
||||
old_path: Final = candidate.request("POST", path, body, key=key, headers=caller_headers)
|
||||
old_subpath: Final = candidate.request("POST", f"{path}/sub", body, key=key, headers=caller_headers)
|
||||
assert old_path.status_code == 404, old_path.text
|
||||
assert old_subpath.status_code == 404, old_subpath.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
renamed: Final = candidate.request("POST", renamed_path, body, key=key, headers=caller_headers)
|
||||
assert renamed.status_code == 200, renamed.text
|
||||
assert renamed.json() == {"received": "/base"}, renamed.text
|
||||
renamed_subpath: Final = candidate.request("POST", f"{renamed_path}/sub", body, key=key, headers=caller_headers)
|
||||
assert renamed_subpath.status_code == 200, renamed_subpath.text
|
||||
assert renamed_subpath.json() == {"received": "/base/sub"}, renamed_subpath.text
|
||||
refused_verbs: Final = tuple(
|
||||
candidate.request("GET", renamed_path, key=key, headers=caller_headers).status_code for _ in range(4)
|
||||
)
|
||||
assert refused_verbs == (405, 405, 405, 405), refused_verbs
|
||||
anonymous: Final = candidate.client.post(f"{renamed_path}/sub", json=body, headers=caller_headers)
|
||||
assert anonymous.status_code == 401, anonymous.text
|
||||
renamed_requests: Final = wire.drain()
|
||||
assert [(request.method, request.target) for request in (*initial_requests, *renamed_requests)] == [
|
||||
("POST", "/base"),
|
||||
("POST", "/base/sub"),
|
||||
("POST", "/base"),
|
||||
("POST", "/base/sub"),
|
||||
], renamed_requests
|
||||
for upstream in (*initial_requests, *renamed_requests):
|
||||
assert upstream.headers["x-static"] == "v1", upstream.headers
|
||||
assert "x-caller" not in upstream.headers, upstream.headers
|
||||
assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers
|
||||
assert key not in upstream.target, upstream.target
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_configured_endpoint_without_auth_serves_its_subpaths_without_a_key(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: an auth=false include_subpath pass-through endpoint answers 401 on every subpath request")
|
||||
marker: Final = uuid.uuid4().hex
|
||||
path: Final = f"/integration-open-subpath-{marker}"
|
||||
with wire_server(_received_reply) as wire, gateway.scenario() as scenario:
|
||||
_create_endpoint(
|
||||
gateway, scenario, {"path": path, "target": f"{wire.url}/open", "include_subpath": True, "auth": False}
|
||||
)
|
||||
response: Final = gateway.client.post(f"{path}/sub", json={"input": marker})
|
||||
assert response.status_code == 200, response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/open/sub")]
|
||||
|
||||
|
||||
def test_deleted_configured_endpoint_leaves_readback_storage_and_runtime_and_team_get_filters(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with (
|
||||
wire_server(_received_reply) as wire,
|
||||
owned_proxy(gateway, tmp_path, {}) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
allowed: Final = _create_endpoint(
|
||||
candidate, scenario, {"path": f"/integration-allowed-{marker}", "target": f"{wire.url}/allowed"}
|
||||
)
|
||||
hidden: Final = _create_endpoint(
|
||||
candidate, scenario, {"path": f"/integration-hidden-{marker}", "target": f"{wire.url}/hidden"}
|
||||
)
|
||||
team: Final = scenario.team(metadata={"allowed_passthrough_routes": [allowed["path"]]})
|
||||
team_view: Final = candidate.request("GET", f"/config/pass_through_endpoint/team/{team}")
|
||||
assert team_view.status_code == 200, team_view.text
|
||||
assert [endpoint["path"] for endpoint in _ENDPOINTS.validate_python(team_view.json()["endpoints"])] == [
|
||||
allowed["path"]
|
||||
], team_view.text
|
||||
|
||||
served: Final = candidate.request("POST", str(hidden["path"]), {"input": marker})
|
||||
assert served.status_code == 200, served.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/hidden")]
|
||||
|
||||
deleted: Final = candidate.request(
|
||||
"DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(hidden["id"])}
|
||||
)
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert _ENDPOINTS.validate_python(deleted.json()["endpoints"]) == [hidden], deleted.text
|
||||
assert _readback(candidate, str(hidden["id"])) == []
|
||||
listed: Final = candidate.request("GET", "/config/pass_through_endpoint")
|
||||
assert listed.status_code == 200, listed.text
|
||||
listed_ids: Final = [endpoint["id"] for endpoint in _ENDPOINTS.validate_python(listed.json()["endpoints"])]
|
||||
assert hidden["id"] not in listed_ids and allowed["id"] in listed_ids, listed.text
|
||||
stored_ids: Final = _stored_endpoint_ids()
|
||||
assert hidden["id"] not in stored_ids and allowed["id"] in stored_ids, stored_ids
|
||||
|
||||
refused: Final = candidate.request("POST", str(hidden["path"]), {"input": marker})
|
||||
assert refused.status_code == 404, refused.text
|
||||
assert wire.drain() == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue