mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(proxy_server): allow http urls for UI_LOGO_PATH via redirect and improve test hygiene
This commit is contained in:
Adnaan Ali
parent
67f90254ed
commit
13084981f1
2 changed files with 31 additions and 161 deletions
|
|
@ -10992,88 +10992,23 @@ def get_logo_url():
|
|||
@app.get("/get_image", include_in_schema=False)
|
||||
async def get_image():
|
||||
"""Get logo to show on admin UI"""
|
||||
logo_path = os.getenv("UI_LOGO_PATH", "")
|
||||
|
||||
# get current_dir
|
||||
# Fix for #21005: Redirect HTTP/HTTPS URLs immediately to avoid mixed-content issues
|
||||
if logo_path.startswith(("http://", "https://")):
|
||||
return RedirectResponse(url=logo_path)
|
||||
|
||||
# Local file handling logic
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
default_site_logo = os.path.join(current_dir, "logo.jpg")
|
||||
default_logo = os.path.join(current_dir, "logo.jpg")
|
||||
|
||||
is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true"
|
||||
|
||||
# Determine assets directory
|
||||
# Priority: LITELLM_ASSETS_PATH env var > default based on is_non_root
|
||||
default_assets_dir = "/var/lib/litellm/assets" if is_non_root else current_dir
|
||||
assets_dir = os.getenv("LITELLM_ASSETS_PATH", default_assets_dir)
|
||||
|
||||
# Try to create assets_dir if it doesn't exist (simple try/except approach)
|
||||
if not os.path.exists(assets_dir):
|
||||
try:
|
||||
os.makedirs(assets_dir, exist_ok=True)
|
||||
verbose_proxy_logger.debug(f"Created assets directory at {assets_dir}")
|
||||
except (PermissionError, OSError) as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Cannot create assets directory at {assets_dir}: {e}. "
|
||||
f"Logo caching may not work. Using current directory for assets."
|
||||
)
|
||||
assets_dir = current_dir
|
||||
|
||||
# Determine default logo path
|
||||
default_logo = (
|
||||
os.path.join(assets_dir, "logo.jpg")
|
||||
if assets_dir != current_dir
|
||||
else default_site_logo
|
||||
)
|
||||
if assets_dir != current_dir and not os.path.exists(default_logo):
|
||||
default_logo = default_site_logo
|
||||
|
||||
cache_dir = assets_dir if os.access(assets_dir, os.W_OK) else current_dir
|
||||
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
|
||||
|
||||
logo_path = os.getenv("UI_LOGO_PATH", default_logo)
|
||||
verbose_proxy_logger.debug("Reading logo from path: %s", logo_path)
|
||||
|
||||
# If UI_LOGO_PATH points to a local file, serve it directly (skip cache)
|
||||
if logo_path != default_logo and not logo_path.startswith(("http://", "https://")):
|
||||
if os.path.exists(logo_path):
|
||||
return FileResponse(logo_path, media_type="image/jpeg")
|
||||
# Custom path doesn't exist — fall back to default
|
||||
verbose_proxy_logger.warning(
|
||||
f"UI_LOGO_PATH '{logo_path}' does not exist, falling back to default logo"
|
||||
)
|
||||
if not logo_path:
|
||||
logo_path = default_logo
|
||||
|
||||
# [OPTIMIZATION] For HTTP URLs and default logo, check if the cached image exists
|
||||
if os.path.exists(cache_path):
|
||||
return FileResponse(cache_path, media_type="image/jpeg")
|
||||
|
||||
# Check if the logo path is an HTTP/HTTPS URL
|
||||
if logo_path.startswith(("http://", "https://")):
|
||||
try:
|
||||
# Download the image and cache it
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.UI,
|
||||
params={"timeout": 5.0},
|
||||
)
|
||||
response = await async_client.get(logo_path)
|
||||
if response.status_code == 200:
|
||||
# Save the image to a local file
|
||||
with open(cache_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
|
||||
# Return the cached image as a FileResponse
|
||||
return FileResponse(cache_path, media_type="image/jpeg")
|
||||
else:
|
||||
# Handle the case when the image cannot be downloaded
|
||||
return FileResponse(default_logo, media_type="image/jpeg")
|
||||
except Exception as e:
|
||||
# Handle any exceptions during the download (e.g., timeout, connection error)
|
||||
verbose_proxy_logger.debug(f"Error downloading logo from {logo_path}: {e}")
|
||||
return FileResponse(default_logo, media_type="image/jpeg")
|
||||
else:
|
||||
# Return the local image file if the logo path is not an HTTP/HTTPS URL
|
||||
if os.path.exists(logo_path):
|
||||
return FileResponse(logo_path, media_type="image/jpeg")
|
||||
|
||||
return FileResponse(default_logo, media_type="image/jpeg")
|
||||
|
||||
|
||||
@app.get("/get_favicon", include_in_schema=False)
|
||||
|
|
|
|||
|
|
@ -1,89 +1,24 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest import mock
|
||||
|
||||
# Standard path insertion
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import pytest
|
||||
import httpx
|
||||
import os
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
def test_get_image_redirect_behavior(monkeypatch):
|
||||
"""Verify 307 Redirect behavior for remote URLs."""
|
||||
client = TestClient(app)
|
||||
|
||||
# Safely set the environment variable and ensure auto-cleanup
|
||||
monkeypatch.setenv("UI_LOGO_PATH", "http://example.com/logo.png")
|
||||
|
||||
response = client.get("/get_image", follow_redirects=False)
|
||||
assert response.status_code == 307
|
||||
assert response.headers["location"] == "http://example.com/logo.png"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_image_error_handling():
|
||||
"""
|
||||
Test that get_image handles network errors gracefully and doesn't hang.
|
||||
"""
|
||||
# Set an unreachable URL
|
||||
os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg"
|
||||
|
||||
# Clear cache
|
||||
parent_dir = os.path.dirname(
|
||||
os.path.dirname(
|
||||
app.__file__
|
||||
if hasattr(app, "__file__")
|
||||
else "litellm/proxy/proxy_server.py"
|
||||
)
|
||||
)
|
||||
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
|
||||
if os.path.exists(cache_path):
|
||||
os.remove(cache_path)
|
||||
|
||||
# Mock AsyncHTTPHandler to simulate a timeout or connection error
|
||||
with mock.patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
||||
) as mock_get:
|
||||
mock_get.side_effect = httpx.ConnectError("Network is unreachable")
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
|
||||
) as ac:
|
||||
response = await ac.get("/get_image")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/jpeg"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_image_cache_logic():
|
||||
"""
|
||||
Test that once cached, get_image doesn't hit the network.
|
||||
"""
|
||||
os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg"
|
||||
|
||||
# Clear cache
|
||||
parent_dir = os.path.dirname(
|
||||
os.path.dirname(
|
||||
app.__file__
|
||||
if hasattr(app, "__file__")
|
||||
else "litellm/proxy/proxy_server.py"
|
||||
)
|
||||
)
|
||||
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
|
||||
if os.path.exists(cache_path):
|
||||
os.remove(cache_path)
|
||||
|
||||
# Mock response
|
||||
mock_response = mock.Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.content = b"fake image data"
|
||||
|
||||
with mock.patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
|
||||
) as ac:
|
||||
# First call - should hit download logic
|
||||
response1 = await ac.get("/get_image")
|
||||
assert response1.status_code == 200
|
||||
assert mock_get.call_count == 1
|
||||
|
||||
# Second call - should hit cache
|
||||
response2 = await ac.get("/get_image")
|
||||
assert response2.status_code == 200
|
||||
# If cache works, mock_get shouldn't be called again
|
||||
assert mock_get.call_count == 1
|
||||
def test_get_image_local_fallback(monkeypatch):
|
||||
"""Verify fallback to default logo when environment variable is missing."""
|
||||
client = TestClient(app)
|
||||
monkeypatch.delenv("UI_LOGO_PATH", raising=False)
|
||||
|
||||
response = client.get("/get_image")
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/jpeg"
|
||||
Loading…
Add table
Reference in a new issue