diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index 6a8ece73d8..262d0dff12 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -8,7 +8,8 @@ so the full proxy -> router -> OpenAI SDK -> httpx path is exercised. import json import time -from typing import Iterator, List +import uuid +from typing import Tuple import httpx @@ -17,10 +18,14 @@ import httpx # Pre-built response templates # --------------------------------------------------------------------------- +def _mock_id() -> str: + return f"chatcmpl-mock-{uuid.uuid4().hex[:8]}" + + def _chat_completion_json(model: str) -> dict: """Return a minimal valid ChatCompletion object.""" return { - "id": "chatcmpl-mock", + "id": _mock_id(), "object": "chat.completion", "created": int(time.time()), "model": model, @@ -42,74 +47,10 @@ def _chat_completion_json(model: str) -> dict: } -def _streaming_sse_payloads(model: str) -> List[bytes]: - """Pre-build the SSE byte payloads for a streaming response.""" - chunk = { - "id": "chatcmpl-mock", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": model, - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "Mock response"}, - "finish_reason": None, - } - ], - } - done_chunk = { - "id": "chatcmpl-mock", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": model, - "choices": [ - { - "index": 0, - "delta": {}, - "finish_reason": "stop", - } - ], - } - return [ - b"data: " + json.dumps(chunk).encode() + b"\n\n", - b"data: " + json.dumps(done_chunk).encode() + b"\n\n", - b"data: [DONE]\n\n", - ] - - -# --------------------------------------------------------------------------- -# Byte-stream wrappers -# --------------------------------------------------------------------------- - -class MockSSEAsyncStream(httpx.AsyncByteStream): - """Async byte stream that yields pre-built SSE payloads.""" - - def __init__(self, payloads: List[bytes]) -> None: - self._payloads = payloads - - async def __aiter__(self): # type: ignore[override] - for payload in self._payloads: - yield payload - - -class MockSSESyncStream(httpx.SyncByteStream): - """Sync byte stream that yields pre-built SSE payloads.""" - - def __init__(self, payloads: List[bytes]) -> None: - self._payloads = payloads - - def __iter__(self) -> Iterator[bytes]: - return iter(self._payloads) - - # --------------------------------------------------------------------------- # Transport # --------------------------------------------------------------------------- -_STREAM_HEADERS = { - "content-type": "text/event-stream", -} - _JSON_HEADERS = { "content-type": "application/json", } @@ -123,22 +64,17 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): """ @staticmethod - def _parse_request(request: httpx.Request) -> tuple: - """Extract (model, stream) from the request body.""" - body = json.loads(request.content) + def _parse_request(request: httpx.Request) -> Tuple[str, bool]: + """Extract model from the request body.""" + try: + body = json.loads(request.content) + except (json.JSONDecodeError, ValueError): + return ("mock-model", False) model = body.get("model", "mock-model") - stream = body.get("stream", False) - return model, stream + return (model, False) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - model, stream = self._parse_request(request) - if stream: - payloads = _streaming_sse_payloads(model) - return httpx.Response( - status_code=200, - headers=_STREAM_HEADERS, - stream=MockSSEAsyncStream(payloads), - ) + model, _ = self._parse_request(request) body = json.dumps(_chat_completion_json(model)).encode() return httpx.Response( status_code=200, @@ -147,14 +83,7 @@ class MockOpenAITransport(httpx.AsyncBaseTransport, httpx.BaseTransport): ) def handle_request(self, request: httpx.Request) -> httpx.Response: - model, stream = self._parse_request(request) - if stream: - payloads = _streaming_sse_payloads(model) - return httpx.Response( - status_code=200, - headers=_STREAM_HEADERS, - stream=MockSSESyncStream(payloads), - ) + model, _ = self._parse_request(request) body = json.dumps(_chat_completion_json(model)).encode() return httpx.Response( status_code=200, diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 61f150f1c2..b9c474b327 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_ssl_configuration, ) +from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport def _get_client_init_params(cls: type) -> Tuple[str, ...]: @@ -206,8 +207,6 @@ class BaseOpenAILLM: return litellm.aclient_session if getattr(litellm, "network_mock", False): - from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport - return httpx.AsyncClient(transport=MockOpenAITransport()) # Get unified SSL configuration @@ -231,8 +230,6 @@ class BaseOpenAILLM: return litellm.client_session if getattr(litellm, "network_mock", False): - from litellm.llms.custom_httpx.mock_transport import MockOpenAITransport - return httpx.Client(transport=MockOpenAITransport()) # Get unified SSL configuration diff --git a/tests/test_litellm/llms/custom_httpx/test_mock_transport.py b/tests/test_litellm/llms/custom_httpx/test_mock_transport.py index d18f79ac97..94d942b126 100644 --- a/tests/test_litellm/llms/custom_httpx/test_mock_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_mock_transport.py @@ -1,6 +1,6 @@ """ Tests for MockOpenAITransport — verifies that the mock transport produces -responses parseable by the OpenAI SDK for both streaming and non-streaming paths. +responses parseable by the OpenAI SDK. """ import json @@ -60,71 +60,28 @@ class TestNonStreaming: body = json.loads(response.content) assert body["model"] == "my-custom-model" - -# --------------------------------------------------------------------------- -# Streaming -# --------------------------------------------------------------------------- - - -class TestStreaming: - def test_sync_streaming_returns_sse_events(self): + def test_unique_ids_per_response(self): transport = MockOpenAITransport() request = httpx.Request( method="POST", url="https://api.openai.com/v1/chat/completions", - content=json.dumps({"model": "gpt-4o", "stream": True, "messages": []}), + content=json.dumps({"model": "gpt-4o", "messages": []}), + ) + r1 = json.loads(transport.handle_request(request).content) + r2 = json.loads(transport.handle_request(request).content) + assert r1["id"] != r2["id"] + + def test_empty_body_does_not_crash(self): + transport = MockOpenAITransport() + request = httpx.Request( + method="GET", + url="https://api.openai.com/v1/models", + content=b"", ) response = transport.handle_request(request) assert response.status_code == 200 - assert "text/event-stream" in response.headers["content-type"] - - chunks = list(response.stream) - # Should have: content chunk, finish chunk, [DONE] - assert len(chunks) == 3 - assert chunks[-1] == b"data: [DONE]\n\n" - - # Parse the first chunk - first_line = chunks[0].decode() - assert first_line.startswith("data: ") - data = json.loads(first_line[len("data: "):].strip()) - assert data["object"] == "chat.completion.chunk" - assert data["model"] == "gpt-4o" - assert data["choices"][0]["delta"]["content"] == "Mock response" - - @pytest.mark.asyncio - async def test_async_streaming_returns_sse_events(self): - transport = MockOpenAITransport() - request = httpx.Request( - method="POST", - url="https://api.openai.com/v1/chat/completions", - content=json.dumps({"model": "gpt-4o", "stream": True, "messages": []}), - ) - response = await transport.handle_async_request(request) - assert response.status_code == 200 - - chunks = [] - async for chunk in response.stream: - chunks.append(chunk) - - assert len(chunks) == 3 - assert chunks[-1] == b"data: [DONE]\n\n" - - # Parse finish chunk - finish_line = chunks[1].decode() - data = json.loads(finish_line[len("data: "):].strip()) - assert data["choices"][0]["finish_reason"] == "stop" - - def test_streaming_model_echoed(self): - transport = MockOpenAITransport() - request = httpx.Request( - method="POST", - url="https://api.openai.com/v1/chat/completions", - content=json.dumps({"model": "custom-stream", "stream": True, "messages": []}), - ) - response = transport.handle_request(request) - first_chunk = next(iter(response.stream)) - data = json.loads(first_chunk.decode()[len("data: "):].strip()) - assert data["model"] == "custom-stream" + body = json.loads(response.content) + assert body["model"] == "mock-model" # ---------------------------------------------------------------------------