mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(health): make the image_edit health probe moderation-safe
The image_edit health probe sent a 512x512 solid-gray PNG with the generic chat prompt "test from litellm", an ambiguous pair OpenAI's gpt-image-1 output moderation sometimes rejects as moderation_blocked, which reported a working deployment as unhealthy. The probe now sends a blue circle on a white background with a descriptive edit prompt, and a provider moderation verdict (ContentPolicyViolationError or a moderation_blocked error body) is treated as proof the endpoint works rather than as an unhealthy deployment.
This commit is contained in:
mateo-berri
parent
f57e4b812c
commit
416984a5e0
2 changed files with 123 additions and 40 deletions
|
|
@ -3,19 +3,25 @@ Helper functions for health check calls.
|
|||
"""
|
||||
|
||||
import base64
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
# Minimal PDF for health checks - base64 encoded 1-page PDF with just "test"
|
||||
TEST_PDF_URL = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="
|
||||
|
||||
# Minimal image for health checks - base64 encoded 512x512 solid-gray PNG
|
||||
TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFlklEQVR42u3VMQEAAAzCMKQjHQ97l0jo0xSAlyIBgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAUgAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAADcDrctaAb6XeXAAAAAASUVORK5CYII="
|
||||
# Minimal image for health checks - base64 encoded 512x512 blue circle on a white background PNG
|
||||
TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAJk0lEQVR42u3VQREAIRADwVWCOmTjBVzwSLorCri6nbkAVBpPACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAZdY+HgEBgIRr/meeGgGA8EMvDAgAuPh6gACAi68HCAA4+mKAAICjLwYIALj7SoAAgLuvBAgA7r4pAQKAu29KgADg7psSIAA4/SYDCADuvikBAoDTbzKAAOD0mwwgADj9JgMIAE6/yQACgNNvMoAA4PSbDCAAOP0mAwgATr/JAAKA668BIAA4/TKAAOD0mwwgALj+pgEIAE6/yQACgOtvGoAA4PSbDCAAuP6mAQgATr/JAAKA628agADg9JsMIAC4/qYBCACuv2kAAoDrbxqAAOD0mwwgALj+pgEIAK6/aQACgOtvGoAA4PqbBiAArr+ZBiAATr+ZDCAArr+ZBiAArr+ZBiAArr+ZBiAArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggAAmAmAAKA62+mAQKA62+mAQKA62/m1xYAXH/TAAQA1980AAHA9TcNQAAEwEwAEADX30wDEADX30wDEADX30wDEADX30wDEAABMBMABMD1N9MABMD1N9MABEAAzAQAAXD9zTQAAXD9zTRAABAAMwEQAFx/Mw0QAFx/Mw0QAATATAAEANffTAMEANffTAMEAAEwEwABwPU30wABQADMBEAAXH8z0wABcP3NTAMEQADMTAAEwPU3Mw0QAAEwMwEQANffTAMQAAEwEwAEwPU30wAEQADMBAABcP3NNAABEAAzAUAAXH8zDUAABMBMABAA199MAwQAATATAAHA9TfTAAFAAMwEQABw/c00QAAQADMBEAAEwEwABMD1NzMNEAABMDMBEADX38w0QAAEwMwEQAAEwMwEQABcfzPTAAEQADMTAAEQADMTAAFw/c1MAwRAAMxMAARAAMxMAATA9TczDRAAATATAARAAMwEAAFw/c00AAEQADMBQAAEwEwAEADX30wDBAABMBMAAUAAzARAABAAMwEQAFx/Mw0QAAEwMwEQAAEwMwEQAAEwMwEQANffzDRAAATAzARAAATAzARAAATAzARAAATAzARAAFx/M9MAARAAMxMAARAAMxMAARAAMxMAARAAMxMAAXD9zUwDBEAAzEwABEAAzAQAARAAMwFAAATATAAQAAEwEwABQADMBEAAEAAzARAAXH8zDRAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAA19/MNEAANMDM9UcABMBMABAAATATAAHwBAJgJgACgACYCYAAIABmAiAA+IvMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAANMDMXH8BEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzAUAABMBMABAADTBz/REAATATAAFAAMwEQAAQADMBEAAEwEwABAANMHP9BUAAzEwABEAAzEwABEAAzEwABEAAzEwABEADzMz1FwABMDMBEAABMDMBEAABMDMBEAANMDPXXwAEwMwEQAAEwMwEQAAEwMwEQAA0wMxcfwEQADMTAAEQADMBQAA0wMz1RwAEwEwAEAABMBMABEADzFx/AUAAzARAABAAMwEQADTAzPUXAATATAAEAAEwEwABQAPMXH8BEAAzEwABEAAzEwAB0AAzc/0FQADMTAAEQAPMzPUXAAEwMwEQAAEwMwEQAA0wM9dfAATAzARAADTAzFx/ARAAMwFAADTAzPVHAATATAAQAA0wc/0RAAEwEwAEQAPMXH8EQADMBAAB0AAz118AEAAzARAANMDM9RcABMBMAAQADTBz/QUAATATAAFAA8xcfwFAA8xcfwFAAMwEQADQADPXXwAEwMwEQAA0wMxcfwHQADNz/QVAAMxMAARAA8xcfwRAA8xcfwRAAMwEAAHQADPXHwHQADPXHwEQADMBQAA0wMz1RwA0wMz1RwAEwEwAEAANMHP9BQANMHP9BQANMHP9BQANMHP9BQABMBMAAUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzPVHABAAEwAEAA0w1x8BQAPM9UcA0ABz/REANMBcfwQADTDXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwGQATOnHwHQADPXHwHQADPXXwDQADPXXwDQADPXXwDQADPXXwCQAXP6EQA0wFx/BAANMNcfAUADzPVHAJABc/oRADTAXH8EABkwpx8BQAPM9UcAkAFz+hEANMBcfwQAGTCnHwFAA8z1RwCQAXP6EQBkwJx+BAANMNcfAUAGzOlHAJABc/oRAGTA6QcBQAacfhAAZMDpRwBABpx+BABkwOlHAEAGnH4EAJTA3UcAQAacfgQAlMDdRwBACdx9BACUwN1HAEAJ3H0EAJTA3UcAQAwcfQQAmmLgsyIA0NIDHw4BgJYe+DQIAISHwVMjAJDQDI+AAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACACAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAIAABNHpialFcmLajuAAAAAElFTkSuQmCC"
|
||||
|
||||
|
||||
IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
|
||||
"Add a small yellow star in the top right corner of this simple drawing of a blue circle on a white background"
|
||||
)
|
||||
|
||||
|
||||
def get_image_file_for_health_check() -> bytes:
|
||||
|
|
@ -121,6 +127,17 @@ class HealthCheckHelpers:
|
|||
else:
|
||||
return await litellm.acompletion(**model_params)
|
||||
|
||||
@staticmethod
|
||||
async def _image_edit_health_check(edit_request: Callable[[], Awaitable["ImageResponse"]]) -> "ImageResponse":
|
||||
import litellm
|
||||
|
||||
try:
|
||||
return await edit_request()
|
||||
except litellm.BadRequestError as e:
|
||||
if isinstance(e, litellm.ContentPolicyViolationError) or "moderation_blocked" in str(e):
|
||||
return litellm.ImageResponse()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def get_mode_handlers(
|
||||
model: str,
|
||||
|
|
@ -195,10 +212,12 @@ class HealthCheckHelpers:
|
|||
**_filter_model_params(model_params=model_params),
|
||||
prompt=prompt,
|
||||
),
|
||||
"image_edit": lambda: litellm.aimage_edit(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
image=get_image_file_for_health_check(),
|
||||
prompt=prompt or "test",
|
||||
"image_edit": lambda: HealthCheckHelpers._image_edit_health_check(
|
||||
edit_request=lambda: litellm.aimage_edit(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
image=get_image_file_for_health_check(),
|
||||
prompt=IMAGE_EDIT_HEALTH_CHECK_PROMPT,
|
||||
),
|
||||
),
|
||||
"video_generation": lambda: litellm.avideo_generation(
|
||||
**_filter_model_params(model_params=model_params),
|
||||
|
|
|
|||
|
|
@ -1,19 +1,40 @@
|
|||
"""Test health check helper functions"""
|
||||
|
||||
import struct
|
||||
import zlib
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
|
||||
from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
|
||||
from litellm.litellm_core_utils.health_check_helpers import (
|
||||
IMAGE_EDIT_HEALTH_CHECK_PROMPT,
|
||||
HealthCheckHelpers,
|
||||
)
|
||||
from litellm.main import ahealth_check
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
|
||||
|
||||
|
||||
def _distinct_rgb_colors(png: bytes) -> set[bytes]:
|
||||
width = int.from_bytes(png[16:20], "big")
|
||||
idat_parts = []
|
||||
offset = 8
|
||||
while offset < len(png):
|
||||
(length,) = struct.unpack(">I", png[offset : offset + 4])
|
||||
if png[offset + 4 : offset + 8] == b"IDAT":
|
||||
idat_parts.append(png[offset + 8 : offset + 8 + length])
|
||||
offset += 12 + length
|
||||
raw = zlib.decompress(b"".join(idat_parts))
|
||||
row_size = 1 + width * 3
|
||||
rows = [raw[i : i + row_size] for i in range(0, len(raw), row_size)]
|
||||
assert all(row[0] == 0 for row in rows)
|
||||
return {bytes(row[i : i + 3]) for row in rows for i in range(1, row_size, 3)}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_edit_health_check_handler_uses_png_and_prompt():
|
||||
async def test_image_edit_health_check_handler_uses_descriptive_prompt_and_multicolor_png():
|
||||
model_params = {"model": "openai/gpt-image-1", "api_key": "sk-test"}
|
||||
mode_handlers = HealthCheckHelpers.get_mode_handlers(
|
||||
model="gpt-image-1",
|
||||
|
|
@ -31,20 +52,76 @@ async def test_image_edit_health_check_handler_uses_png_and_prompt():
|
|||
model="gpt-image-1",
|
||||
custom_llm_provider="openai",
|
||||
model_params=model_params,
|
||||
prompt="edit this image",
|
||||
prompt="test from litellm",
|
||||
)["image_edit"]()
|
||||
|
||||
assert mock_aimage_edit.call_count == 2
|
||||
default_call = mock_aimage_edit.call_args_list[0].kwargs
|
||||
explicit_call = mock_aimage_edit.call_args_list[1].kwargs
|
||||
assert default_call["model"] == "openai/gpt-image-1"
|
||||
assert default_call["prompt"] == "test"
|
||||
assert explicit_call["prompt"] == "edit this image"
|
||||
image = default_call["image"]
|
||||
for handler_call in mock_aimage_edit.call_args_list:
|
||||
assert handler_call.kwargs["model"] == "openai/gpt-image-1"
|
||||
assert handler_call.kwargs["prompt"] == IMAGE_EDIT_HEALTH_CHECK_PROMPT
|
||||
image = mock_aimage_edit.call_args_list[0].kwargs["image"]
|
||||
assert isinstance(image, bytes)
|
||||
assert image.startswith(b"\x89PNG")
|
||||
assert int.from_bytes(image[16:20], "big") == 512
|
||||
assert int.from_bytes(image[20:24], "big") == 512
|
||||
assert len(_distinct_rgb_colors(image)) >= 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ahealth_check_image_edit_treats_content_policy_violation_as_healthy():
|
||||
moderation_error = litellm.ContentPolicyViolationError(
|
||||
message="Your request was rejected as a result of our safety system.",
|
||||
model="gpt-image-1",
|
||||
llm_provider="openai",
|
||||
)
|
||||
with patch( # test-quality-ok: the public health-check path has no dependency injection seam
|
||||
"litellm.aimage_edit", new_callable=AsyncMock, side_effect=moderation_error
|
||||
):
|
||||
result = await ahealth_check(
|
||||
{"model": "gpt-image-1", "api_key": "sk-test"},
|
||||
mode="image_edit",
|
||||
)
|
||||
|
||||
assert "error" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ahealth_check_image_edit_treats_moderation_blocked_code_as_healthy():
|
||||
moderation_blocked = litellm.BadRequestError(
|
||||
message=(
|
||||
'{"error": {"code": "moderation_blocked", "message": "Your request was blocked", '
|
||||
'"moderation_stage": "output", "type": "invalid_request_error"}}'
|
||||
),
|
||||
model="gpt-image-1",
|
||||
llm_provider="openai",
|
||||
)
|
||||
with patch( # test-quality-ok: the public health-check path has no dependency injection seam
|
||||
"litellm.aimage_edit", new_callable=AsyncMock, side_effect=moderation_blocked
|
||||
):
|
||||
result = await ahealth_check(
|
||||
{"model": "gpt-image-1", "api_key": "sk-test"},
|
||||
mode="image_edit",
|
||||
)
|
||||
|
||||
assert "error" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ahealth_check_image_edit_still_fails_on_non_moderation_errors():
|
||||
auth_error = litellm.AuthenticationError(
|
||||
message="Incorrect API key provided",
|
||||
llm_provider="openai",
|
||||
model="gpt-image-1",
|
||||
)
|
||||
with patch( # test-quality-ok: the public health-check path has no dependency injection seam
|
||||
"litellm.aimage_edit", new_callable=AsyncMock, side_effect=auth_error
|
||||
):
|
||||
result = await ahealth_check(
|
||||
{"model": "gpt-image-1", "api_key": "sk-bad"},
|
||||
mode="image_edit",
|
||||
)
|
||||
|
||||
assert "error" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -88,9 +165,7 @@ def test_update_model_params_with_health_check_tracking_information():
|
|||
|
||||
# Verify that litellm_metadata was added
|
||||
assert "litellm_metadata" in result
|
||||
assert result["litellm_metadata"]["tags"] == [
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
|
||||
]
|
||||
assert result["litellm_metadata"]["tags"] == [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME]
|
||||
|
||||
# Verify the auth setup was called
|
||||
mock_add_auth.assert_called_once()
|
||||
|
|
@ -169,16 +244,12 @@ async def test_ahealth_check_failure_masks_raw_request_headers():
|
|||
if "Authorization" in headers:
|
||||
auth_header = headers["Authorization"]
|
||||
# Should be masked (e.g., "Be****90" or similar)
|
||||
assert (
|
||||
auth_header != f"Bearer {test_api_key}"
|
||||
), "Authorization header must be masked"
|
||||
assert (
|
||||
auth_header != test_api_key
|
||||
), "API key must not appear in Authorization header"
|
||||
assert auth_header != f"Bearer {test_api_key}", "Authorization header must be masked"
|
||||
assert auth_header != test_api_key, "API key must not appear in Authorization header"
|
||||
# Masked headers typically have asterisks or are truncated
|
||||
assert "*" in auth_header or len(auth_header) < len(
|
||||
f"Bearer {test_api_key}"
|
||||
), f"Authorization header should be masked but got: {auth_header}"
|
||||
assert "*" in auth_header or len(auth_header) < len(f"Bearer {test_api_key}"), (
|
||||
f"Authorization header should be masked but got: {auth_header}"
|
||||
)
|
||||
|
||||
# Content-Type should remain unmasked (not sensitive)
|
||||
if "Content-Type" in headers:
|
||||
|
|
@ -257,9 +328,7 @@ async def test_batch_health_check_skips_bridge_when_no_logging_obj():
|
|||
"litellm_metadata": litellm_metadata,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.alist_batches", new_callable=AsyncMock, return_value={}
|
||||
) as mock_alist:
|
||||
with patch("litellm.alist_batches", new_callable=AsyncMock, return_value={}) as mock_alist:
|
||||
await HealthCheckHelpers._batch_health_check(
|
||||
custom_llm_provider="openai",
|
||||
model_params={"model": "openai/gpt-4"},
|
||||
|
|
@ -283,9 +352,7 @@ async def test_batch_health_check_uses_alist_batches_for_supported_providers():
|
|||
"litellm_metadata": litellm_metadata,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.alist_batches", new_callable=AsyncMock, return_value={}
|
||||
) as mock_alist:
|
||||
with patch("litellm.alist_batches", new_callable=AsyncMock, return_value={}) as mock_alist:
|
||||
await HealthCheckHelpers._batch_health_check(
|
||||
custom_llm_provider=provider,
|
||||
model_params={"model": f"{provider}/some-model"},
|
||||
|
|
@ -344,9 +411,7 @@ async def test_realtime_health_check_uses_model_level_vertex_params():
|
|||
|
||||
fake_vertex_base = MagicMock()
|
||||
fake_vertex_base.get_vertex_region = MagicMock(return_value="us-central1")
|
||||
fake_vertex_base._ensure_access_token_async = AsyncMock(
|
||||
return_value=("model-level-token", "model-level-project")
|
||||
)
|
||||
fake_vertex_base._ensure_access_token_async = AsyncMock(return_value=("model-level-token", "model-level-project"))
|
||||
connect_calls = []
|
||||
|
||||
with (
|
||||
|
|
@ -381,8 +446,7 @@ async def test_realtime_health_check_uses_model_level_vertex_params():
|
|||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert connect_calls[0]["url"] == (
|
||||
"wss://us-central1-aiplatform.googleapis.com/ws/"
|
||||
"google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
"wss://us-central1-aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
assert connect_calls[0]["additional_headers"] == {
|
||||
"Authorization": "Bearer model-level-token",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue