mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
fix(proxy): set model on non-streaming RejectedRequestError responses
When a pre-call hook returns a string, the chat and text completion routes answer with it through RejectedRequestError. The non-streaming response was a bare ModelResponse/TextCompletionResponse with model=None, so OpenAI SDK clients failed to parse the chat response. Take the model from the request data, falling back to the exception's, the same way the streaming branch does.
This commit is contained in:
parent
dddbaa803d
commit
7793917aae
2 changed files with 83 additions and 0 deletions
|
|
@ -12824,6 +12824,7 @@ async def chat_completion(
|
|||
request_data=_data,
|
||||
)
|
||||
_chat_response = litellm.ModelResponse()
|
||||
_chat_response.model = _data.get("model") or e.model or ""
|
||||
_chat_response.choices[0].message.content = e.message
|
||||
|
||||
if data.get("stream", None) is not None and data["stream"] is True:
|
||||
|
|
@ -13009,6 +13010,7 @@ async def completion(
|
|||
else:
|
||||
_response = litellm.TextCompletionResponse()
|
||||
_response.choices[0].text = e.message
|
||||
_response.model = _data.get("model") or e.model or ""
|
||||
return _response
|
||||
except Exception as e:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
|
|
|
|||
81
tests/unit/proxy/test_rejected_request_response_model.py
Normal file
81
tests/unit/proxy/test_rejected_request_response_model.py
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
"""Regression test for the non-streaming RejectedRequestError response.
|
||||
|
||||
A pre-call hook that returns a string raises ``RejectedRequestError``, and the
|
||||
route answers with that string instead of calling the LLM. The non-streaming
|
||||
response used to be a bare ``ModelResponse()`` with ``model=None``, which the
|
||||
OpenAI SDK rejects (``model: Input should be a valid string``).
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request, Response
|
||||
from openai.types.chat import ChatCompletion
|
||||
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import chat_completion, completion
|
||||
|
||||
|
||||
async def _run_rejected_request(route, body, exception_model="my-model"):
|
||||
exception = RejectedRequestError(
|
||||
message="cached answer",
|
||||
model=exception_model,
|
||||
llm_provider="",
|
||||
request_data=dict(body),
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.read_request_body",
|
||||
new_callable=AsyncMock,
|
||||
return_value=dict(body),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.base_process_llm_request",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=exception,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
|
||||
):
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
return await route(
|
||||
request=MagicMock(spec=Request),
|
||||
fastapi_response=MagicMock(spec=Response),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_chat_completion_response_has_model():
|
||||
response = await _run_rejected_request(
|
||||
chat_completion,
|
||||
{"model": "my-model", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
|
||||
assert response.model == "my-model"
|
||||
ChatCompletion.model_validate(response.model_dump())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_text_completion_response_has_model():
|
||||
response = await _run_rejected_request(completion, {"model": "my-model", "prompt": "hi"})
|
||||
|
||||
# finish_reason is still None here, so this is not validated against
|
||||
# openai.types.Completion the way the chat response is.
|
||||
assert response.model == "my-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route, body",
|
||||
[
|
||||
(chat_completion, {"model": "my-model", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
(completion, {"model": "my-model", "prompt": "hi"}),
|
||||
],
|
||||
)
|
||||
async def test_rejected_response_falls_back_to_request_model(route, body):
|
||||
# RejectedRequestError is public, so a custom hook can raise it without a model.
|
||||
response = await _run_rejected_request(route, body, exception_model=None)
|
||||
|
||||
assert response.model == "my-model"
|
||||
Loading…
Add table
Reference in a new issue