scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

test(integration): regression tests for August provider translation and streaming bugs (#42621)

Browse source 
* test(integration): Bedrock batch files upload completions and responses records as user messages (Pylon #6882)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): client Anthropic OAuth token never replaces Bedrock SigV4 authorization (Pylon #6888)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bridge /v1/messages and /v1/responses streams through empty-choices chunks (Pylon #6992)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): prepend azure content-filter metadata chunk to the messages stream (Pylon #6992)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): fireworks routers/ slug reaches the provider as accounts/fireworks/routers/<id> (Pylon #7030)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock hidden thinking tokens are not reported as text tokens (Pylon #7067)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): azure_ai FLUX.2-flex image generation targets the flex provider path with the BFL body (Pylon #7092)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): keep Databricks streaming usage and cache reads in the client stream and spend log (Pylon #7094)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): openai-compatible image edits forward provider-specific form fields to the backend (Pylon #7122)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): replayed intercepted web search turn reaches Bedrock as text through /v1/messages (Pylon #7181)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): streamed web search turn capped by max_agentic_loops ends the turn with snippets and ordered blocks (Pylon #7230)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock rerank keeps forwarded client headers out of the SigV4 signature (Pylon #7284)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): azure_ai rerank authenticates with an Entra token when no api key is set (Pylon #7303)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): perplexity stream with cost breakdown object completes and bills total_cost (Pylon #7331)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): azure_ai strips Anthropic message fields before the Foundry request (Pylon #7336)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): capped intercepted web search ends the turn without an internal tool_use block (Pylon #7378)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock passthrough converse-stream keeps event-stream content-type (Pylon #7482)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): v1/messages success exposes v3 priority rate limit headers (Pylon #7532)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): config deployment dropped by a stale boot cost map is restored after reload (Pylon #7564)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock Mantle project id reaches the provider as anthropic-workspace-id (Pylon #7583)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): prefixed Opus 4.8 reasoning_effort reaches Bedrock as adaptive thinking (Pylon #7586)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): dashscope chat forwards reasoning_effort to the provider (Pylon #7606)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): failing stream logging callback still releases the max_parallel_requests slot (Pylon #7608)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): gen 5 Claude Bedrock Invoke tool search sends the Bedrock beta field (Pylon #7642)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): databricks ai gateway api_base requests OAuth token from workspace origin (Pylon #7724)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): deepseek vision image content list reaches the provider unchanged (Pylon #7729)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock Mantle context overflow surfaces as 400 prompt is too long (Pylon #7732)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Codex history items reach Bedrock Mantle as supported Responses input types (Pylon #7783)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): chat over responses deployment returns finish_reason length when output tokens run out (Pylon #7784)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock_mantle rewrites Codex history items before the Responses wire (Pylon #7812)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): advisor sub-call on /v1/messages uses the configured advisor deployment (Pylon #7828)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): tencent thinking reaches the provider body instead of failing the request (Pylon #7834)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): xAI chat web search reaches /v1/responses with instructions and nested filters (Pylon #7835)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): send Bedrock Converse config blocks once at top level (Pylon #7839)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock converse sends gpt-5 reasoning_effort as reasoning.effort (Pylon #7850)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock cohere.embed-english-v3 embeddings accept encoding_format and dimensions (Pylon #7963)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): streamed chat completions emit SSE keepalive pings while the upstream is silent before its first token (Pylon #7987)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(integration): format the TTFT keepalive regression test (Pylon #7987)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): openai chat drops tool_choice when the request has no tools (Pylon #8022)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): stream whose first chunk has no choices falls back and bills the fallback (Pylon #8006)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* Revert "test(integration): Bedrock Mantle project id reaches the provider as anthropic-workspace-id (Pylon #7583)"

This reverts commit 864b65811f.

* test(integration): reconcile contracts manifest for bundled regression tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): isolate proxy config writes in bundled regression tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): address review feedback on keepalive, cost map reload and websearch order

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: kerry <kerry@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] 2026-09-22 20:52:00 -07:00 • committed by GitHub
parent 3eb7e45615
commit 5c0b374f0a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 3370 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

120
tests/integration/contracts.json
View file

@ -1789,6 +1789,126 @@
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
],
"tests/integration/providers/test_anthropic_advisor_wire.py::test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_anthropic_unauthenticated": [
"providers.anthropic_messages_advisor.sub_call_uses_the_configured_advisor_deployment"
],
"tests/integration/providers/test_azure_ai_chat_wire.py::test_azure_ai_strips_thinking_blocks_and_cache_control_from_forwarded_messages": [
"providers.azure_ai.anthropic_message_fields_are_stripped_before_foundry"
],
"tests/integration/providers/test_azure_ai_flux2_image_wire.py::test_azure_flux2_flex_generation_hits_flex_provider_path_not_pro": [
"other.provider_wire.azure_ai.flux2_flex_generation_targets_flex_path_with_bfl_body"
],
"tests/integration/providers/test_azure_ai_rerank_auth_wire.py::test_azure_ai_rerank_with_entra_token_and_no_api_key_sends_bearer_to_provider": [
"other.provider_wire.azure_ai.rerank_entra_token_without_api_key_reaches_provider"
],
"tests/integration/providers/test_bedrock_auth_wire.py::test_client_anthropic_oauth_authorization_header_does_not_replace_bedrock_sigv4_signature": [
"providers.bedrock_auth.client_anthropic_oauth_token_never_replaces_sigv4_authorization"
],
"tests/integration/providers/test_bedrock_batch_files_wire.py::test_completions_and_responses_batch_records_upload_as_anthropic_user_messages": [
"other.provider_wire.bedrock.batch_file_completions_and_responses_records_reach_s3_as_user_messages"
],
"tests/integration/providers/test_bedrock_claude_thinking_wire.py::test_prefixed_opus_4_8_reasoning_effort_reaches_bedrock_as_adaptive_thinking_not_budget_tokens": [
"other.provider_wire.bedrock.prefixed_opus_4_8_reasoning_effort_sends_adaptive_thinking"
],
"tests/integration/providers/test_bedrock_converse_config_blocks_wire.py::test_guardrail_and_performance_config_are_not_duplicated_inside_inference_config": [
"other.provider_wire.bedrock.converse_config_blocks_sent_once_at_top_level"
],
"tests/integration/providers/test_bedrock_embedding_wire.py::test_cohere_embed_english_v3_accepts_encoding_format_and_dimensions": [
"other.provider_wire.bedrock.cohere_embed_english_v3_accepts_encoding_format"
],
"tests/integration/providers/test_bedrock_gpt5_reasoning_wire.py::test_gpt5_reasoning_effort_is_accepted_and_sent_as_converse_reasoning_effort": [
"providers.bedrock_converse.gpt5_reasoning_effort_reaches_provider_as_reasoning_effort"
],
"tests/integration/providers/test_bedrock_invoke_tool_search_wire.py::test_gen5_claude_bedrock_invoke_messages_tool_search_sends_bedrock_beta_field": [
"providers.bedrock_invoke.tool_search_gen5_claude_sends_bedrock_beta_and_reports_support"
],
"tests/integration/providers/test_bedrock_mantle_codex_input_wire.py::test_codex_agent_message_context_compaction_and_local_shell_call_reach_mantle_as_supported_items": [
"other.provider_wire.bedrock_mantle.codex_history_items_reach_mantle_as_supported_types"
],
"tests/integration/providers/test_bedrock_mantle_responses_wire.py::test_codex_agent_message_compaction_and_local_shell_items_are_rewritten_for_mantle": [
"providers.bedrock_mantle.codex_history_items_reach_the_wire_as_supported_input_items"
],
"tests/integration/providers/test_bedrock_mantle_wire.py::test_bedrock_mantle_context_overflow_returns_400_saying_prompt_is_too_long": [
"other.provider_wire.bedrock_mantle.context_overflow_is_reported_as_prompt_too_long"
],
"tests/integration/providers/test_bedrock_messages_web_search_replay_wire.py::test_replayed_intercepted_web_search_turn_reaches_bedrock_as_text_and_answers": [
"providers.bedrock_messages.replayed_intercepted_web_search_turn_is_flattened_to_text"
],
"tests/integration/providers/test_bedrock_passthrough_stream_wire.py::test_bedrock_passthrough_converse_stream_response_carries_event_stream_content_type": [
"other.provider_wire.bedrock.passthrough_stream_keeps_event_stream_content_type"
],
"tests/integration/providers/test_bedrock_rerank_wire.py::test_forwarded_client_header_on_rerank_is_excluded_from_the_sigv4_signature": [
"providers.bedrock_rerank.forwarded_client_headers_are_sent_unsigned"
],
"tests/integration/providers/test_bedrock_thinking_tokens_wire.py::test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens": [
"other.provider_wire.bedrock.hidden_thinking_tokens_are_not_reported_as_text"
],
"tests/integration/providers/test_dashscope_chat_wire.py::test_dashscope_chat_forwards_reasoning_effort_none_to_the_provider": [
"other.provider_wire.dashscope.reasoning_effort_reaches_provider"
],
"tests/integration/providers/test_databricks_chat_wire.py::test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log": [
"other.provider_wire.databricks.stream_usage_and_cache_reads_reach_client_and_spend_log"
],
"tests/integration/providers/test_databricks_oauth_wire.py::test_databricks_ai_gateway_api_base_requests_oauth_token_from_workspace_origin": [
"other.provider_wire.databricks.oauth_token_url_uses_workspace_origin_for_ai_gateway_api_base"
],
"tests/integration/providers/test_deepseek_vision_wire.py::test_deepseek_vision_forwards_image_url_content_list_instead_of_collapsing_to_text": [
"other.provider_wire.deepseek.vision_image_content_list_reaches_provider"
],
"tests/integration/providers/test_fireworks_ai_router_slug_wire.py::test_fireworks_router_slug_chat_sends_router_resource_not_models_path": [
"other.provider_wire.fireworks_ai.router_slug_chat_sends_router_resource_name"
],
"tests/integration/providers/test_fireworks_ai_router_slug_wire.py::test_fireworks_router_slug_text_completion_sends_router_resource_not_models_path": [
"other.provider_wire.fireworks_ai.router_slug_text_completion_sends_router_resource_name"
],
"tests/integration/providers/test_openai_chat_wire.py::test_openai_chat_tool_choice_without_tools_is_not_forwarded": [
"providers.openai_chat_wire.tool_choice_without_tools_is_dropped_before_the_wire"
],
"tests/integration/providers/test_openai_image_edit_wire.py::test_openai_compatible_image_edit_forwards_seed_form_field_to_backend": [
"other.provider_wire.openai.image_edit_forwards_provider_specific_form_fields"
],
"tests/integration/providers/test_responses_bridge_incomplete.py::test_chat_over_responses_deployment_returns_length_when_output_tokens_run_out": [
"other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length"
],
"tests/integration/providers/test_tencent_chat_wire.py::test_tencent_thinking_is_sent_in_provider_body_instead_of_failing_the_request[reasoning_effort_none]": [
"other.provider_wire.tencent.thinking_reaches_provider_in_request_body"
],
"tests/integration/providers/test_tencent_chat_wire.py::test_tencent_thinking_is_sent_in_provider_body_instead_of_failing_the_request[thinking_enabled]": [
"other.provider_wire.tencent.thinking_reaches_provider_in_request_body"
],
"tests/integration/providers/test_websearch_interception_wire.py::test_capped_websearch_interception_loop_ends_turn_instead_of_exposing_internal_tool_use": [
"other.provider_wire.anthropic.websearch_interception_capped_loop_ends_turn_without_internal_tool_use"
],
"tests/integration/providers/test_websearch_interception_wire.py::test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks": [
"other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results"
],
"tests/integration/providers/test_xai_web_search_wire.py::test_xai_chat_web_search_is_sent_to_responses_with_instructions_and_nested_filters": [
"other.provider_wire.xai.chat_web_search_reaches_responses_with_instructions_and_filters"
],
"tests/integration/routing/test_priority_rate_limit_headers.py::test_non_streaming_v1_messages_success_carries_v3_priority_rate_limit_headers": [
"other.routing.priority_rate_limits.v1_messages_success_exposes_v3_priority_headers"
],
"tests/integration/routing/test_stale_cost_map_boot.py::test_config_deployment_dropped_by_stale_boot_cost_map_is_restored_after_reload": [
"other.routing.cost_map.config_deployment_dropped_by_stale_boot_map_is_restored_after_reload"
],
"tests/integration/streaming/test_stream_contracts.py::test_messages_stream_completes_through_trailing_empty_choices_usage_chunk": [
"other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream"
],
"tests/integration/streaming/test_stream_contracts.py::test_perplexity_stream_with_cost_breakdown_object_completes_and_bills_total_cost": [
"other.streaming.usage.provider_cost_object_completes_stream_and_bills_total_cost"
],
"tests/integration/streaming/test_stream_contracts.py::test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bills_the_fallback": [
"other.streaming.fallback.empty_leading_chunk_then_disconnect_streams_fallback_with_usage_and_spend"
],
"tests/integration/streaming/test_stream_contracts.py::test_responses_stream_completes_through_empty_choices_metadata_and_usage_chunks": [
"other.streaming.responses_bridge.empty_choices_chunks_complete_stream"
],
"tests/integration/streaming/test_stream_parallel_slot_release.py::test_failing_stream_logging_callback_does_not_leak_max_parallel_requests_slot": [
"streaming.max_parallel_requests.slot_released_when_stream_logging_callback_fails"
],
"tests/integration/streaming/test_ttft_keepalive.py::test_stream_emits_sse_ping_comments_before_the_first_data_frame_while_upstream_is_silent": [
"streaming.keepalive.sse_pings_fill_silent_time_to_first_token"
],
"tests/integration/observability/test_callback_delivery.py::test_response_survives_raising_success_deployment_hook[chat]": [
"other.observability.callbacks.raising_success_deployment_hook_keeps_response"
],

113
tests/integration/providers/test_anthropic_advisor_wire.py Normal file
View file

@ -0,0 +1,113 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
_ADVISOR_KEY: Final = "synthetic-advisor-key"
_QUESTION: Final = "which index should this query use"
_ADVICE: Final = "use the composite index on (tenant_id, created_at)"
_FINAL_ANSWER: Final = "done, the composite index is the right one"
_ADVISOR_CALL_MESSAGE: Final = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "advisor-call",
"type": "function",
"function": {"name": "advisor", "arguments": json.dumps({"question": _QUESTION})},
}
],
}
_FINAL_MESSAGE: Final = {"role": "assistant", "content": _FINAL_ANSWER}
def _chat_completion(identity: str, message: dict[str, object], finish_reason: str) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{identity}",
"object": "chat.completion",
"created": 1,
"model": "llama-3.3-70b-versatile",
"choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
}
).encode()
)
def _executor_reply(body: dict[str, object], identity: str) -> Reply:
messages: Final = body["messages"]
assert isinstance(messages, list)
if any(message.get("role") == "tool" for message in messages):
assert messages[-1]["content"] == _ADVICE
return _chat_completion(identity, _FINAL_MESSAGE, "stop")
tools: Final = body["tools"]
assert isinstance(tools, list)
assert tools[0]["function"]["name"] == "advisor"
return _chat_completion(identity, _ADVISOR_CALL_MESSAGE, "tool_calls")
@pytest.mark.covers("providers.anthropic_messages_advisor.sub_call_uses_the_configured_advisor_deployment")
def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_anthropic_unauthenticated(
gateway: Gateway,
) -> None:
identity: Final = "advisor-wire-" + uuid.uuid4().hex
def respond(request: Request) -> Reply:
body: Final = json.loads(request.body)
if request.target == "/v1/chat/completions":
assert request.headers["authorization"] == "Bearer integration-provider-key"
return _executor_reply(body, identity)
assert request.target == "/v1/messages"
assert request.headers["x-api-key"] == _ADVISOR_KEY
assert body["model"] == "claude-opus-4-1-20250805"
assert body["messages"] == [
{"role": "user", "content": "please plan the migration"},
{"role": "user", "content": _QUESTION},
]
assert "tools" not in body
return Reply(
body=json.dumps(
{
"id": f"msg-{identity}",
"type": "message",
"role": "assistant",
"model": "claude-opus-4-1-20250805",
"content": [{"type": "text", "text": _ADVICE}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
executor: Final = scenario.model(model="hosted_vllm/llama-3.3-70b", api_base=wire.url + "/v1")
advisor: Final = scenario.model(
model="anthropic/claude-opus-4-1-20250805", api_base=wire.url, api_key=_ADVISOR_KEY
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": executor,
"max_tokens": 64,
"messages": [{"role": "user", "content": "please plan the migration"}],
"tools": [{"type": "advisor_20260301", "name": "advisor", "model": advisor}],
},
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["content"] == [{"type": "text", "text": _FINAL_ANSWER}], response.text
assert body["stop_reason"] == "end_turn", response.text
assert [request.target for request in wire.drain()] == [
"/v1/chat/completions",
"/v1/messages",
"/v1/chat/completions",
]

85
tests/integration/providers/test_azure_ai_chat_wire.py Normal file
View file

@ -0,0 +1,85 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "kimi-k2-thinking"
_API_KEY: Final = "synthetic-azure-ai-key"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_THINKING_BLOCK: Final[JsonValue] = {
"type": "thinking",
"thinking": "The user wants the sum of 17 and 26.",
"signature": "synthetic-signature",
}
_HISTORY_WITH_ANTHROPIC_FIELDS: Final[JsonValue] = [
{
"role": "system",
"content": "You are a calculator.",
"cache_control": {"type": "ephemeral"},
},
{"role": "user", "content": "What is 17 + 26?"},
{
"role": "assistant",
"content": "43",
"thinking_blocks": [_THINKING_BLOCK],
"provider_specific_fields": {"citations": None},
},
{"role": "user", "content": "And doubled?"},
]
_HISTORY_AS_OPENAI_SPEC: Final[JsonValue] = [
{"role": "system", "content": "You are a calculator."},
{"role": "user", "content": "What is 17 + 26?"},
{"role": "assistant", "content": "43"},
{"role": "user", "content": "And doubled?"},
]
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "86"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 31, "completion_tokens": 2, "total_tokens": 33},
}
).encode()
@pytest.mark.covers("providers.azure_ai.anthropic_message_fields_are_stripped_before_foundry")
def test_azure_ai_strips_thinking_blocks_and_cache_control_from_forwarded_messages(gateway: Gateway) -> None:
identity: Final = f"azure-ai-strip-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == _HISTORY_AS_OPENAI_SPEC
return Reply(body=_completion(identity))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"azure_ai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": _HISTORY_WITH_ANTHROPIC_FIELDS},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"role": "assistant", "content": "86"},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

48
tests/integration/providers/test_azure_ai_flux2_image_wire.py Normal file
View file

@ -0,0 +1,48 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_FLEX_MODEL: Final = "azure_ai/FLUX.2-flex"
_PROMPT: Final = "a red fox in the snow"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@pytest.mark.covers("other.provider_wire.azure_ai.flux2_flex_generation_targets_flex_path_with_bfl_body")
def test_azure_flux2_flex_generation_hits_flex_provider_path_not_pro(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview"
assert request.headers["api-key"] == "synthetic-azure-key"
assert _JSON_OBJECT.validate_json(request.body) == {
"model": "FLUX.2-flex",
"prompt": _PROMPT,
"num_images": 2,
"width": 1536,
"height": 1024,
"guidance": 4.5,
"steps": 32,
}
return Reply(body=json.dumps({"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=_FLEX_MODEL, api_base=wire.url, api_key="synthetic-azure-key", api_version="preview"
)
response: Final = gateway.request(
"POST",
"/v1/images/generations",
{"model": model, "prompt": _PROMPT, "n": 2, "size": "1536x1024", "guidance": 4.5, "steps": 32},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["data"] == [
{"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
{"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
]
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview")
]

46
tests/integration/providers/test_azure_ai_rerank_auth_wire.py Normal file
View file

@ -0,0 +1,46 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "azure_ai/Cohere-rerank-v4.0-fast"
ENTRA_TOKEN: Final = "synthetic-entra-access-token"
QUERY: Final = "which document mentions the gateway"
DOCUMENTS: Final = ("the gateway proxies rerank calls", "unrelated synthetic text")
RESPONSE: Final = json.dumps(
{
"id": "synthetic-rerank-id",
"results": [{"index": 0, "relevance_score": 0.91}, {"index": 1, "relevance_score": 0.03}],
"meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}},
}
).encode()
def entra_rerank_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/providers/cohere/v2/rerank"
assert request.headers["authorization"] == f"Bearer {ENTRA_TOKEN}"
assert "api-key" not in request.headers
body: Final = json.loads(request.body)
assert body == {"model": "Cohere-rerank-v4.0-fast", "query": QUERY, "documents": list(DOCUMENTS), "top_n": 2}
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.azure_ai.rerank_entra_token_without_api_key_reaches_provider")
def test_azure_ai_rerank_with_entra_token_and_no_api_key_sends_bearer_to_provider(gateway: Gateway) -> None:
with wire_server(entra_rerank_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=f"{wire.url}/providers/cohere/v2",
azure_ad_token=ENTRA_TOKEN,
model_info={"mode": "rerank"},
)
response: Final = gateway.request(
"POST", "/v1/rerank", {"model": model, "query": QUERY, "documents": list(DOCUMENTS), "top_n": 2}
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert [(result["index"], result["relevance_score"]) for result in body["results"]] == [(0, 0.91), (1, 0.03)]
assert len(wire.drain()) == 1, "Expected exactly one provider rerank call"

164
tests/integration/providers/test_bedrock_auth_wire.py
View file

@ -7,18 +7,22 @@ from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps({
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}).encode()
ACCESS_KEY: Final = "AKIAINTEGRATION000002"
CLIENT_OAUTH_TOKEN: Final = "Bearer sk-ant-oat01-synthetic-client-subscription-token"
RESPONSE: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}
).encode()
def bearer_peer(request: Request) -> Reply:
@ -34,31 +38,58 @@ def bearer_peer(request: Request) -> Reply:
@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain")
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
import litellm
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
for name in tuple(name for name in os.environ if name.startswith("AWS_")):
monkeypatch.delenv(name, raising=False)
for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items():
for name, value in {
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}.items():
monkeypatch.setenv(name, value)
with wire_server(bearer_peer) as wire:
with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"):
await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0)
await asyncio.to_thread(
litellm.completion,
model=MODEL,
aws_profile_name="integration-profile-must-not-be-read",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
messages=[{"role": "user", "content": "synthetic credential control"}],
timeout=5,
num_retries=0,
)
assert wire.drain() == ()
for source in ("argument", "environment"):
if source == "environment":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN)
parameters: Final = {
"model": MODEL, "api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0,
"messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"model": MODEL,
"api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url,
"timeout": 5,
"num_retries": 0,
"messages": [
{"role": "system", "content": "synthetic system"},
{"role": "user", "content": "synthetic bearer request"},
],
"max_tokens": 16,
}
for asynchronous in (False, True):
result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters)
result: Final = (
await litellm.acompletion(**parameters)
if asynchronous
else await asyncio.to_thread(litellm.completion, **parameters)
)
assert result.choices[0].message.content == "bedrock wire control"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
@ -66,28 +97,57 @@ async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credential
@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload")
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None:
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(
gateway: Gateway, tmp_path: Path
) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
with wire_server(bearer_peer) as wire:
parameters: Final = {
"model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url,
"model": MODEL,
"api_key": "os.environ/INTEGRATION_BEARER_TOKEN",
"aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url,
}
alias: Final = f"integration-yaml-{uuid.uuid4().hex}"
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "bedrock.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
overrides: Final = {
"INTEGRATION_BEARER_TOKEN": TOKEN,
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}
with (
owned_proxy(
gateway,
tmp_path,
overrides,
config=path,
remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")),
) as candidate,
candidate.scenario() as scenario,
):
database_model: Final = scenario.model(**parameters)
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {
"model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16, "cache": {"no-cache": True},
})
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [
{"role": "system", "content": "synthetic system"},
{"role": "user", "content": "synthetic bearer request"},
],
"max_tokens": 16,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
@ -95,5 +155,61 @@ def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload
if generation == 0:
entries: Final = candidate.get("/model/info")["data"]
target: Final = next(entry for entry in entries if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}})
response: Final = candidate.request(
"PATCH",
f"/model/{target['model_info']['id']}/update",
{"model_info": {"description": "bearer reload"}},
)
assert response.status_code == 200, response.text
INVOKE_MODEL: Final = "bedrock/invoke/anthropic.claude-3-haiku-20240307-v1:0"
INVOKE_RESPONSE: Final = json.dumps(
{
"id": "msg_synthetic",
"type": "message",
"role": "assistant",
"model": "anthropic.claude-3-haiku-20240307-v1:0",
"content": [{"type": "text", "text": "bedrock invoke wire control"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 4},
}
).encode()
def sigv4_invoke_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/"), dict(
request.headers
)
assert CLIENT_OAUTH_TOKEN not in request.headers.values(), dict(request.headers)
assert json.loads(request.body)["messages"] == [{"role": "user", "content": "synthetic oauth isolation request"}]
return Reply(body=INVOKE_RESPONSE)
@pytest.mark.covers("providers.bedrock_auth.client_anthropic_oauth_token_never_replaces_sigv4_authorization")
def test_client_anthropic_oauth_authorization_header_does_not_replace_bedrock_sigv4_signature(gateway: Gateway) -> None:
with wire_server(sigv4_invoke_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=INVOKE_MODEL,
api_key=None,
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key="synthetic-secret-key-for-testing",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
api_base=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic oauth isolation request"}],
"max_tokens": 16,
},
headers={"Authorization": CLIENT_OAUTH_TOKEN, "x-litellm-api-key": f"Bearer {gateway.key}"},
)
assert response.status_code == 200, response.text
assert response.json()["content"] == [{"type": "text", "text": "bedrock invoke wire control"}], response.text
assert len(wire.drain()) == 1, response.text

78
tests/integration/providers/test_bedrock_batch_files_wire.py Normal file
View file

@ -0,0 +1,78 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0"
BUCKET: Final = "integration-batch-bucket"
PROMPT: Final = "synthetic completions prompt"
RESPONSES_INPUT: Final = "synthetic responses input"
INPUT_LINES: Final = (
{
"custom_id": "completions-record",
"method": "POST",
"url": "/v1/completions",
"body": {"model": MODEL, "prompt": PROMPT, "max_tokens": 64},
},
{
"custom_id": "responses-record",
"method": "POST",
"url": "/v1/responses",
"body": {"model": MODEL, "input": RESPONSES_INPUT, "max_output_tokens": 16},
},
)
EXPECTED_S3_OBJECT: Final = (
{
"recordId": "completions-record",
"modelInput": {
"messages": [{"role": "user", "content": [{"type": "text", "text": PROMPT}]}],
"max_tokens": 64,
"anthropic_version": "bedrock-2023-05-31",
},
},
{
"recordId": "responses-record",
"modelInput": {
"messages": [{"role": "user", "content": [{"type": "text", "text": RESPONSES_INPUT}]}],
"max_tokens": 16,
"anthropic_version": "bedrock-2023-05-31",
},
},
)
def s3_peer(request: Request) -> Reply:
assert request.method == "PUT" and request.target.startswith(f"/{BUCKET}/"), request.target
assert request.headers["authorization"].startswith("AWS4-HMAC-SHA256 ")
return Reply(body=b"")
@pytest.mark.covers(
"other.provider_wire.bedrock.batch_file_completions_and_responses_records_reach_s3_as_user_messages"
)
def test_completions_and_responses_batch_records_upload_as_anthropic_user_messages(gateway: Gateway) -> None:
with wire_server(s3_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=None,
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
aws_region_name="us-east-1",
s3_bucket_name=BUCKET,
s3_endpoint_url=wire.url,
)
jsonl: Final = "\n".join(json.dumps(line, separators=(",", ":")) for line in INPUT_LINES) + "\n"
response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model},
{"file": ("in.jsonl", jsonl.encode(), "application/jsonl")},
)
assert response.status_code == 200, response.text
assert response.json()["object"] == "file" and response.json()["purpose"] == "batch", response.text
uploads: Final = wire.drain()
assert len(uploads) == 1, f"Expected exactly one S3 PUT, saw {[upload.target for upload in uploads]}"
stored: Final = tuple(json.loads(line) for line in uploads[0].body.decode().splitlines() if line.strip())
assert stored == EXPECTED_S3_OBJECT

60
tests/integration/providers/test_bedrock_claude_thinking_wire.py Normal file
View file

@ -0,0 +1,60 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/invoke/us.anthropic.claude-opus-4-8"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps(
{
"id": "msg_adaptive_control",
"type": "message",
"role": "assistant",
"model": "us.anthropic.claude-opus-4-8",
"content": [{"type": "text", "text": "adaptive thinking control"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 5},
}
).encode()
def adaptive_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/us.anthropic.claude-opus-4-8/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"type": "text", "text": "synthetic effort request"}]}]
assert body["thinking"]["type"] == "adaptive", body
assert body["output_config"] == {"effort": "high"}, body
assert "budget_tokens" not in json.dumps(body), body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.prefixed_opus_4_8_reasoning_effort_sends_adaptive_thinking")
def test_prefixed_opus_4_8_reasoning_effort_reaches_bedrock_as_adaptive_thinking_not_budget_tokens(
gateway: Gateway,
) -> None:
with wire_server(adaptive_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
aws_region_name="us-east-1",
api_base=wire.url,
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic effort request"}],
"max_tokens": 4096,
"reasoning_effort": "high",
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "adaptive thinking control"
assert response.json()["usage"]["prompt_tokens"] == 12 and response.json()["usage"]["completion_tokens"] == 5
assert len(wire.drain()) == 1

46
tests/integration/providers/test_bedrock_converse_config_blocks_wire.py Normal file
View file

@ -0,0 +1,46 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE, TOKEN
GUARDRAIL: Final = {"guardrailIdentifier": "integration-guardrail", "guardrailVersion": "DRAFT", "trace": "enabled"}
PERFORMANCE: Final = {"latency": "optimized"}
def converse_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
body: Final = json.loads(request.body)
assert body["inferenceConfig"] == {"maxTokens": 16, "temperature": 0.2}, body
assert body["guardrailConfig"] == GUARDRAIL, body
assert body["performanceConfig"] == PERFORMANCE, body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.converse_config_blocks_sent_once_at_top_level")
def test_guardrail_and_performance_config_are_not_duplicated_inside_inference_config(gateway: Gateway) -> None:
with wire_server(converse_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
guardrailConfig=GUARDRAIL,
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic guardrail request"}],
"max_tokens": 16,
"temperature": 0.2,
"performanceConfig": PERFORMANCE,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert len(wire.drain()) == 1, response.text

58
tests/integration/providers/test_bedrock_embedding_wire.py Normal file
View file

@ -0,0 +1,58 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/cohere.embed-english-v3"
TOKEN: Final = "synthetic-bedrock-bearer"
INPUT: Final = "hello world"
VECTOR: Final = [0.1, 0.2, 0.3]
RESPONSE: Final = json.dumps(
{
"embeddings": {"float": [VECTOR]},
"id": "synthetic-cohere-embed",
"response_type": "embeddings_by_type",
"texts": [INPUT],
}
).encode()
def cohere_english_v3_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/cohere.embed-english-v3/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
assert json.loads(request.body) == {
"texts": [INPUT],
"input_type": "search_document",
"embedding_types": ["float"],
"output_dimension": 512,
}
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.cohere_embed_english_v3_accepts_encoding_format")
def test_cohere_embed_english_v3_accepts_encoding_format_and_dimensions(gateway: Gateway) -> None:
with wire_server(cohere_english_v3_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL,
api_key=TOKEN,
api_base=wire.url,
aws_region_name="us-east-1",
)
for encoding_format in ("float", "base64"):
response: Final = gateway.request(
"POST",
"/v1/embeddings",
{
"model": model,
"input": INPUT,
"encoding_format": encoding_format,
"dimensions": 512,
},
)
assert response.status_code == 200, f"encoding_format={encoding_format}: {response.text}"
assert response.json()["data"] == [
{"object": "embedding", "index": 0, "embedding": VECTOR, "type": "float"},
], response.text
assert len(wire.drain()) == 1, f"encoding_format={encoding_format} never reached Bedrock"

50
tests/integration/providers/test_bedrock_gpt5_reasoning_wire.py Normal file
View file

@ -0,0 +1,50 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/us.openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "gpt-5 reasoning wire control"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14},
"metrics": {"latencyMs": 1},
}
).encode()
def gpt5_converse_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/us.openai.gpt-5.6-sol/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic reasoning request"}]}]
assert body["additionalModelRequestFields"] == {"reasoning": {"effort": "high"}}, body
assert body["inferenceConfig"] == {"maxTokens": 16}, body
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_converse.gpt5_reasoning_effort_reaches_provider_as_reasoning_effort")
def test_gpt5_reasoning_effort_is_accepted_and_sent_as_converse_reasoning_effort(gateway: Gateway) -> None:
with wire_server(gpt5_converse_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": "synthetic reasoning request"}],
"reasoning_effort": "high",
"max_tokens": 16,
"cache": {"no-cache": True},
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "gpt-5 reasoning wire control", response.text
assert response.json()["usage"]["total_tokens"] == 14, response.text
assert len(wire.drain()) == 1

88
tests/integration/providers/test_bedrock_invoke_tool_search_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
MODEL_ID: Final = "us.anthropic.claude-sonnet-5"
TOKEN: Final = "synthetic-bedrock-bearer"
TOOL_SEARCH_TOOL: Final = {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}
DEFERRED_TOOL: Final = {
"name": "get_weather",
"description": "Weather lookup",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
"defer_loading": True,
}
RESPONSE: Final = json.dumps(
{
"id": "msg_tool_search_control",
"type": "message",
"role": "assistant",
"model": MODEL_ID,
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_control",
"name": "tool_search_tool_regex",
"input": {"pattern": "weather"},
},
{
"type": "tool_search_tool_result",
"tool_use_id": "srvtoolu_control",
"content": {
"type": "tool_search_tool_search_result",
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}],
},
},
{"type": "text", "text": "tool search wire control"},
],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
}
).encode()
def tool_search_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == f"/model/{MODEL_ID}/invoke", request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["anthropic_beta"] == ["tool-search-tool-2025-10-19"], body
assert body["messages"] == [{"role": "user", "content": "find the weather tool"}]
assert body["tools"] == [TOOL_SEARCH_TOOL, DEFERRED_TOOL], body["tools"]
assert body["max_tokens"] == 64
assert "model" not in body
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_invoke.tool_search_gen5_claude_sends_bedrock_beta_and_reports_support")
def test_gen5_claude_bedrock_invoke_messages_tool_search_sends_bedrock_beta_field(gateway: Gateway) -> None:
with wire_server(tool_search_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/invoke/{MODEL_ID}",
api_key=TOKEN,
aws_region_name="us-east-1",
api_base=wire.url,
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"messages": [{"role": "user", "content": "find the weather tool"}],
"tools": [TOOL_SEARCH_TOOL, DEFERRED_TOOL],
},
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["content"][2] == {"type": "text", "text": "tool search wire control"}, response.text
assert body["stop_reason"] == "end_turn"
assert body["usage"]["input_tokens"] == 12 and body["usage"]["output_tokens"] == 6
assert len(wire.drain()) == 1
entries: Final = gateway.get("/v1/model/info")["data"]
assert isinstance(entries, list)
info: Final = next(entry for entry in entries if isinstance(entry, dict) and entry["model_name"] == model)
assert isinstance(info["model_info"], dict)
assert info["model_info"]["supports_tool_search"] is True, info["model_info"]

88
tests/integration/providers/test_bedrock_mantle_codex_input_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MODEL: Final = "bedrock_mantle/openai.gpt-5.6-sol"
TOKEN: Final = "synthetic-mantle-bearer"
CIPHERTEXT: Final = "synthetic-compaction-ciphertext"
CALL_ID: Final = "call_synthetic_shell"
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
ACTION: Final[dict[str, JsonValue]] = {"type": "exec", "command": ["ls", "-la"], "timeout_ms": 1000}
RESPONSE: Final = json.dumps(
{
"id": "resp_synthetic_mantle",
"object": "response",
"created_at": 1789788253,
"status": "completed",
"model": "openai.gpt-5.6-sol",
"output": [
{
"type": "message",
"id": "msg_synthetic_mantle",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "mantle wire control", "annotations": []}],
}
],
"usage": {"input_tokens": 21, "output_tokens": 4, "total_tokens": 25},
}
).encode()
def user_turn(text: str) -> JsonValue:
return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}
def codex_history(marker: str) -> tuple[JsonValue, ...]:
return (
user_turn(f"first turn {marker}"),
{"type": "agent_message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent reply"}]},
{"type": "context_compaction", "encrypted_content": CIPHERTEXT},
{"type": "local_shell_call", "call_id": CALL_ID, "status": "completed", "action": ACTION},
{"type": "function_call_output", "call_id": CALL_ID, "output": "synthetic shell output"},
user_turn(f"next turn {marker}"),
)
def mantle_history(marker: str) -> tuple[JsonValue, ...]:
return (
user_turn(f"first turn {marker}"),
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent reply"}]},
{"type": "compaction", "encrypted_content": CIPHERTEXT},
{"type": "function_call", "call_id": CALL_ID, "name": "local_shell", "arguments": json.dumps(ACTION)},
{"type": "function_call_output", "call_id": CALL_ID, "output": "synthetic shell output"},
user_turn(f"next turn {marker}"),
)
@pytest.mark.covers("other.provider_wire.bedrock_mantle.codex_history_items_reach_mantle_as_supported_types")
def test_codex_agent_message_context_compaction_and_local_shell_call_reach_mantle_as_supported_items(
gateway: Gateway,
) -> None:
marker: Final = uuid.uuid4().hex
expected_input: Final = list(mantle_history(marker))
def mantle_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = JSON_OBJECT.validate_json(request.body)
assert body["model"] == "openai.gpt-5.6-sol", body
assert body["input"] == expected_input, body["input"]
return Reply(body=RESPONSE)
with wire_server(mantle_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=MODEL, api_key=TOKEN, api_base=wire.url, aws_region_name="us-east-2")
response: Final = gateway.request(
"POST", "/v1/responses", {"model": model, "input": list(codex_history(marker)), "store": False}
)
assert response.status_code == 200, response.text
assert response.json()["output"][0]["content"][0]["text"] == "mantle wire control", response.text
assert response.json()["usage"]["total_tokens"] == 25, response.text
forwarded: Final = wire.drain()
assert len(forwarded) == 1, forwarded
assert JSON_OBJECT.validate_json(forwarded[0].body)["input"] == expected_input, forwarded[0].body

106
tests/integration/providers/test_bedrock_mantle_responses_wire.py Normal file
View file

@ -0,0 +1,106 @@
import json
from collections.abc import Callable
from typing import Final
from uuid import uuid4
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_MODEL: Final = "bedrock_mantle/openai.gpt-5.6-sol"
_TOKEN: Final = "synthetic-mantle-bearer"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_SHELL_ACTION: Final[dict[str, JsonValue]] = {"type": "exec", "command": ["ls", "-la"], "timeout_ms": 1000}
_OUTPUT_MESSAGE: Final[dict[str, JsonValue]] = {
"type": "message",
"id": "msg_mantle",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "mantle wire control", "annotations": []}],
}
_RESPONSE: Final = json.dumps(
{
"id": "resp_mantle",
"object": "response",
"status": "completed",
"created_at": 1700000000,
"model": "gpt-5.6-sol",
"output": [_OUTPUT_MESSAGE],
"usage": {
"input_tokens": 11,
"output_tokens": 4,
"total_tokens": 15,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
},
}
).encode()
def _codex_history(marker: str) -> list[JsonValue]:
return [
{"type": "message", "role": "user", "content": f"delegate to a subagent {marker}"},
{
"type": "agent_message",
"id": "msg_agent",
"content": [{"type": "text", "text": "sub-agent said "}, {"type": "text", "encrypted_content": "hello"}],
},
{"type": "context_compaction", "id": "cmp_1", "encrypted_content": "compacted-history"},
{
"type": "local_shell_call",
"id": "lsc_1",
"call_id": "call_shell",
"status": "completed",
"action": _SHELL_ACTION,
},
{"type": "function_call_output", "call_id": "call_shell", "output": "total 0"},
]
def _mantle_history(marker: str) -> list[JsonValue]:
return [
{"type": "message", "role": "user", "content": f"delegate to a subagent {marker}"},
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "sub-agent said hello"}]},
{"type": "compaction", "encrypted_content": "compacted-history"},
{
"type": "function_call",
"call_id": "call_shell",
"name": "local_shell",
"arguments": json.dumps(_SHELL_ACTION),
},
{"type": "function_call_output", "call_id": "call_shell", "output": "total 0"},
]
def _mantle_peer(marker: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
assert request.headers["authorization"] == f"Bearer {_TOKEN}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["input"] == _mantle_history(marker), json.dumps(body["input"])
return Reply(body=_RESPONSE)
return respond
@pytest.mark.covers("providers.bedrock_mantle.codex_history_items_reach_the_wire_as_supported_input_items")
def test_codex_agent_message_compaction_and_local_shell_items_are_rewritten_for_mantle(gateway: Gateway) -> None:
marker: Final = uuid4().hex
with wire_server(_mantle_peer(marker)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_TOKEN, aws_region_name="us-east-1")
response: Final = gateway.request(
"POST", "/v1/responses", {"model": model, "input": _codex_history(marker), "stream": False}
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["output"] == [
{
**_OUTPUT_MESSAGE,
"phase": None,
"content": [
{"type": "output_text", "text": "mantle wire control", "annotations": [], "logprobs": None}
],
}
], response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/openai/v1/responses")]

53
tests/integration/providers/test_bedrock_mantle_wire.py Normal file
View file

@ -0,0 +1,53 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "openai.gpt-5.6-sol"
_API_KEY: Final = "synthetic-mantle-bearer"
_PROMPT: Final = "synthetic long conversation control"
_PROMPT_TOKENS: Final = 1055489
_MODEL_MAXIMUM: Final = 1050000
_RESPONSES_PATH: Final = "/openai/v1/responses"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_OVERFLOW_BODY: Final = json.dumps(
{
"error": {
"code": "validation_error",
"message": f"prompt tokens ({_PROMPT_TOKENS}) exceed model maximum ({_MODEL_MAXIMUM}) for {_BACKEND}",
"type": "invalid_request_error",
}
}
).encode()
def _overflow_peer(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == _RESPONSES_PATH
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert _PROMPT in json.dumps(body["input"]), body
return Reply(status=400, body=_OVERFLOW_BODY)
@pytest.mark.covers("other.provider_wire.bedrock_mantle.context_overflow_is_reported_as_prompt_too_long")
def test_bedrock_mantle_context_overflow_returns_400_saying_prompt_is_too_long(gateway: Gateway) -> None:
with wire_server(_overflow_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"bedrock_mantle/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 400, response.text
error: Final = _JSON_OBJECT.validate_json(response.content)["error"]
assert isinstance(error, dict), response.text
assert error["code"] == "400", response.text
message: Final = error["message"]
assert isinstance(message, str), response.text
assert f"prompt is too long: {_PROMPT_TOKENS} tokens > {_MODEL_MAXIMUM} maximum" in message, response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", _RESPONSES_PATH)]

88
tests/integration/providers/test_bedrock_messages_web_search_replay_wire.py Normal file
View file

@ -0,0 +1,88 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
BEDROCK_MODEL: Final = "us.anthropic.claude-opus-5-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
SNIPPET: Final = "synthetic snippet about the integration harness"
INTERCEPTED_TURN: Final = (
{"type": "server_tool_use", "id": "srvtoolu_synthetic", "name": "web_search", "input": {"query": "harness docs"}},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_synthetic",
"content": [
{
"type": "web_search_result",
"url": "https://example.test/harness",
"title": "Harness",
"page_age": None,
"encrypted_content": "",
"snippet": SNIPPET,
},
],
},
{"type": "text", "text": "The harness is documented at example.test"},
)
FLATTENED_TURN: Final = (
{
"type": "text",
"text": f"Web search results for 'harness docs':\n\nTitle: Harness\nURL: https://example.test/harness\nSnippet: {SNIPPET}",
},
{"type": "text", "text": "The harness is documented at example.test"},
)
REPLY: Final = json.dumps(
{
"id": "msg_synthetic_replay",
"type": "message",
"role": "assistant",
"model": BEDROCK_MODEL,
"content": [{"type": "text", "text": "replay accepted"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 30, "output_tokens": 3},
}
).encode()
def bedrock_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == f"/model/{BEDROCK_MODEL}/invoke"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] == [
{"role": "user", "content": "where is the harness documented"},
{"role": "assistant", "content": list(FLATTENED_TURN)},
{"role": "user", "content": "and what does it say"},
], request.body.decode()
assert "tools" not in body, request.body.decode()
return Reply(body=REPLY)
@pytest.mark.covers("providers.bedrock_messages.replayed_intercepted_web_search_turn_is_flattened_to_text")
def test_replayed_intercepted_web_search_turn_reaches_bedrock_as_text_and_answers(gateway: Gateway) -> None:
with wire_server(bedrock_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/{BEDROCK_MODEL}",
api_key=TOKEN,
api_base=wire.url,
aws_region_name="us-east-1",
)
response: Final = gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"messages": [
{"role": "user", "content": "where is the harness documented"},
{"role": "assistant", "content": list(INTERCEPTED_TURN)},
{"role": "user", "content": "and what does it say"},
],
},
headers={"x-api-key": gateway.key, "anthropic-version": "2023-06-01"},
)
assert response.status_code == 200, response.text
assert response.json()["content"] == [{"type": "text", "text": "replay accepted"}], response.text
assert len(wire.drain()) == 1

42
tests/integration/providers/test_bedrock_passthrough_stream_wire.py Normal file
View file

@ -0,0 +1,42 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.upstream import _aws_event_frame
from integration._support.wire import Reply, Request, wire_server
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
_REQUEST_BODY: Final = {"messages": [{"role": "user", "content": [{"text": "synthetic passthrough stream"}]}]}
_EVENTS: Final = (
("messageStart", {"role": "assistant"}),
("contentBlockDelta", {"delta": {"text": "bedrock stream control"}, "contentBlockIndex": 0}),
("messageStop", {"stopReason": "end_turn"}),
("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}),
)
_STREAM_BYTES: Final = b"".join(_aws_event_frame(kind, payload, "sc", "u") for kind, payload in _EVENTS)
def event_stream_peer(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == f"/model/{_MODEL_ID}/converse-stream"
assert json.loads(request.body)["messages"] == _REQUEST_BODY["messages"]
return Reply(body=_STREAM_BYTES, content_type=_EVENT_STREAM)
@pytest.mark.covers("other.provider_wire.bedrock.passthrough_stream_keeps_event_stream_content_type")
def test_bedrock_passthrough_converse_stream_response_carries_event_stream_content_type(gateway: Gateway) -> None:
with wire_server(event_stream_peer) as wire, gateway.scenario() as scenario:
deployment: Final = scenario.model(
model=f"bedrock/{_MODEL_ID}",
api_base=wire.url,
aws_access_key_id="AKIASCRIPTEDPROVIDER",
aws_secret_access_key="scripted-secret",
aws_region_name="us-east-1",
)
response: Final = gateway.request("POST", f"/bedrock/model/{deployment}/converse-stream", _REQUEST_BODY)
assert response.status_code == 200, response.text
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

92
tests/integration/providers/test_bedrock_rerank_wire.py Normal file
View file

@ -0,0 +1,92 @@
import json
import os
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"
ACCESS_KEY: Final = "AKIAINTEGRATION000002"
FORWARDED_FOR: Final = "203.0.113.5"
RESPONSE: Final = json.dumps(
{"results": [{"index": 1, "relevanceScore": 0.9}, {"index": 0, "relevanceScore": 0.1}]}
).encode()
def signed_headers(authorization: str) -> tuple[str, ...]:
return tuple(authorization.split("SignedHeaders=")[1].split(",")[0].split(";"))
def rerank_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/rerank"
assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/")
assert signed_headers(request.headers["authorization"]) == ("content-type", "host", "x-amz-date"), request.headers[
"authorization"
]
assert request.headers["x-forwarded-for"] == FORWARDED_FOR
body: Final = json.loads(request.body)
assert body["queries"] == [{"textQuery": {"text": "synthetic rerank query"}, "type": "TEXT"}]
assert body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["modelConfiguration"] == {
"modelArn": "arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"
}
assert body["rerankingConfiguration"]["bedrockRerankingConfiguration"]["numberOfResults"] == 2
return Reply(body=RESPONSE)
@pytest.mark.covers("providers.bedrock_rerank.forwarded_client_headers_are_sent_unsigned")
def test_forwarded_client_header_on_rerank_is_excluded_from_the_sigv4_signature(
gateway: Gateway, tmp_path: Path
) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["general_settings"]["forward_client_headers_to_llm_api"] = True
path: Final = tmp_path / "forwarding.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {
"AWS_CONFIG_FILE": str(empty),
"AWS_SHARED_CREDENTIALS_FILE": str(empty),
"AWS_EC2_METADATA_DISABLED": "true",
"LITELLM_RUST": "false",
}
with wire_server(rerank_peer) as wire:
with (
owned_proxy(
gateway,
tmp_path,
overrides,
config=path,
remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")),
) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(
model=MODEL,
api_key=None,
api_base=None,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
aws_access_key_id=ACCESS_KEY,
aws_secret_access_key="synthetic-rerank-secret-key-for-testing",
)
response: Final = candidate.request(
"POST",
"/v1/rerank",
{
"model": model,
"query": "synthetic rerank query",
"documents": ["first synthetic document", "second synthetic document"],
"top_n": 2,
},
headers={"x-forwarded-for": FORWARDED_FOR},
)
assert response.status_code == 200, response.text
assert response.json()["results"] == [
{"index": 1, "relevance_score": 0.9},
{"index": 0, "relevance_score": 0.1},
], response.text
assert len(wire.drain()) == 1

94
tests/integration/providers/test_bedrock_thinking_tokens_wire.py Normal file
View file

@ -0,0 +1,94 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
MODEL: Final = "bedrock/converse/global.anthropic.claude-opus-4-8"
TOKEN: Final = "synthetic-bedrock-bearer"
PROMPT: Final = "How many prime numbers are less than 30? Think it through, then answer with just the number."
RESPONSES_PROMPT: Final = "How many prime numbers are less than 30? Answer with just the number."
REDACTED_DATA: Final = "RWRhY3RlZC1ieS1CZWRyb2Nr"
INPUT_TOKENS: Final = 31
OUTPUT_TOKENS: Final = 257
RESPONSE: Final = json.dumps(
{
"output": {
"message": {
"role": "assistant",
"content": [{"reasoningContent": {"redactedContent": REDACTED_DATA}}, {"text": "10"}],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": INPUT_TOKENS,
"outputTokens": OUTPUT_TOKENS,
"totalTokens": INPUT_TOKENS + OUTPUT_TOKENS,
},
"metrics": {"latencyMs": 1},
}
).encode()
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]])
def redacted_thinking_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/global.anthropic.claude-opus-4-8/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
body: Final = json.loads(request.body)
assert body["messages"] in (
[{"role": "user", "content": [{"text": PROMPT}]}],
[{"role": "user", "content": [{"text": RESPONSES_PROMPT}]}],
), body
assert body["additionalModelRequestFields"]["thinking"]["type"] == "adaptive", body
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.hidden_thinking_tokens_are_not_reported_as_text")
def test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens(gateway: Gateway) -> None:
with wire_server(redacted_thinking_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url
)
chat: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": PROMPT}],
"max_tokens": 4000,
"reasoning_effort": "max",
},
)
assert chat.status_code == 200, chat.text
chat_body: Final = _JSON_OBJECT.validate_json(chat.content)
message: Final = _JSON_OBJECT.validate_python(_JSON_LIST.validate_python(chat_body["choices"])[0]["message"])
assert message["content"] == "10", chat.text
assert message["thinking_blocks"] == [{"type": "redacted_thinking", "data": REDACTED_DATA}], chat.text
usage: Final = _JSON_OBJECT.validate_python(chat_body["usage"])
assert usage["completion_tokens"] == OUTPUT_TOKENS, chat.text
details: Final = _JSON_OBJECT.validate_python(usage["completion_tokens_details"])
assert details == {}, chat.text
assert len(wire.drain()) == 1
responses: Final = gateway.request(
"POST",
"/v1/responses",
{"model": model, "input": RESPONSES_PROMPT, "max_output_tokens": 4000, "reasoning": {"effort": "max"}},
)
assert responses.status_code == 200, responses.text
responses_body: Final = _JSON_OBJECT.validate_json(responses.content)
output: Final = _JSON_LIST.validate_python(responses_body["output"])
reasoning_items: Final = tuple(item for item in output if item["type"] == "reasoning")
assert len(reasoning_items) == 1, responses.text
assert reasoning_items[0]["encrypted_content"] == json.dumps(
[{"type": "redacted_thinking", "data": REDACTED_DATA}], separators=(",", ":")
), responses.text
responses_usage: Final = _JSON_OBJECT.validate_python(responses_body["usage"])
assert responses_usage["output_tokens"] == OUTPUT_TOKENS, responses.text
assert _JSON_OBJECT.validate_python(responses_usage["output_tokens_details"])["reasoning_tokens"] == 0, (
responses.text
)
assert len(wire.drain()) == 1

62
tests/integration/providers/test_dashscope_chat_wire.py Normal file
View file

@ -0,0 +1,62 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "qwen3.7-plus"
_API_KEY: Final = "synthetic-dashscope-key"
_PROMPT: Final = "What is 3^3?"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "27"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 17, "completion_tokens": 5, "total_tokens": 22},
}
).encode()
@pytest.mark.covers("other.provider_wire.dashscope.reasoning_effort_reaches_provider")
def test_dashscope_chat_forwards_reasoning_effort_none_to_the_provider(gateway: Gateway) -> None:
identity: Final = f"dashscope-reasoning-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
assert _JSON_OBJECT.validate_json(request.body) == {
"model": _BACKEND,
"messages": [{"role": "user", "content": _PROMPT}],
"reasoning_effort": "none",
}
return Reply(body=_completion(identity))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"dashscope/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}], "reasoning_effort": "none"},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"role": "assistant", "content": "27", "provider_specific_fields": {"refusal": None}},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

129
tests/integration/providers/test_databricks_chat_wire.py Normal file
View file

@ -0,0 +1,129 @@
import json
import uuid
from collections.abc import Mapping
from typing import Final
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
_BACKEND: Final = "databricks-glm-5-2"
_API_KEY: Final = "synthetic-databricks-key"
_PROMPT: Final = "Summarise the cached briefing in one sentence."
_PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = {
"prompt_tokens": 12011,
"completion_tokens": 8,
"total_tokens": 12019,
"cache_read_input_tokens": 12002,
"cache_creation_input_tokens": 0,
}
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
class _PromptTokensDetails(BaseModel):
model_config = ConfigDict(extra="ignore")
cached_tokens: int | None = None
class _Usage(BaseModel):
model_config = ConfigDict(extra="ignore")
prompt_tokens: int
completion_tokens: int
total_tokens: int
prompt_tokens_details: _PromptTokensDetails | None = None
class _Delta(BaseModel):
model_config = ConfigDict(extra="ignore")
content: str | None = None
class _Choice(BaseModel):
model_config = ConfigDict(extra="ignore")
delta: _Delta
class _Chunk(BaseModel):
model_config = ConfigDict(extra="ignore")
id: str
choices: tuple[_Choice, ...]
usage: _Usage | None = None
def _frame(identity: str, choices: list[Mapping[str, object]], usage: Mapping[str, JsonValue] | None = None) -> bytes:
value: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _BACKEND,
"choices": choices,
**({} if usage is None else {"usage": usage}),
}
return b"data: " + json.dumps(value).encode() + b"\n\n"
@pytest.mark.covers("other.provider_wire.databricks.stream_usage_and_cache_reads_reach_client_and_spend_log")
def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gateway: Gateway) -> None:
identity: Final = f"databricks-stream-{uuid.uuid4().hex}"
frames: Final = (
_frame(
identity, [{"index": 0, "delta": {"role": "assistant", "content": "The briefing "}, "finish_reason": None}]
),
_frame(identity, [{"index": 0, "delta": {"content": "is short."}, "finish_reason": None}]),
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
_frame(identity, [], usage=_PROVIDER_USAGE),
b"data: [DONE]\n\n",
)
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
assert body["stream"] is True
return Reply(content_type="text/event-stream", chunks=frames)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": _PROMPT}],
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read()
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
assert lines[-1] == "data: [DONE]", lines
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
assert {chunk.id for chunk in chunks} == {identity}
assert (
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
== "The briefing is short."
)
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert len(usages) == 1, lines
assert (
usages[0].prompt_tokens,
usages[0].completion_tokens,
usages[0].total_tokens,
usages[0].prompt_tokens_details.cached_tokens if usages[0].prompt_tokens_details is not None else None,
) == (12011, 8, 12019, 12002), lines
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
rows: Final = eventually(
lambda: read_rows(
'SELECT prompt_tokens, completion_tokens, total_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019)

92
tests/integration/providers/test_databricks_oauth_wire.py Normal file
View file

@ -0,0 +1,92 @@
import base64
import json
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import parse_qs
import pytest
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_MODEL: Final = "databricks/synthetic-vendor.chat-model.v1"
_CLIENT_ID: Final = "synthetic-databricks-client-id"
_CLIENT_SECRET: Final = "synthetic-databricks-client-secret"
_ACCESS_TOKEN: Final = "synthetic-databricks-oauth-token"
_PROMPT: Final = "Which workspace issued this token?"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _basic_credentials(client_id: str, client_secret: str) -> str:
return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _MODEL.removeprefix("databricks/"),
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "the workspace origin"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13},
}
).encode()
@pytest.mark.covers("other.provider_wire.databricks.oauth_token_url_uses_workspace_origin_for_ai_gateway_api_base")
def test_databricks_ai_gateway_api_base_requests_oauth_token_from_workspace_origin(
gateway: Gateway, tmp_path: Path
) -> None:
identity: Final = f"databricks-oauth-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
if request.target == "/oidc/v1/token":
assert request.method == "POST"
assert request.headers["authorization"] == _basic_credentials(_CLIENT_ID, _CLIENT_SECRET)
assert request.headers["content-type"] == "application/x-www-form-urlencoded"
assert parse_qs(request.body.decode()) == {"grant_type": ["client_credentials"], "scope": ["all-apis"]}
return Reply(
body=json.dumps({"access_token": _ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode()
)
if request.target == "/ai-gateway/mlflow/v1/chat/completions":
assert request.method == "POST"
assert request.headers["authorization"] == f"Bearer {_ACCESS_TOKEN}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _MODEL.removeprefix("databricks/")
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
return Reply(body=_completion(identity))
return Reply(status=401, body=json.dumps({"error": f"unauthenticated path {request.target}"}).encode())
overrides: Final = {"DATABRICKS_CLIENT_ID": _CLIENT_ID, "DATABRICKS_CLIENT_SECRET": _CLIENT_SECRET}
with wire_server(respond) as wire, owned_proxy(gateway, tmp_path, overrides) as candidate:
with candidate.scenario() as scenario:
model: Final = scenario.model(model=_MODEL, api_base=f"{wire.url}/ai-gateway/mlflow/v1", api_key=None)
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {"content": "the workspace origin", "role": "assistant"},
}
]
assert payload["usage"] == {"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13}
assert [(request.method, request.target) for request in wire.drain()] == [
("POST", "/oidc/v1/token"),
("POST", "/ai-gateway/mlflow/v1/chat/completions"),
]

40
tests/integration/providers/test_deepseek_vision_wire.py Normal file
View file

@ -0,0 +1,40 @@
from typing import Final
import httpx
import pytest
from pydantic import JsonValue
from tests.integration._support.client import JSON_OBJECT, Gateway, object_value
_VISION_MODEL: Final = "deepseek-v4-flash-vision-exp"
_API_KEY: Final = "synthetic-deepseek-key"
_VISION_CONTENT: Final[JsonValue] = [
{"type": "text", "text": "what is in this image?"},
{"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}},
]
@pytest.mark.covers("other.provider_wire.deepseek.vision_image_content_list_reaches_provider")
def test_deepseek_vision_forwards_image_url_content_list_instead_of_collapsing_to_text(gateway: Gateway) -> None:
with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream:
upstream.get("/__observations").raise_for_status()
model: Final = scenario.model(
model=f"deepseek/{_VISION_MODEL}",
api_key=_API_KEY,
model_info={"mode": "chat", "supports_vision": True},
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _VISION_CONTENT}]},
)
assert response.status_code == 200, response.text
observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
assert isinstance(observations, list)
assert len(observations) == 1, response.text
observed: Final = object_value(observations[0])
assert observed["path"] == "/v1/chat/completions", response.text
assert observed["authorization"] == f"Bearer {_API_KEY}", response.text
body: Final = object_value(observed["body"])
assert body["model"] == _VISION_MODEL, response.text
assert body["messages"] == [{"role": "user", "content": _VISION_CONTENT}], response.text

84
tests/integration/providers/test_fireworks_ai_router_slug_wire.py Normal file
View file

@ -0,0 +1,84 @@
import json
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_ROUTER_SLUG: Final = "routers/glm-latest"
_ROUTER_RESOURCE: Final = "accounts/fireworks/routers/glm-latest"
_API_KEY: Final = "synthetic-fireworks-key"
_PROMPT: Final = "route me through the router"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _provider_body(request: Request, target: str) -> dict[str, JsonValue]:
assert request.method == "POST"
assert request.target == target
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
return _JSON_OBJECT.validate_json(request.body)
@pytest.mark.covers("other.provider_wire.fireworks_ai.router_slug_chat_sends_router_resource_name")
def test_fireworks_router_slug_chat_sends_router_resource_not_models_path(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
body: Final = _provider_body(request, "/chat/completions")
assert body["model"] == _ROUTER_RESOURCE, body
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
return Reply(
body=json.dumps(
{
"id": "fw-router-chat",
"object": "chat.completion",
"created": 1,
"model": _ROUTER_RESOURCE,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"fireworks_ai/{_ROUTER_SLUG}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["choices"] == [
{"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": "routed"}}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
@pytest.mark.covers("other.provider_wire.fireworks_ai.router_slug_text_completion_sends_router_resource_name")
def test_fireworks_router_slug_text_completion_sends_router_resource_not_models_path(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
body: Final = _provider_body(request, "/completions")
assert body["model"] == _ROUTER_RESOURCE, body
assert body["prompt"] == _PROMPT
return Reply(
body=json.dumps(
{
"id": "fw-router-text",
"object": "text_completion",
"created": 1,
"model": _ROUTER_RESOURCE,
"choices": [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"fireworks_ai/{_ROUTER_SLUG}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request("POST", "/v1/completions", {"model": model, "prompt": _PROMPT})
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["choices"] == [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]

66
tests/integration/providers/test_openai_chat_wire.py Normal file
View file

@ -0,0 +1,66 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "gpt-5.4-mini"
_API_KEY: Final = "synthetic-openai-key"
_PROMPT: Final = "Summarize this conversation in one sentence."
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _completion(identity: str, content: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 19, "completion_tokens": 7, "total_tokens": 26},
}
).encode()
@pytest.mark.covers("providers.openai_chat_wire.tool_choice_without_tools_is_dropped_before_the_wire")
def test_openai_chat_tool_choice_without_tools_is_not_forwarded(gateway: Gateway) -> None:
identity: Final = f"openai-toolless-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
assert "tool_choice" not in body, body
assert "tools" not in body, body
return Reply(body=_completion(identity, "One sentence."))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}], "tool_choice": "none"},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {
"role": "assistant",
"content": "One sentence.",
"provider_specific_fields": {"refusal": None},
},
"provider_specific_fields": {},
}
]
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

75
tests/integration/providers/test_openai_image_edit_wire.py Normal file
View file

@ -0,0 +1,75 @@
import json
from email.message import Message
from email.parser import BytesParser
from email.policy import HTTP
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import BaseModel
_PNG_BYTES: Final = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00"
b"\x1f\x15\xc4\x89\x00\x00\x00\rIDAT\x08\xd7c\xf8\xcf\xc0\xf0\x1f\x00\x05\x00\x01\xff"
b"\x89\x99=\x1d\x00\x00\x00\x00IEND\xaeB`\x82"
)
_PROMPT: Final = "turn the red circle green"
_EDITED_IMAGE_B64: Final = "aW50ZWdyYXRpb24tZWRpdGVkLWltYWdl"
class _Image(BaseModel):
b64_json: str
class _ImageResponse(BaseModel):
data: tuple[_Image, ...]
def _multipart_parts(request: Request) -> tuple[Message, ...]:
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
assert parsed.is_multipart(), request.headers["content-type"]
return tuple(parsed.iter_parts())
def _text_fields(parts: tuple[Message, ...]) -> dict[str, str]:
return {
part.get_param("name", header="content-disposition"): part.get_payload(decode=True).decode()
for part in parts
if part.get_filename() is None
}
def _file_fields(parts: tuple[Message, ...]) -> dict[str, bytes]:
return {
part.get_param("name", header="content-disposition"): part.get_payload(decode=True)
for part in parts
if part.get_filename() is not None
}
@pytest.mark.covers("other.provider_wire.openai.image_edit_forwards_provider_specific_form_fields")
def test_openai_compatible_image_edit_forwards_seed_form_field_to_backend(gateway: Gateway) -> None:
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/v1/images/edits"
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
parts: Final = _multipart_parts(request)
assert _text_fields(parts) == {"model": "gpt-image-1", "prompt": _PROMPT, "seed": "42"}
assert _file_fields(parts) == {"image[]": _PNG_BYTES}
return Reply(body=json.dumps({"created": 1700000000, "data": [{"b64_json": _EDITED_IMAGE_B64}]}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-image-1", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
)
response: Final = gateway.request_multipart(
"/v1/images/edits",
{"model": model, "prompt": _PROMPT, "seed": "42"},
{"image": ("red_circle.png", _PNG_BYTES, "image/png")},
)
assert response.status_code == 200, response.text
payload: Final = _ImageResponse.model_validate_json(response.content)
assert [image.b64_json for image in payload.data] == [_EDITED_IMAGE_B64], response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/images/edits")]

64
tests/integration/providers/test_responses_bridge_incomplete.py Normal file
View file

@ -0,0 +1,64 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
@pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length")
def test_chat_over_responses_deployment_returns_length_when_output_tokens_run_out(gateway: Gateway) -> None:
identity: Final = "responses-incomplete-" + uuid.uuid4().hex
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/responses", request.target
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
body: Final = json.loads(request.body)
assert body["model"] == "gpt-5.3-codex"
assert body["max_output_tokens"] == 16
assert body["reasoning"] == {"effort": "high"}
assert body["input"] == [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": f"explain the plan in detail {identity}"}],
}
]
return Reply(
body=json.dumps(
{
"id": f"resp_{identity}",
"object": "response",
"created_at": 1789788253,
"status": "incomplete",
"incomplete_details": {"reason": "max_output_tokens"},
"model": "gpt-5.3-codex",
"output": [{"type": "reasoning", "id": f"rs_{identity}", "summary": []}],
"usage": {"input_tokens": 12, "output_tokens": 16, "total_tokens": 28},
}
).encode()
)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/responses/gpt-5.3-codex", api_base=wire.url, api_key="synthetic-openai-key"
)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"explain the plan in detail {identity}"}],
"reasoning_effort": "high",
"max_completion_tokens": 16,
},
)
assert response.status_code == 200, response.text
body: Final = response.json()
assert len(wire.drain()) == 1
assert [choice["finish_reason"] for choice in body["choices"]] == ["length"], response.text
assert body["choices"][0]["message"]["content"] == "", response.text
assert body["choices"][0]["message"]["role"] == "assistant", response.text
assert body["usage"]["prompt_tokens"] == 12 and body["usage"]["completion_tokens"] == 16, response.text
assert body["usage"]["total_tokens"] == 28, response.text

86
tests/integration/providers/test_tencent_chat_wire.py Normal file
View file

@ -0,0 +1,86 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "deepseek-v4-pro"
_API_KEY: Final = "synthetic-tencent-key"
_PROMPT: Final = "What is 17 + 26? Answer with just the number."
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_REASONING_REQUESTS: Final[tuple[tuple[str, dict[str, JsonValue], dict[str, JsonValue]], ...]] = (
("thinking_enabled", {"thinking": {"type": "enabled"}}, {"type": "enabled"}),
("reasoning_effort_none", {"reasoning_effort": "none"}, {"type": "disabled"}),
)
def _completion(identity: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _BACKEND,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43."},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64},
}
).encode()
@pytest.mark.covers("other.provider_wire.tencent.thinking_reaches_provider_in_request_body")
@pytest.mark.parametrize(
("reasoning_params", "expected_thinking"),
tuple(case[1:] for case in _REASONING_REQUESTS),
ids=tuple(case[0] for case in _REASONING_REQUESTS),
)
def test_tencent_thinking_is_sent_in_provider_body_instead_of_failing_the_request(
gateway: Gateway, reasoning_params: dict[str, JsonValue], expected_thinking: dict[str, JsonValue]
) -> None:
identity: Final = f"tencent-thinking-{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/chat/completions"
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
assert request.headers["content-type"] == "application/json"
assert _JSON_OBJECT.validate_json(request.body) == {
"model": _BACKEND,
"messages": [{"role": "user", "content": _PROMPT}],
"thinking": expected_thinking,
}
return Reply(body=_completion(identity))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"tencent/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": _PROMPT}], **reasoning_params},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
assert payload["id"] == identity
assert payload["choices"] == [
{
"finish_reason": "stop",
"index": 0,
"message": {
"role": "assistant",
"content": "43",
"reasoning_content": "17 plus 26 is 43.",
"provider_specific_fields": {"refusal": None},
},
"provider_specific_fields": {},
}
]
assert payload["usage"] == {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64}
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]

299
tests/integration/providers/test_websearch_interception_wire.py Normal file
View file

@ -0,0 +1,299 @@
import json
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
INVOKE_TARGET: Final = f"/model/{BEDROCK_MODEL}/invoke"
SEARCH_TARGET: Final = "/tavily/search"
SEARCH_RESULT: Final = {
"title": "Synthetic result",
"url": "https://example.test/result",
"content": "the snippet text",
}
def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]:
frames: Final = tuple(frame for frame in text.split("\n\n") if frame.strip())
return tuple(
(
next(line.removeprefix("event: ") for line in frame.splitlines() if line.startswith("event: ")),
json.loads(next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))),
)
for frame in frames
)
@pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results")
def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks(
gateway: Gateway, tmp_path: Path
) -> None:
def respond(request: Request) -> Reply:
assert request.method == "POST", request.target
body: Final = json.loads(request.body)
if request.target == SEARCH_TARGET:
assert request.headers["authorization"] == "Bearer synthetic-tavily-key"
assert body["query"] == "query-0", body
return Reply(body=json.dumps({"query": "query-0", "results": [SEARCH_RESULT]}).encode())
assert request.target == INVOKE_TARGET
assert request.headers["authorization"] == "Bearer synthetic-bedrock-token"
assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"]
assert "stream" not in body, body
depth: Final = sum(
1
for message in body["messages"]
if isinstance(message["content"], list)
for block in message["content"]
if block["type"] == "tool_result"
)
if depth == 1:
assert body["messages"][2]["content"] == [
{
"type": "tool_result",
"tool_use_id": "toolu_0",
"content": "Title: Synthetic result\nURL: https://example.test/result\nSnippet: the snippet text",
}
], body["messages"]
return Reply(
body=json.dumps(
{
"id": f"msg_{depth}",
"type": "message",
"role": "assistant",
"model": BEDROCK_MODEL,
"content": [
{"type": "text", "text": f"turn-{depth}"},
{
"type": "tool_use",
"id": f"toolu_{depth}",
"name": "litellm_web_search",
"input": {"query": f"query-{depth}"},
},
],
"stop_reason": "tool_use",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 4},
}
).encode()
)
with wire_server(respond) as wire:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["search_tools"] = [
{
"search_tool_name": "integration-search",
"litellm_params": {
"search_provider": "tavily",
"api_key": "synthetic-tavily-key",
"api_base": wire.url + "/tavily",
},
}
]
config["litellm_settings"].update(
{
"callbacks": ["websearch_interception"],
"websearch_interception_params": {
"enabled_providers": ["bedrock"],
"search_tool_name": "integration-search",
"max_agentic_loops": 1,
},
}
)
path: Final = tmp_path / "websearch.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(
model=f"bedrock/{BEDROCK_MODEL}",
api_key="synthetic-bedrock-token",
api_base=wire.url,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint=wire.url,
)
response: Final = candidate.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"stream": True,
"messages": [{"role": "user", "content": "search control"}],
"tools": [{"type": "web_search_20250305", "name": "web_search"}],
},
)
assert response.status_code == 200, response.text
events: Final = sse_events(response.text)
assert [name for name, _ in events][:1] == ["message_start"], response.text
assert [name for name, _ in events][-2:] == ["message_delta", "message_stop"], response.text
for position, (name, event) in enumerate(events):
if name == "content_block_stop":
assert event["index"] in {
earlier_event["index"]
for earlier, earlier_event in events[:position]
if earlier == "content_block_start"
}, response.text
started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start")
search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use")
assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text
assert started[-1] == {"type": "text", "text": ""}, response.text
assert started[:-1] == tuple(
block
for search_id in search_ids
for block in (
{"type": "server_tool_use", "id": search_id, "name": "web_search", "input": {"query": "query-0"}},
{
"type": "web_search_tool_result",
"tool_use_id": search_id,
"content": [
{
"type": "web_search_result",
"url": "https://example.test/result",
"title": "Synthetic result",
"page_age": None,
"encrypted_content": "",
"snippet": "the snippet text",
}
],
},
)
), response.text
assert (
"".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1"
), response.text
assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [
"end_turn"
], response.text
assert "litellm_web_search" not in response.text, response.text
assert [request.target for request in wire.drain()] == [INVOKE_TARGET, SEARCH_TARGET, INVOKE_TARGET]
import threading
import uuid
from typing import Final
from urllib.parse import parse_qs, urlsplit
import httpx
import pytest
from integration._support.client import Gateway, eventually
_QUERY: Final = "integration capped search"
_TEXT_BLOCK: Final = {"type": "text", "text": "searching once more"}
_NOT_INTERCEPTED: Final = "native tool reached the provider"
_SEARCH_RESULT_BLOCK: Final = {
"type": "web_search_result",
"url": "https://owned.invalid/a",
"title": "Owned result",
"page_age": None,
"encrypted_content": "",
"snippet": "owned snippet",
}
def _search_tool_use(identity: str) -> dict[str, object]:
return {"type": "tool_use", "id": identity, "name": "litellm_web_search", "input": {"query": _QUERY}}
def _anthropic_reply(identity: str, content: list[dict[str, object]], stop_reason: str) -> Reply:
return Reply(
body=json.dumps(
{
"id": identity,
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20250929",
"content": content,
"stop_reason": stop_reason,
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 4},
}
).encode()
)
@pytest.mark.covers(
"other.provider_wire.anthropic.websearch_interception_capped_loop_ends_turn_without_internal_tool_use"
)
def test_capped_websearch_interception_loop_ends_turn_instead_of_exposing_internal_tool_use(
gateway: Gateway, tmp_path: Path
) -> None:
identity: Final = "websearch-wire-" + uuid.uuid4().hex
searched: Final = threading.Event()
def respond(request: Request) -> Reply:
parts: Final = urlsplit(request.target)
if request.method == "GET" and parts.path == "/search":
assert parse_qs(parts.query)["q"] == [_QUERY], request.target
searched.set()
return Reply(
body=json.dumps(
{
"results": [
{"title": "Owned result", "url": "https://owned.invalid/a", "content": "owned snippet"}
]
}
).encode()
)
assert request.method == "POST" and parts.path == "/v1/messages", request.target
body: Final = json.loads(request.body)
if any(tool.get("type") == "web_search_20250305" for tool in body["tools"]):
return _anthropic_reply(identity, [{"type": "text", "text": _NOT_INTERCEPTED}], "end_turn")
assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"]
return _anthropic_reply(identity, [_TEXT_BLOCK, _search_tool_use(identity)], "tool_use")
def send(candidate: Gateway, model: str) -> httpx.Response:
return candidate.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 64,
"messages": [{"role": "user", "content": identity + " attempt " + uuid.uuid4().hex}],
"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}],
},
)
def searched_through_proxy(response: httpx.Response) -> bool:
return searched.is_set() and _NOT_INTERCEPTED not in response.text
with wire_server(respond) as wire:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["search_tools"] = [
{
"search_tool_name": "integration-searxng",
"litellm_params": {"search_provider": "searxng", "api_base": wire.url},
}
]
config["litellm_settings"].update(
{
"callbacks": ["websearch_interception"],
"websearch_interception_params": {
"enabled": True,
"enabled_providers": ["anthropic"],
"search_tool_name": "integration-searxng",
},
}
)
path: Final = tmp_path / "websearch.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(
model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key"
)
response: Final = eventually(lambda: send(candidate, model), searched_through_proxy, seconds=40)
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["stop_reason"] == "end_turn", response.text
content: Final = body["content"]
assert [block["type"] for block in content] == ["server_tool_use", "web_search_tool_result", "text"], (
response.text
)
assert content[0]["name"] == "web_search" and content[0]["input"] == {"query": _QUERY}, response.text
assert content[1]["tool_use_id"] == content[0]["id"], response.text
assert content[1]["content"] == [_SEARCH_RESULT_BLOCK], response.text
assert content[2] == _TEXT_BLOCK, response.text
targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain())
assert targets[-3:] == (("POST", "/v1/messages"), ("GET", "/search"), ("POST", "/v1/messages")), targets

84
tests/integration/providers/test_xai_web_search_wire.py Normal file
View file

@ -0,0 +1,84 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
_BACKEND: Final = "grok-4.6-web-search-unmapped"
_API_KEY: Final = "synthetic-xai-key"
_SYSTEM_PROMPT: Final = "Answer in one short sentence and cite the source."
_ALLOWED_DOMAINS: Final = ("weather.example.com", "news.example.org")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def _responses_reply(identity: str, text: str) -> bytes:
return json.dumps(
{
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": _BACKEND,
"output": [
{
"type": "message",
"id": f"msg-{identity}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
],
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [{"type": "web_search"}],
"usage": {"input_tokens": 23, "output_tokens": 41, "total_tokens": 64},
}
).encode()
@pytest.mark.covers("other.provider_wire.xai.chat_web_search_reaches_responses_with_instructions_and_filters")
def test_xai_chat_web_search_is_sent_to_responses_with_instructions_and_nested_filters(gateway: Gateway) -> None:
identity: Final = f"xai-web-search-{uuid.uuid4().hex}"
user_prompt: Final = f"What is the weather in Paris today? Request {identity}."
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert request.target == "/v1/responses", request.target
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
body: Final = _JSON_OBJECT.validate_json(request.body)
assert body["model"] == _BACKEND
assert body["instructions"] == _SYSTEM_PROMPT
assert body["input"] == [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": user_prompt}]}
]
assert body["tools"] == [{"type": "web_search", "filters": {"allowed_domains": list(_ALLOWED_DOMAINS)}}]
assert "web_search_options" not in body
return Reply(body=_responses_reply(identity, "Sunny, 21C."))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"xai/{_BACKEND}", api_base=f"{wire.url}/v1", api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [
{"role": "system", "content": _SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
"web_search_options": {"filters": {"allowed_domains": list(_ALLOWED_DOMAINS)}},
},
)
assert response.status_code == 200, response.text
payload: Final = _JSON_OBJECT.validate_json(response.content)
choices: Final = payload["choices"]
assert isinstance(choices, list) and len(choices) == 1, response.text
choice: Final = choices[0]
assert isinstance(choice, dict), response.text
message: Final = choice["message"]
assert isinstance(message, dict), response.text
assert (message["role"], message["content"]) == ("assistant", "Sunny, 21C."), response.text
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/responses")]

88
tests/integration/routing/test_priority_rate_limit_headers.py Normal file
View file

@ -0,0 +1,88 @@
import json
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
MODEL_RPM: Final = 40
MODEL_TPM: Final = 1000
PREMIUM_SHARE: Final = 0.5
UPSTREAM_REPLY: Final = json.dumps(
{
"id": "msg_priority_headers",
"type": "message",
"role": "assistant",
"model": ANTHROPIC_MODEL,
"content": [{"type": "text", "text": "priority header control"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 4},
}
).encode()
@pytest.mark.covers("other.routing.priority_rate_limits.v1_messages_success_exposes_v3_priority_headers")
def test_non_streaming_v1_messages_success_carries_v3_priority_rate_limit_headers(
gateway: Gateway, tmp_path: Path
) -> None:
probe: Final = "priority header probe " + uuid.uuid4().hex
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/messages"
assert request.headers["x-api-key"] == "synthetic-anthropic-key"
assert json.loads(request.body) == {
"model": ANTHROPIC_MODEL,
"messages": [{"role": "user", "content": probe}],
"max_tokens": 16,
"stream": False,
}
return Reply(body=UPSTREAM_REPLY)
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["litellm_settings"] = {
**configuration["litellm_settings"],
"callbacks": ["dynamic_rate_limiter_v3"],
"priority_reservation": {"premium": PREMIUM_SHARE},
}
path: Final = tmp_path / "priority.yaml"
path.write_text(yaml.safe_dump(configuration))
with (
wire_server(respond) as wire,
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(
model=f"anthropic/{ANTHROPIC_MODEL}",
api_base=wire.url,
api_key="synthetic-anthropic-key",
rpm=MODEL_RPM,
tpm=MODEL_TPM,
)
key: Final = scenario.key(metadata={"priority": "premium"})
response: Final = candidate.request(
"POST",
"/v1/messages",
{"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": probe}]},
key=key,
)
assert response.status_code == 200, response.text
assert response.json()["content"] == [{"type": "text", "text": "priority header control"}], response.text
assert len(wire.drain()) == 1
expected: Final = {
"x-litellm-priority": "premium",
"x-litellm-rate-limiter-version": "v3",
"x-ratelimit-model_saturation_check-limit-requests": str(MODEL_RPM),
"x-ratelimit-model_saturation_check-remaining-requests": str(MODEL_RPM - 1),
"x-ratelimit-priority_model-limit-requests": str(int(MODEL_RPM * PREMIUM_SHARE)),
"x-ratelimit-priority_model-remaining-requests": str(int(MODEL_RPM * PREMIUM_SHARE) - 1),
"x-ratelimit-priority_model-limit-tokens": str(int(MODEL_TPM * PREMIUM_SHARE)),
"x-ratelimit-priority_model-remaining-tokens": str(int(MODEL_TPM * PREMIUM_SHARE) - 1),
}
observed: Final = {name: response.headers.get(name) for name in expected}
assert observed == expected, response.headers

86
tests/integration/routing/test_stale_cost_map_boot.py Normal file
View file

@ -0,0 +1,86 @@
import json
import threading
import uuid
from pathlib import Path
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path:
config: Final = directory / "stale_cost_map_config.yaml"
config.write_text(
json.dumps(
{
"model_list": [
{
"model_name": model,
"litellm_params": {"model": model, "api_base": upstream_url + "/v1", "api_key": "sk-upstream"},
}
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
"store_model_in_db": True,
},
"router_settings": {"disable_cooldowns": True},
}
)
)
return config
@pytest.mark.covers("other.routing.cost_map.config_deployment_dropped_by_stale_boot_map_is_restored_after_reload")
def test_config_deployment_dropped_by_stale_boot_cost_map_is_restored_after_reload(
gateway: Gateway, tmp_path: Path
) -> None:
model: Final = "integration-fresh-" + uuid.uuid4().hex
remote_map: Final = json.dumps(
{model: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0, "output_cost_per_token": 0}}
).encode()
fresh_map_published: Final = threading.Event()
def respond(request: Request) -> Reply:
assert request.target == "/model_prices.json", request
return Reply(body=remote_map) if fresh_map_published.is_set() else Reply(status=503, body=b"{}")
overrides: Final = {"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"}
with (
wire_server(respond) as peer,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
config: Final = _proxy_config(tmp_path, model, gateway.upstream_url)
with owned_proxy(
gateway,
tmp_path,
{**overrides, "LITELLM_MODEL_COST_MAP_URL": peer.url + "/model_prices.json"},
config=config,
remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",),
) as candidate:
assert model not in tuple(entry["id"] for entry in candidate.get("/v1/models")["data"])
fresh_map_published.set()
reload: Final = candidate.request("POST", "/reload/model_cost_map")
assert reload.status_code == 200, reload.text
eventually(
lambda: tuple(str(entry["id"]) for entry in candidate.get("/v1/models")["data"]),
lambda served: model in served,
seconds=30,
)
upstream.get("/__observations").raise_for_status()
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "stale cost map control"}]},
)
assert response.status_code == 200, response.text
assert upstream.get("/__observations").json()["requests"] == [
{
"path": "/v1/chat/completions",
"authorization": "Bearer sk-upstream",
"body": {"model": model, "messages": [{"role": "user", "content": "stale cost map control"}]},
}
]

502
tests/integration/streaming/test_stream_contracts.py
View file

@ -2,25 +2,47 @@ import asyncio
import json
import threading
import uuid
from pathlib import Path
from typing import Final
import pytest
from hypothesis import Phase, example, given, settings, strategies as st
from openai import OpenAI
import yaml
from hypothesis import Phase, example, given, settings
from hypothesis import strategies as st
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, wire_server
from openai import OpenAI
def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
value: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
value: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n"
def text_stream(identity: str) -> tuple[bytes, ...]:
usage: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}
return (frame(identity, {"role": "assistant", "content": "Hello "}), frame(identity, {"content": "雪 café"}), frame(identity, {}, finish="stop"), b"data: " + json.dumps(usage).encode() + b"\n\n", b"data: [DONE]\n\n")
usage: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
return (
frame(identity, {"role": "assistant", "content": "Hello "}),
frame(identity, {"content": "雪 café"}),
frame(identity, {}, finish="stop"),
b"data: " + json.dumps(usage).encode() + b"\n\n",
b"data: [DONE]\n\n",
)
@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage")
@ -37,14 +59,27 @@ def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage
boundaries: Final = (0, *sorted(cuts), len(body))
pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:]))
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "partition control"}], stream=True, stream_options={"include_usage": True}, timeout=5, num_retries=0)
stream: Final = litellm.completion(
model="openai/gpt-4o-mini",
api_base=wire.url + "/v1",
api_key="synthetic-stream-key",
messages=[{"role": "user", "content": "partition control"}],
stream=True,
stream_options={"include_usage": True},
timeout=5,
num_retries=0,
)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert (
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
)
assert {chunk.id for chunk in chunks} == {"stream-partition-control"}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["stop"]
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
"stop"
]
usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
assert len(usages) == 1
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
@ -59,25 +94,63 @@ def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None:
identity: Final = "stream-tools-control"
deltas: Final = (
{"role": "assistant", "tool_calls": [{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}}, {"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}}]},
{"tool_calls": [{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}}, {"index": 0, "function": {"arguments": '{"x":1,'}}]},
{"tool_calls": [{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}}, {"index": 1, "function": {"arguments": '"y":4}'}}]},
{
"role": "assistant",
"tool_calls": [
{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}},
{"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}},
],
},
{
"tool_calls": [
{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}},
{"index": 0, "function": {"arguments": '{"x":1,'}},
]
},
{
"tool_calls": [
{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}},
{"index": 1, "function": {"arguments": '"y":4}'}},
]
},
)
frames: Final = (
*tuple(frame(identity, delta) for delta in deltas),
frame(identity, {}, finish="tool_calls"),
b"data: [DONE]\n\n",
)
frames: Final = (*tuple(frame(identity, delta) for delta in deltas), frame(identity, {}, finish="tool_calls"), b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "tool control"}], stream=True, timeout=5, num_retries=0)
stream: Final = litellm.completion(
model="openai/gpt-4o-mini",
api_base=wire.url + "/v1",
api_key="synthetic-stream-key",
messages=[{"role": "user", "content": "tool control"}],
stream=True,
timeout=5,
num_retries=0,
)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
events: Final = tuple((choice.index, tool) for chunk in chunks for choice in chunk.choices for tool in (choice.delta.tool_calls or ()))
for index, name, call_id, arguments in ((0, "add", "call-add", {"x": 1, "y": 2}), (1, "multiply", "call-multiply", {"x": 3, "y": 4})):
events: Final = tuple(
(choice.index, tool)
for chunk in chunks
for choice in chunk.choices
for tool in (choice.delta.tool_calls or ())
)
for index, name, call_id, arguments in (
(0, "add", "call-add", {"x": 1, "y": 2}),
(1, "multiply", "call-multiply", {"x": 3, "y": 4}),
):
selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index))
assert "".join(tool.id or "" for tool in selected) == call_id
assert "".join(tool.function.name or "" for tool in selected) == name
assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments
assert {tool.index for _, tool in events} == {0, 1}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["tool_calls"]
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == [
"tool_calls"
]
assert len(wire.drain()) == 1
@ -86,13 +159,27 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat
with gateway.scenario() as scenario:
for include in (None, False, True):
identity: Final = "stream-usage-" + uuid.uuid4().hex
with wire_server(lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire:
model: Final = scenario.model(api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
with OpenAI(api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0) as client:
stream: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": identity}], stream=True, **({} if include is None else {"stream_options": {"include_usage": include}}))
with wire_server(
lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))
) as wire:
model: Final = scenario.model(
api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
)
with OpenAI(
api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0
) as client:
stream: Final = client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": identity}],
stream=True,
**({} if include is None else {"stream_options": {"include_usage": include}}),
)
with stream:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert (
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
== "Hello 雪 café"
)
assert {chunk.id for chunk in chunks} == {identity}
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert len(usages) == (1 if include else 0)
@ -101,28 +188,369 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat
requests: Final = wire.drain()
assert len(requests) == 1
assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True
rows: Final = eventually(lambda identity=identity: read_rows('SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70)
rows: Final = eventually(
lambda identity=identity: read_rows(
'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4
assert float(rows[0]["spend"]) == pytest.approx(0.019)
@pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream")
def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None:
identity: Final = "messages-empty-choices-" + uuid.uuid4().hex
metadata: Final = (
b"data: "
+ json.dumps(
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [],
"prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
},
ensure_ascii=False,
).encode()
+ b"\n\n"
)
frames: Final = (metadata, *text_stream(identity))
with (
wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(model="azure/gpt-4o-mini", api_base=wire.url + "/v1")
with gateway.client.stream(
"POST",
"/v1/messages",
json={
"model": model,
"max_tokens": 64,
"stream": True,
"messages": [{"role": "user", "content": identity}],
},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read().decode()
events: Final = tuple(
json.loads(line.removeprefix("data: ")) for line in response.iter_lines() if line.startswith("data: ")
)
assert tuple(event["type"] for event in events) == (
"message_start",
"content_block_start",
"content_block_delta",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
), f"observed events: {events!r}"
assert (
"".join(event["delta"]["text"] for event in events if event["type"] == "content_block_delta") == "Hello 雪 café"
)
message_delta: Final = next(event for event in events if event["type"] == "message_delta")
assert message_delta["usage"] == {"input_tokens": 11, "output_tokens": 4}
requests: Final = wire.drain()
assert len(requests) == 1
outbound: Final = json.loads(requests[0].body)
assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
f"observed outbound body: {outbound!r}"
)
@pytest.mark.covers("other.streaming.responses_bridge.empty_choices_chunks_complete_stream")
def test_responses_stream_completes_through_empty_choices_metadata_and_usage_chunks(gateway: Gateway) -> None:
identity: Final = "responses-empty-choices-" + uuid.uuid4().hex
metadata: Final = (
b"data: "
+ json.dumps(
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [],
"prompt_filter_results": [{"prompt_index": 0, "content_filter_results": {}}],
},
ensure_ascii=False,
).encode()
+ b"\n\n"
)
frames: Final = (metadata, *text_stream(identity))
with (
wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1")
with gateway.client.stream(
"POST",
"/v1/responses",
json={"model": model, "input": identity, "stream": True},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read().decode()
events: Final = tuple(
json.loads(line.removeprefix("data: "))
for line in response.iter_lines()
if line.startswith("data: ") and line != "data: [DONE]"
)
assert (
"".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == "Hello 雪 café"
), f"observed events: {events!r}"
assert tuple(event["type"] for event in events if event["type"] != "response.output_text.delta") == (
"response.created",
"response.in_progress",
"response.output_item.added",
"response.content_part.added",
"response.output_text.done",
"response.content_part.done",
"response.output_item.done",
"response.completed",
), f"observed events: {events!r}"
assert events[-1]["type"] == "response.completed"
assert events[-1]["response"]["usage"] == {
"input_tokens": 11,
"output_tokens": 4,
"output_tokens_details": {"reasoning_tokens": 0, "text_tokens": 4},
"total_tokens": 15,
}
requests: Final = wire.drain()
assert len(requests) == 1
outbound: Final = json.loads(requests[0].body)
assert outbound["stream"] is True and outbound["stream_options"] == {"include_usage": True}, (
f"observed outbound body: {outbound!r}"
)
def provider_cost_object_stream(identity: str, total_cost: float) -> tuple[bytes, ...]:
cost: Final = {
"input_tokens_cost": 0.0001,
"output_tokens_cost": 0.0002,
"request_cost": 0.012,
"total_cost": total_cost,
}
usage: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "sonar",
"choices": [],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15, "cost": cost},
}
return (
frame(identity, {"role": "assistant", "content": "Hello "}),
frame(identity, {"content": "from search"}),
frame(identity, {}, finish="stop"),
b"data: " + json.dumps(usage).encode() + b"\n\n",
b"data: [DONE]\n\n",
)
def sse_data_lines(text: str) -> tuple[str, ...]:
return tuple(line.removeprefix("data: ") for line in text.splitlines() if line.startswith("data: "))
@pytest.mark.covers("other.streaming.usage.provider_cost_object_completes_stream_and_bills_total_cost")
def test_perplexity_stream_with_cost_breakdown_object_completes_and_bills_total_cost(gateway: Gateway) -> None:
identity: Final = "stream-cost-object-" + uuid.uuid4().hex
total_cost: Final = 0.0123
with (
gateway.scenario() as scenario,
wire_server(
lambda request: Reply(
content_type="text/event-stream", chunks=provider_cost_object_stream(identity, total_cost)
)
) as wire,
):
model: Final = scenario.model(model="perplexity/sonar", api_base=wire.url + "/v1")
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": identity}],
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
text: Final = response.read().decode()
assert response.status_code == 200, text
lines: Final = sse_data_lines(text)
assert lines[-1] == "[DONE]", text
events: Final = tuple(json.loads(line) for line in lines[:-1])
assert [event for event in events if "error" in event] == [], text
assert (
"".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"])
== "Hello from search"
), text
assert [
choice.get("finish_reason")
for event in events
for choice in event["choices"]
if choice.get("finish_reason")
] == ["stop"], text
usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
assert len(usages) == 1, text
assert (usages[0]["prompt_tokens"], usages[0]["completion_tokens"], usages[0]["total_tokens"]) == (11, 4, 15), (
text
)
requests: Final = wire.drain()
assert len(requests) == 1
outbound: Final = json.loads(requests[0].body)
assert outbound["model"] == "sonar" and outbound["stream"] is True, outbound
assert outbound["messages"] == [{"role": "user", "content": identity}], outbound
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (11, 4)
assert float(rows[0]["spend"]) == pytest.approx(total_cost)
@pytest.mark.covers(
"other.streaming.fallback.empty_leading_chunk_then_disconnect_streams_fallback_with_usage_and_spend"
)
def test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bills_the_fallback(
gateway: Gateway,
tmp_path: Path,
) -> None:
identity: Final = "stream-empty-fallback-" + uuid.uuid4().hex
empty_first: Final = (
b"data: "
+ json.dumps(
{
"id": identity + "-primary",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [],
"usage": {"prompt_tokens": 11, "completion_tokens": 0, "total_tokens": 11},
}
).encode()
+ b"\n\n"
)
with (
wire_server(
lambda request: Reply(
content_type="text/event-stream",
chunks=(empty_first, b":" + b"x" * 4_000_000 + b"\n\n", empty_first),
abort_after=2,
)
) as primary,
wire_server(
lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))
) as fallback,
):
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
{
"model_name": name,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "synthetic-fallback-key",
"api_base": server.url + "/v1",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
},
}
for name, server in (("primary", primary), ("fallback", fallback))
]
config["router_settings"] = {
"num_retries": 0,
"disable_cooldowns": True,
"fallbacks": [{"primary": ["fallback"]}],
}
path: Final = tmp_path / "fallbacks.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate:
body: Final = {
"model": "primary",
"messages": [{"role": "user", "content": identity}],
"stream": True,
"stream_options": {"include_usage": True},
}
with candidate.client.stream(
"POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {candidate.key}"}
) as response:
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data:"))
assert response.status_code == 200, lines
assert lines[-1] == "data: [DONE]", lines
events: Final = tuple(json.loads(line.removeprefix("data:")) for line in lines[:-1])
assert all("error" not in event for event in events), lines
assert (
"".join(choice["delta"].get("content") or "" for event in events for choice in event["choices"])
== "Hello 雪 café"
), lines
usages: Final = tuple(event["usage"] for event in events if event.get("usage") is not None)
assert (usages[-1]["prompt_tokens"], usages[-1]["completion_tokens"]) == (11, 4), lines
assert tuple(
json.loads(request.body)["messages"]
for request in primary.drain()
if request.target.endswith("/chat/completions")
) == (body["messages"],)
assert tuple(
json.loads(request.body)["messages"]
for request in fallback.drain()
if request.target.endswith("/chat/completions")
) == (body["messages"],)
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, prompt_tokens, completion_tokens, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(identity,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["status"]) == (11, 4, "success"), (
rows
)
assert float(rows[0]["spend"]) == pytest.approx(0.019), rows
@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
import litellm
for truncated in (True, False):
with wire_server(lambda request, truncated=truncated: Reply(content_type="text/event-stream", chunks=text_stream("stream-truncated"), abort_after=1 if truncated else None)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "truncation control"}], stream=True, timeout=5, num_retries=0)
with wire_server(
lambda request, truncated=truncated: Reply(
content_type="text/event-stream",
chunks=text_stream("stream-truncated"),
abort_after=1 if truncated else None,
)
) as wire:
stream: Final = litellm.completion(
model="openai/gpt-4o-mini",
api_base=wire.url + "/v1",
api_key="synthetic-stream-key",
messages=[{"role": "user", "content": "truncation control"}],
stream=True,
timeout=5,
num_retries=0,
)
try:
if truncated:
with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read") as failure:
with pytest.raises(
litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read"
) as failure:
tuple(stream)
assert isinstance(failure.value.original_exception, litellm.APIConnectionError)
assert failure.value.generated_content == "Hello "
assert failure.value.is_pre_first_chunk is False
else:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert (
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
== "Hello 雪 café"
)
assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices)
finally:
asyncio.run(stream.aclose())
@ -134,9 +562,23 @@ def test_client_cancellation_releases_the_actual_provider_connection() -> None:
import litellm
gate: Final = threading.Event()
frames: Final = (frame("stream-cancel", {"role": "assistant", "content": "first"}), b":" + b"x" * 4_000_000 + b"\n\n", b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "cancellation control"}], stream=True, timeout=5, num_retries=0)
frames: Final = (
frame("stream-cancel", {"role": "assistant", "content": "first"}),
b":" + b"x" * 4_000_000 + b"\n\n",
b"data: [DONE]\n\n",
)
with wire_server(
lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)
) as wire:
stream: Final = litellm.completion(
model="openai/gpt-4o-mini",
api_base=wire.url + "/v1",
api_key="synthetic-stream-key",
messages=[{"role": "user", "content": "cancellation control"}],
stream=True,
timeout=5,
num_retries=0,
)
try:
first: Final = next(stream)
assert first.choices[0].delta.content == "first"

86
tests/integration/streaming/test_stream_parallel_slot_release.py Normal file
View file

@ -0,0 +1,86 @@
import json
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
def frame(identity: str, delta: dict[str, str], *, finish: str | None = None) -> bytes:
event: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return b"data: " + json.dumps(event).encode() + b"\n\n"
@pytest.mark.covers("streaming.max_parallel_requests.slot_released_when_stream_logging_callback_fails")
def test_failing_stream_logging_callback_does_not_leak_max_parallel_requests_slot(
gateway: Gateway, tmp_path: Path
) -> None:
identity: Final = "stream-slot-" + uuid.uuid4().hex
prompt: Final = "slot release control " + identity
def analyzer(request: Request) -> Reply:
assert request.target == "/analyze"
assert json.loads(request.body)["text"] == prompt
return Reply(status=500, body=json.dumps({"error": "synthetic analyzer outage"}).encode())
def provider(request: Request) -> Reply:
assert request.target == "/v1/chat/completions"
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": prompt}]
assert body["stream"] is True
return Reply(
content_type="text/event-stream",
chunks=(
frame(identity, {"role": "assistant", "content": "Hello"}),
frame(identity, {"content": " slot"}),
frame(identity, {}, finish="stop"),
b"data: [DONE]\n\n",
),
)
with wire_server(analyzer) as policy, wire_server(provider) as upstream:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["guardrails"] = [
{
"guardrail_name": identity,
"litellm_params": {
"guardrail": "presidio",
"mode": "logging_only",
"default_on": True,
"presidio_filter_scope": "input",
"pii_entities_config": {"EMAIL_ADDRESS": "MASK"},
"presidio_analyzer_api_base": policy.url + "/",
"presidio_anonymizer_api_base": policy.url + "/",
},
}
]
path: Final = tmp_path / "failing_logging_guardrail.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(api_base=upstream.url + "/v1")
key: Final = scenario.key(max_parallel_requests=1)
body: Final = {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True}
first: Final = candidate.request("POST", "/v1/chat/completions", body, key=key)
assert first.status_code == 200, first.text
assert first.text.endswith("data: [DONE]\n\n"), first.text
assert len(upstream.drain()) == 1
eventually(lambda: policy.received.qsize(), lambda count: count >= 1)
assert {scan.target for scan in policy.drain()} == {"/analyze"}
second: Final = eventually(
lambda: candidate.request("POST", "/v1/chat/completions", body, key=key),
lambda response: response.status_code == 200,
seconds=20,
return_last_on_timeout=True,
)
assert second.status_code == 200, second.text
assert second.text.endswith("data: [DONE]\n\n"), second.text

62
tests/integration/streaming/test_ttft_keepalive.py Normal file
View file

@ -0,0 +1,62 @@
import json
import threading
import uuid
from collections.abc import Callable, Iterable, Iterator
from typing import Final
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from integration.streaming.test_stream_contracts import text_stream
KEEPALIVE_SECONDS: Final = 1
def _reply_after_first_ping(identity: str, first_ping_seen: threading.Event) -> Callable[[Request], Reply]:
def respond(_request: Request) -> Reply:
first_ping_seen.wait(timeout=10)
return Reply(content_type="text/event-stream", chunks=text_stream(identity))
return respond
def _frames_setting(first_ping_seen: threading.Event, lines: Iterable[str]) -> Iterator[str]:
for line in lines:
if line == ": ping":
first_ping_seen.set()
yield line
@pytest.mark.covers("streaming.keepalive.sse_pings_fill_silent_time_to_first_token")
def test_stream_emits_sse_ping_comments_before_the_first_data_frame_while_upstream_is_silent(
gateway: Gateway,
) -> None:
identity: Final = "stream-ttft-keepalive-" + uuid.uuid4().hex
first_ping_seen: Final = threading.Event()
with gateway.scenario() as scenario:
with wire_server(_reply_after_first_ping(identity, first_ping_seen)) as wire:
model: Final = scenario.model(api_base=wire.url + "/v1", keepalive_seconds=KEEPALIVE_SECONDS)
with gateway.client.stream(
"POST",
"/v1/chat/completions",
json={"model": model, "messages": [{"role": "user", "content": identity}], "stream": True},
headers={"Authorization": f"Bearer {gateway.key}"},
) as response:
assert response.status_code == 200, response.read().decode()
frames: Final = tuple(
_frames_setting(first_ping_seen, (line for line in response.iter_lines() if line))
)
first_data: Final = next(index for index, line in enumerate(frames) if line.startswith("data:"))
assert first_data >= 1, f"No keepalive reached the client before the first data frame: {frames}"
assert frames[:first_data] == (": ping",) * first_data, frames
assert frames[-1] == "data: [DONE]", frames
deltas: Final = tuple(json.loads(line.removeprefix("data: ")) for line in frames[first_data:-1])
assert (
"".join(choice["delta"].get("content", "") for chunk in deltas for choice in chunk["choices"])
== "Hello 雪 café"
), frames
requests: Final = wire.drain()
assert len(requests) == 1
outbound: Final = json.loads(requests[0].body)
assert outbound["model"] == "gpt-4o-mini" and outbound["stream"] is True, outbound
assert "keepalive_seconds" not in outbound, outbound