diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 7f68b72a522..a9413fe72a6 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -250,11 +250,8 @@ class ProxyInitializationHelpers: click.echo(f"\nLiteLLM: Making a test ChatCompletions request to your proxy. Model={request_model}") import openai - api_base = f"http://{host}:{port}" - if isinstance(test, str): - api_base = test - else: - raise ValueError("Invalid test value") + bare_host: Final = f"[{host}]" if ":" in host and not host.startswith("[") else host + api_base: Final = test if isinstance(test, str) else f"http://{bare_host}:{port}" client: Final = openai.OpenAI(api_key="My API Key", base_url=api_base) response: Final = client.chat.completions.create( diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 89a35923be0..ff670d6a5ed 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -74,7 +74,7 @@ class TestProxyInitializationHelpers: @patch("openai.OpenAI") @patch("click.echo") @patch("builtins.print") - def test_run_test_chat_completion(self, mock_print, mock_echo, mock_openai): + def test_run_test_chat_completion(self, mock_print, mock_echo, mock_openai): # test-quality-ok: pins the string-URL passthrough to the client; the flag path is asserted at the HTTP boundary in the stub-server test below # Setup mock_client = MagicMock() mock_openai.return_value = mock_client @@ -90,12 +90,6 @@ class TestProxyInitializationHelpers: ] # Execute - with pytest.raises(ValueError, match="Invalid test value"): - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", True - ) - - # Test with valid string test value ProxyInitializationHelpers._run_test_chat_completion( "localhost", 8000, "gpt-3.5-turbo", "http://test-url" ) @@ -106,6 +100,101 @@ class TestProxyInitializationHelpers: ) mock_client.chat.completions.create.assert_called() + @pytest.mark.parametrize("host", ["127.0.0.1", "::1"]) + def test_run_test_chat_completion_flag_uses_host_and_port(self, host: str) -> None: + import json + import socket + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + seen: Final[list[tuple[str, dict[str, object]]]] = [] + + class Stub(BaseHTTPRequestHandler): + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + seen.append((self.path, body)) + if self.path.endswith("/chat/completions") and body.get("stream"): + payload = ( + 'data: {"id":"c","object":"chat.completion.chunk","created":1,' + '"model":"m","choices":[{"index":0,"delta":{"content":"hi"},' + '"finish_reason":null}]}\n\ndata: [DONE]\n\n' + ) + ctype = "text/event-stream" + elif self.path.endswith("/chat/completions"): + payload = json.dumps( + { + "id": "c", + "object": "chat.completion", + "created": 1, + "model": "m", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + ) + ctype = "application/json" + else: + payload = json.dumps( + { + "id": "c", + "object": "text_completion", + "created": 1, + "model": "m", + "choices": [{"index": 0, "text": "hi", "finish_reason": "stop"}], + } + ) + ctype = "application/json" + data = payload.encode() + self.send_response(200) + self.send_header("Content-Type", ctype) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, *args): + pass + + class StubServer(HTTPServer): + address_family = socket.AF_INET6 if ":" in host else socket.AF_INET + + try: + server: Final = StubServer((host, 0), Stub) + except OSError: + pytest.skip(f"cannot bind {host} on this machine") + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + ProxyInitializationHelpers._run_test_chat_completion( + host, server.server_port, "my-model", True + ) + finally: + server.shutdown() + server.server_close() + + chat: Final = [body for path, body in seen if path.endswith("/chat/completions")] + assert [b.get("stream", False) for b in chat] == [False, True] + assert all(b["model"] == "my-model" for b in chat) + + @patch("openai.OpenAI") + def test_test_flag_on_cli_sends_request_to_host_and_port(self, mock_openai): + from click.testing import CliRunner + + result = CliRunner().invoke( + run_server, + ["--test", "--host", "127.0.0.1", "--port", "4000", "--model", "my-model"], + ) + + assert result.exit_code == 0, result.output + mock_openai.assert_called_once_with( + api_key="My API Key", base_url="http://127.0.0.1:4000" + ) + first_call = mock_openai.return_value.chat.completions.create.call_args_list[0] + assert first_call.kwargs["model"] == "my-model" + def test_get_default_unvicorn_init_args(self): # Test without log_config args = ProxyInitializationHelpers._get_default_unvicorn_init_args(