diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 201ecee5b44..1c03942bce4 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2508,8 +2508,8 @@ _TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( _STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" _STORED_PROMPT_REFUSAL: Final = ( - "This key is limited to specific models, and the gateway cannot check which model a stored prompt runs, " - "so frames that reference a stored prompt are refused for this key." + "This key is limited to specific models, so it cannot run a stored prompt. " + "The gateway cannot check which model a stored prompt runs." ) _FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e6a0b1be28c..8d74b29fe45 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7774,10 +7774,13 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t relayed = await _relay_client_frames([json.dumps(frame)], gate) + error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] + close_reason: Final = relayed.client.close.await_args.kwargs["reason"] assert gate.checked == [STORED_PROMPT_PROBE_MODEL] assert relayed.upstream.send.await_args_list == [] - assert "stored prompt" in json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] assert relayed.client.close.await_args.kwargs["code"] == 1008 + assert "stored prompt" in close_reason + assert close_reason.endswith(".") and error_message.startswith(close_reason) @pytest.mark.asyncio