diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0ddfa1210c2..4857973d51d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -329,6 +329,20 @@ class ChunkProcessor: models[0] if models else None, ) + @staticmethod + def _get_upstream_provider(chunks: Sequence[object]) -> str | None: + def _provider_from_chunk(chunk: object) -> str | None: + hidden_params: Final = get_hidden_params(chunk) + if hidden_params is None: + return None + provider: Final[object] = hidden_params.get("upstream_provider") + return provider if isinstance(provider, str) and provider else None + + return next( + (provider for chunk in chunks if (provider := _provider_from_chunk(chunk)) is not None), + None, + ) + @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, @@ -506,10 +520,22 @@ class ChunkProcessor: chunks, first_chunk_model, ) - if provider_response_model is not None: - response._hidden_params = dict( # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter - response._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params getter - provider_response_model=provider_response_model, + upstream_provider: Final = self._get_upstream_provider(chunks) + provider_hidden_params: Final[Mapping[str, str]] = MappingProxyType( + { + key: value + for key, value in ( + ("provider_response_model", provider_response_model), + ("upstream_provider", upstream_provider), + ) + if value is not None + } + ) + if provider_hidden_params: + response_hidden_params: Final[Mapping[str, builtins.object]] = get_hidden_params(response) or {} + set_hidden_params( + response, + dict(response_hidden_params, **provider_hidden_params), ) return response diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9710956c97e..484edd869cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -277,6 +277,7 @@ class _VertexChunkLike(Protocol): class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None + upstream_provider: str | None = None def _provider_response_model(chunk: object) -> str | None: @@ -293,9 +294,17 @@ def _parsed_provider_hidden_params(hidden: object) -> _ParsedChunkHiddenParams | return None +def _provider_upstream_provider(chunk: object) -> str | None: + parsed: Final = _parsed_provider_hidden_params(get_hidden_params(chunk)) + if parsed is None or not parsed.upstream_provider: + return None + return parsed.upstream_provider + + def _provider_hidden_params( chunk: object, provider_response_model: str | None, + upstream_provider: str | None, ) -> Mapping[str, object] | None: hidden: Final = get_hidden_params(chunk) parsed: Final = _parsed_provider_hidden_params(hidden) @@ -308,6 +317,10 @@ def _provider_hidden_params( for key, value in ( ("provider_response_model", provider_response_model), ("provider_specific_fields", provider_specific_fields), + ( + "upstream_provider", + parsed.upstream_provider if parsed is not None and parsed.upstream_provider else upstream_provider, + ), ) if value is not None } @@ -360,6 +373,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None + self._upstream_provider: str | None = None self.received_finish_reason: str | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ @@ -1653,10 +1667,17 @@ class CustomStreamWrapper: provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model + upstream_provider: Final = _provider_upstream_provider(chunk) + if self._upstream_provider is None and upstream_provider is not None: + self._upstream_provider = upstream_provider prebuilt_delta: Final = self._prebuilt_plain_delta(chunk) plain_delta_created: bool = prebuilt_delta is not None # rebind-ok: the first chunk builds its delta below model_response = self.model_response_creator( - hidden_params=_provider_hidden_params(chunk, self._provider_response_model), + hidden_params=_provider_hidden_params( + chunk, + self._provider_response_model, + self._upstream_provider, + ), delta=prebuilt_delta, ) response_obj: dict[str, Any] = {} diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index cca6815399f..69f3d76c034 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -14,7 +14,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.core_helpers import finite_cost -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -212,19 +212,27 @@ class OpenrouterConfig(OpenAIGPTConfig): # OpenRouter returns cost information in the usage object when usage.include=true try: response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + response_provider: Final = response_json.get("provider") + if isinstance(response_provider, str) and response_provider: + set_hidden_params( + model_response, + {**(get_hidden_params(model_response) or {}), "upstream_provider": response_provider}, + ) if "usage" in response_json and response_json["usage"]: response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost") if response_cost is not None: # Store cost in hidden params for the cost calculator to use if not hasattr(model_response, HIDDEN_PARAMS_ATTR): set_hidden_params(model_response, {}) - hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata - dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + cost_hidden_params: Final = ( + cast( # cast-ok: preserve mapping operations on dynamic response metadata + dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + ) ) - if "additional_headers" not in hidden_params: - hidden_params["additional_headers"] = {} + if "additional_headers" not in cost_hidden_params: + cost_hidden_params["additional_headers"] = {} additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata - dict[str, object], hidden_params["additional_headers"] + dict[str, object], cost_hidden_params["additional_headers"] ) reported_cost: Final = ( finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None @@ -282,7 +290,7 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") new_choices.append(choice) - return ModelResponseStream( + model_response: Final = ModelResponseStream( id=chunk["id"], object="chat.completion.chunk", created=chunk["created"], @@ -290,6 +298,12 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): model=chunk["model"], choices=new_choices, ) + typed_chunk: Final[Mapping[str, object]] = cast( # cast-ok: stream chunks are decoded JSON objects + Mapping[str, object], chunk + ) + upstream_provider: Final[object] = typed_chunk.get("provider") + if isinstance(upstream_provider, str) and upstream_provider: + set_hidden_params(model_response, {"upstream_provider": upstream_provider}) except KeyError as e: raise OpenRouterException( message=f"KeyError: {e}, Got unexpected response from OpenRouter: {chunk}", @@ -298,3 +312,5 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): ) except Exception as e: raise e + else: + return model_response diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f165355da5a..73b12399af4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3758,6 +3758,7 @@ class SpendLogsMetadata(TypedDict): kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None] internal_call_origin: InternalCallOrigin | None billed_retry_attempt: ReadOnly[NotRequired[bool | None]] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_roi_estimator: ReadOnly[NotRequired[bool | None]] is_silent_experiment: ReadOnly[NotRequired[bool | None]] guardrail_information: list[StandardLoggingGuardrailInformation] | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f8d5cba4229..306e70f7029 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -194,6 +195,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "autorouter_savings_estimate", "autorouter_baseline_observation", "used_client_oauth_token", + "upstream_provider", KUBERNETES_POD_ROUTING_KEY, "litellm_roi_estimator", "is_silent_experiment", @@ -203,6 +205,10 @@ _STAMPED_METADATA_KEYS: Final = frozenset( ) +class _UpstreamProviderMetadata(TypedDict, total=False): + upstream_provider: ReadOnly[str] + + def _stored_user_email(metadata: Mapping[str, object]) -> str | None: email: Final = metadata.get("user_api_key_user_email") return email if isinstance(email, str) and email.strip() else None @@ -221,6 +227,7 @@ def _get_spend_logs_metadata( model_map_information: StandardLoggingModelInformation | None = None, cold_storage_object_key: str | None = None, litellm_overhead_time_ms: float | None = None, + upstream_provider: str | None = None, cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, @@ -232,6 +239,9 @@ def _get_spend_logs_metadata( kubernetes_pod_routing: StandardLoggingKubernetesPodRouting | None = None, is_silent_experiment: bool = False, ) -> SpendLogsMetadata: + upstream_provider_metadata: Final[_UpstreamProviderMetadata] = ( + {"upstream_provider": upstream_provider} if upstream_provider is not None else {} + ) if metadata is None: return SpendLogsMetadata( user_api_key=None, @@ -265,6 +275,7 @@ def _get_spend_logs_metadata( eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, + **upstream_provider_metadata, attempted_retries=None, max_retries=None, attempted_fallbacks=None, @@ -290,6 +301,7 @@ def _get_spend_logs_metadata( **MappingProxyType( {key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS} ), + **upstream_provider_metadata, autorouter_savings=autorouter_savings, autorouter_savings_estimate=autorouter_savings_estimate, autorouter_baseline_observation=autorouter_baseline_observation, @@ -706,6 +718,14 @@ def get_logging_payload( if standard_logging_payload is not None: hidden_params: Final = standard_logging_payload.get("hidden_params", {}) litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + upstream_provider_value: Final[object] = ( + standard_logging_payload["hidden_params"]["upstream_provider"] + if standard_logging_payload is not None + and "hidden_params" in standard_logging_payload + and "upstream_provider" in standard_logging_payload["hidden_params"] + else None + ) + upstream_provider: Final = upstream_provider_value if isinstance(upstream_provider_value, str) else None logged_provider: Final = ( kwargs.get("custom_llm_provider") @@ -807,6 +827,7 @@ def get_logging_payload( else None ), litellm_overhead_time_ms=litellm_overhead_time_ms, + upstream_provider=upstream_provider, cost_breakdown=( standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None ), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f8115571dfd..db865c7488a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3460,6 +3460,7 @@ class StandardLoggingHiddenParams(TypedDict): batch_models: list[str] | None batch_successful_requests: ReadOnly[int | None] batch_failed_requests: ReadOnly[int | None] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None diff --git a/tests/integration/spend/test_openrouter_upstream_provider.py b/tests/integration/spend/test_openrouter_upstream_provider.py new file mode 100644 index 00000000000..fa17a364aea --- /dev/null +++ b/tests/integration/spend/test_openrouter_upstream_provider.py @@ -0,0 +1,216 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +_BACKEND: Final = "openai/gpt-4o-mini" +_PROMPT: Final = "upstream provider trace" +_PROVIDER: Final = "serving-site-a" +_ANSWER: Final = "hello" +_USAGE: Final[dict[str, JsonValue]] = { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + "cost": 0.00015, +} + + +def _reply( + request_id: str, + *, + response_provider: bool, + stream_provider_mode: Literal["all", "last", "none"], +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target.split("?", 1)[0].endswith("/chat/completions"), request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["messages"] == [{"role": "user", "content": _PROMPT}], body + if body.get("stream") is True: + common: Final = { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + } + all_provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} + first_provider: Final = all_provider if stream_provider_mode == "all" else {} + finish_provider: Final = ( + all_provider if stream_provider_mode in ("all", "last") else {} + ) + usage_provider: Final = all_provider if stream_provider_mode == "all" else {} + chunks: Final[tuple[dict[str, JsonValue], ...]] = ( + { + **common, + **first_provider, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": _ANSWER}, "finish_reason": None}], + }, + { + **common, + **finish_provider, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + {**common, **usage_provider, "choices": [], "usage": _USAGE}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + + (b"data: [DONE]\n\n",), + ) + + provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} if response_provider else {} + return Reply( + body=json.dumps( + { + "id": request_id, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + **provider, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _ANSWER}, + "finish_reason": "stop", + } + ], + "usage": _USAGE, + } + ).encode() + ) + + return respond + + +def _stream_chunks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _send( + gateway: Gateway, + scenario: Scenario, + wire_url: str, + request_id: str, + *, + openrouter: bool, + stream: bool, +) -> httpx.Response: + model: Final = scenario.model( + model=f"{'openrouter' if openrouter else 'openai'}/{_BACKEND}", + api_base=wire_url, + api_key="synthetic-provider-key", + ) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + assert response.status_code == 200, response.text + returned_id: Final = ( + string_value(_stream_chunks(response)[0]["id"]) + if stream + else string_value(JSON_OBJECT.validate_json(response.content)["id"]) + ) + assert returned_id == request_id, response.text + return response + + +def _spend_metadata(request_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return object_value(rows[0]["metadata"]) + + +@pytest.mark.parametrize( + ("stream", "stream_provider_mode"), + [ + pytest.param(False, "none", id="nonstream"), + pytest.param(True, "all", id="stream-every-chunk"), + pytest.param(True, "last", id="stream-later-chunk"), + ], +) +def test_openrouter_provider_is_persisted_in_spend_metadata_regression( + gateway: Gateway, stream: bool, stream_provider_mode: Literal["all", "last", "none"] +) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode=stream_provider_mode, + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=stream) + metadata: Final = _spend_metadata(request_id) + + assert metadata.get("upstream_provider") == _PROVIDER, metadata + + +def test_openrouter_client_response_keeps_provider_in_body_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + response: Final = _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=False) + + body: Final = JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) + message: Final = object_value(object_value(choices[0])["message"]) + assert body["provider"] == _PROVIDER + assert "upstream_provider" not in body + assert message["content"] == _ANSWER + + +def test_non_openrouter_spend_metadata_has_no_provider_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=False, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=False, stream=False) + metadata: Final = _spend_metadata(request_id) + + assert "upstream_provider" not in metadata, metadata diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d0e31cdd57a..f7c2082b683 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12574,9 +12574,10 @@ def test_get_hidden_params(): assert result["api_base"] is None assert result["response_cost"] is None assert result["additional_headers"] is None + assert result.get("upstream_provider") is None # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") # Test with valid params hidden_params = { @@ -12584,6 +12585,7 @@ def test_get_hidden_params(): "cache_key": "test-cache", "api_base": "https://api.test.com", "response_cost": 0.001, + "upstream_provider": "serving-site-a", "additional_headers": { "x-ratelimit-limit-requests": "2000", "x-ratelimit-remaining-requests": "1999", @@ -12596,8 +12598,9 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 + assert result.get("upstream_provider") == "serving-site-a" # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 9df45b92686..bbf1897b3a7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -910,6 +910,71 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" +def test_stream_chunk_builder_finds_upstream_provider_in_later_chunk_regression() -> None: + first: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + later: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" world", role=None), + ) + ], + ) + first_hidden_params: Final = {"origin": "first"} + later_hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunks: Final = [ + {**first.model_dump(), "_hidden_params": first_hidden_params}, + {**later.model_dump(), "_hidden_params": later_hidden_params}, + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + assert first_hidden_params.get("upstream_provider") is None + assert later_hidden_params.get("upstream_provider") == "serving-site-a" + + +def test_stream_chunk_builder_preserves_source_provider_metadata_guard() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunk_dict: Final = {**chunk.model_dump(), "_hidden_params": hidden_params} + + response: Final = stream_chunk_builder(chunks=[chunk_dict]) + + assert response is not None + assert chunk_dict["_hidden_params"]["upstream_provider"] == "serving-site-a" + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): """Vertex AI metadata on streaming chunks must appear on assembled response.""" grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 888c47505a4..3d659641f54 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -872,3 +872,120 @@ def test_openrouter_streaming_reasoning_is_preserved() -> None: assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages) assert assembled is not None assert assembled.choices[0].message.reasoning_content == "Consider the evidence. Apply the rule." + + +def _transform_openrouter_response(body: dict[str, object]) -> litellm.ModelResponse: + return OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-4o-mini", + raw_response=httpx.Response(200, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_openrouter_nonstream_provider_without_usage_is_preserved_regression() -> None: + response: Final = _transform_openrouter_response( + { + "id": "gen-provider-no-usage", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + ) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_invalid_provider_keeps_reported_cost_guard( + provider_fields: dict[str, object], case: str +) -> None: + response: Final = _transform_openrouter_response( + { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + "cost": 0.00015, + }, + } + ) + + assert response._hidden_params.get("upstream_provider") is None + assert response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER] == 0.00015 + + +def test_openrouter_streaming_chunk_preserves_provider_regression() -> None: + chunk: Final = { + "id": "gen-provider-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_streaming_chunk_ignores_invalid_provider_guard(provider_fields: dict[str, object], case: str) -> None: + chunk: Final = { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") is None diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 1ea445fda0d..a91cdd512fe 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2394,6 +2394,65 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): ) +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_carries_upstream_provider_from_hidden_params_regression() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload_with_provider: Final = { + **dict(standard_logging_payload), + "hidden_params": { + **dict(standard_logging_payload["hidden_params"]), + "upstream_provider": "serving-site-a", + }, + } + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": payload_with_provider, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert json.loads(payload["metadata"]).get("upstream_provider") == "serving-site-a" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_omits_upstream_provider_when_hidden_param_absent_guard() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert "upstream_provider" not in json.loads(payload["metadata"]) + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_strips_null_bytes_from_request_tags(): @@ -6911,7 +6970,10 @@ def test_spend_logs_payload(model_id: str | None): ) assert len(payload["request_id"]) > 0 - expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - {"is_silent_experiment"} + expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - { + "is_silent_experiment", + "upstream_provider", + } assert "metadata" in payload assert isinstance(payload["metadata"], str)