diff --git a/litellm/proxy/common_utils/banner.py b/litellm/proxy/common_utils/banner.py index 25ad5e295f6..6926f3f013b 100644 --- a/litellm/proxy/common_utils/banner.py +++ b/litellm/proxy/common_utils/banner.py @@ -1,6 +1,6 @@ +import sys from typing import Final -# LiteLLM ASCII banner LITELLM_BANNER: Final = """ ██╗ ██╗████████╗███████╗██╗ ██╗ ███╗ ███╗ ██║ ██║╚══██╔══╝██╔════╝██║ ██║ ████╗ ████║ ██║ ██║ ██║ █████╗ ██║ ██║ ██╔████╔██║ @@ -8,12 +8,38 @@ LITELLM_BANNER: Final = """ ██╗ ██╗████████╗ ███████╗██║ ██║ ███████╗███████╗███████╗██║ ╚═╝ ██║ ╚══════╝╚═╝ ╚═╝ ╚══════╝╚══════╝╚══════╝╚═╝ ╚═╝""" +LITELLM_BANNER_ASCII: Final = r""" _ _____ _______ ______ _ _ __ __ + | | |_ _||__ __|| ____|| | | | | \/ | + | | | | | | | |__ | | | | | \ / | + | | | | | | | __| | | | | | |\/| | + | |____ _| |_ | | | |____ | |____ | |____ | | | | + |______||_____| |_| |______||______||______||_| |_|""" -def show_banner(): - """Display the LiteLLM CLI banner.""" + +def banner_for_encoding(encoding: str | None) -> str: + """Return the banner variant that ``encoding`` can represent.""" + if encoding is None: + return LITELLM_BANNER + try: + LITELLM_BANNER.encode(encoding) + except (UnicodeEncodeError, LookupError): + return LITELLM_BANNER_ASCII + return LITELLM_BANNER + + +def _echo(text: str) -> None: try: import click - click.echo(f"\n{LITELLM_BANNER}\n") + click.echo(text) except ImportError: - print("\n") # noqa: T201 + print(text) # noqa: T201 # the banner is CLI output, and click is what we would print through + + +def show_banner() -> None: + """Display the LiteLLM CLI banner. Never raises: the banner must not block startup.""" + encoding: Final[str | None] = getattr(sys.stdout, "encoding", None) + try: + _echo(f"\n{banner_for_encoding(encoding)}\n") + except UnicodeEncodeError: + pass diff --git a/tests/test_litellm/proxy/common_utils/test_banner.py b/tests/test_litellm/proxy/common_utils/test_banner.py new file mode 100644 index 00000000000..9f2f96db82b --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_banner.py @@ -0,0 +1,61 @@ +from contextlib import redirect_stdout +from pathlib import Path + +import pytest + +from litellm.proxy.common_utils.banner import ( + LITELLM_BANNER, + LITELLM_BANNER_ASCII, + banner_for_encoding, + show_banner, +) + + +def _banner_written_to_stdout(encoding: str, destination: Path) -> str: + with destination.open("w", encoding=encoding, errors="strict") as stream, redirect_stdout(stream): + show_banner() + return destination.read_text(encoding=encoding) + + +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16", "cp1252", "cp437", "cp850", "ascii", "latin-1", "koi8-r"]) +def test_banner_for_encoding_returns_a_banner_the_encoding_can_represent(encoding: str): + """ + The banner we hand to stdout must always survive that stream's codec. + + Regression for the proxy exiting at startup on Windows when stdout is redirected: + the ANSI code page there is cp1252, which has none of the box-drawing characters + the unicode banner is drawn with. + """ + assert banner_for_encoding(encoding).encode(encoding) + + +def test_banner_for_encoding_prefers_the_unicode_banner_when_the_codec_allows_it(): + assert banner_for_encoding("utf-8") == LITELLM_BANNER + + +def test_banner_for_encoding_degrades_when_the_codec_cannot_represent_the_unicode_banner(): + assert banner_for_encoding("cp1252") == LITELLM_BANNER_ASCII + + +def test_banner_for_encoding_degrades_for_a_codec_python_does_not_know(): + assert banner_for_encoding("not-a-real-codec") == LITELLM_BANNER_ASCII + + +def test_banner_for_encoding_keeps_the_unicode_banner_when_the_stream_has_no_encoding(): + assert banner_for_encoding(None) == LITELLM_BANNER + + +def test_show_banner_does_not_raise_when_stdout_cannot_encode_the_unicode_banner(tmp_path: Path): + """ + show_banner() runs as the first statement of run_server, so anything it raises takes + the whole proxy down before it binds a port. + """ + written = _banner_written_to_stdout("cp1252", tmp_path / "stdout.log") + + assert LITELLM_BANNER_ASCII in written + + +def test_show_banner_still_writes_the_unicode_banner_to_a_utf8_stdout(tmp_path: Path): + written = _banner_written_to_stdout("utf-8", tmp_path / "stdout.log") + + assert LITELLM_BANNER in written