mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(proxy): resolve team callbacks when key logging is empty and report GCS flush failures
An empty key-level logging list is now treated as unset, so team logging, deprecated team callback_settings and default_team_settings apply. /key/health tests the effective callbacks, flushes the GCS logger explicitly and reports failed uploads. Failed GCS batches stay queued for the next flush instead of being dropped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
0f03fc2985
commit
7b459ac1bc
9 changed files with 413 additions and 71 deletions
|
|
@ -3,6 +3,7 @@ import hashlib
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
|
|
@ -68,10 +69,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
||||
if logging_payload is None:
|
||||
raise ValueError("standard_logging_object not found in kwargs")
|
||||
# When queue is at maxsize, flush immediately to make room (no blocking, no data dropped)
|
||||
if self.log_queue.full():
|
||||
await self.flush_queue()
|
||||
await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj))
|
||||
await self._enqueue(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj))
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("GCS Bucket logging error: %s", e)
|
||||
|
|
@ -87,14 +85,35 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
||||
if logging_payload is None:
|
||||
raise ValueError("standard_logging_object not found in kwargs")
|
||||
# When queue is at maxsize, flush immediately to make room (no blocking, no data dropped)
|
||||
if self.log_queue.full():
|
||||
await self.flush_queue()
|
||||
await self.log_queue.put(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj))
|
||||
await self._enqueue(GCSLogQueueItem(payload=logging_payload, kwargs=kwargs, response_obj=response_obj))
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("GCS Bucket logging error: %s", e)
|
||||
|
||||
async def _enqueue(self, item: GCSLogQueueItem) -> None:
|
||||
if self.log_queue.full():
|
||||
await self.flush_queue()
|
||||
if self.log_queue.full():
|
||||
self.log_queue.get_nowait()
|
||||
verbose_logger.error("GCS Bucket log queue still full after flush, dropped the oldest queued event")
|
||||
self.log_queue.put_nowait(item)
|
||||
|
||||
def _requeue(self, items: Sequence[GCSLogQueueItem]) -> None:
|
||||
dropped: Final = sum(1 for item in items if not self._put_nowait_or_drop(item))
|
||||
verbose_logger.error(
|
||||
"GCS Bucket upload failed for %s events, %s kept in queue for the next flush, %s dropped (queue full)",
|
||||
len(items),
|
||||
len(items) - dropped,
|
||||
dropped,
|
||||
)
|
||||
|
||||
def _put_nowait_or_drop(self, item: GCSLogQueueItem) -> bool:
|
||||
try:
|
||||
self.log_queue.put_nowait(item)
|
||||
except asyncio.QueueFull:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _drain_queue_batch(self) -> list[GCSLogQueueItem]:
|
||||
"""
|
||||
Drain items from the queue (non-blocking), respecting batch_size limit.
|
||||
|
|
@ -219,17 +238,19 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
verbose_logger.exception("GCS Bucket error logging batch payload to GCS bucket: %s", e)
|
||||
return (success_count, error_count)
|
||||
|
||||
async def _send_individual_logs(self, items: list[GCSLogQueueItem]) -> None:
|
||||
async def _send_individual_logs(self, items: list[GCSLogQueueItem]) -> GCSFlushResult:
|
||||
"""
|
||||
Send each log individually as separate GCS objects (legacy behavior).
|
||||
This is used when GCS_USE_BATCHED_LOGGING is disabled.
|
||||
"""
|
||||
for item in items:
|
||||
await self._send_single_log_item(item)
|
||||
failed_items: Final = tuple([item for item in items if not await self._send_single_log_item(item)])
|
||||
if failed_items:
|
||||
self._requeue(failed_items)
|
||||
return GCSFlushResult(sent=len(items) - len(failed_items), failed=len(failed_items))
|
||||
|
||||
async def _send_single_log_item(self, item: GCSLogQueueItem) -> None:
|
||||
async def _send_single_log_item(self, item: GCSLogQueueItem) -> bool:
|
||||
"""
|
||||
Send a single log item to GCS as an individual object.
|
||||
Send a single log item to GCS as an individual object. Returns whether the upload succeeded.
|
||||
"""
|
||||
try:
|
||||
gcs_logging_config: Final[GCSLoggingConfig] = await self.get_gcs_logging_config(item["kwargs"])
|
||||
|
|
@ -254,8 +275,25 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("GCS Bucket error logging individual payload to GCS bucket: %s", e)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def async_send_batch(self):
|
||||
async def _send_grouped_batches(self, items: list[GCSLogQueueItem]) -> GCSFlushResult:
|
||||
results: Final = tuple(
|
||||
[
|
||||
(group_items, await self._send_grouped_batch(group_items, config_key))
|
||||
for config_key, group_items in self._group_items_by_config(items).items()
|
||||
]
|
||||
)
|
||||
for group_items, (_, group_failed) in results:
|
||||
if group_failed:
|
||||
self._requeue(group_items)
|
||||
return GCSFlushResult(
|
||||
sent=sum(group_sent for _, (group_sent, _) in results),
|
||||
failed=sum(group_failed for _, (_, group_failed) in results),
|
||||
)
|
||||
|
||||
async def async_send_batch(self) -> None:
|
||||
"""
|
||||
Process queued logs - sends logs to GCS Bucket.
|
||||
|
||||
|
|
@ -264,18 +302,17 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
|
||||
If disabled, sends each log individually as separate GCS objects (legacy behavior).
|
||||
"""
|
||||
await self._send_queued_events()
|
||||
|
||||
async def _send_queued_events(self) -> GCSFlushResult:
|
||||
items_to_process: Final = self._drain_queue_batch()
|
||||
|
||||
if not items_to_process:
|
||||
return
|
||||
return GCSFlushResult(sent=0, failed=0)
|
||||
|
||||
if self.use_batched_logging:
|
||||
grouped_items: Final = self._group_items_by_config(items_to_process)
|
||||
|
||||
for config_key, group_items in grouped_items.items():
|
||||
await self._send_grouped_batch(group_items, config_key)
|
||||
else:
|
||||
await self._send_individual_logs(items_to_process)
|
||||
return await self._send_grouped_batches(items_to_process)
|
||||
return await self._send_individual_logs(items_to_process)
|
||||
|
||||
def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str:
|
||||
"""
|
||||
|
|
@ -355,12 +392,16 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
def _get_object_date_from_datetime(self, datetime_obj: datetime) -> str:
|
||||
return datetime_obj.strftime("%Y-%m-%d")
|
||||
|
||||
async def flush_queue(self):
|
||||
async def flush_queue(self) -> None:
|
||||
"""
|
||||
Override flush_queue to work with asyncio.Queue.
|
||||
"""
|
||||
await self.async_send_batch()
|
||||
await self.flush_queue_and_report()
|
||||
|
||||
async def flush_queue_and_report(self) -> GCSFlushResult:
|
||||
result: Final = await self._send_queued_events()
|
||||
self.last_flush_time = time.time()
|
||||
return result
|
||||
|
||||
async def periodic_flush(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import enum
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
|
||||
|
|
@ -4441,7 +4441,7 @@ class CurrentItemRateLimit(TypedDict):
|
|||
|
||||
|
||||
class LoggingCallbackStatus(TypedDict, total=False):
|
||||
callbacks: list[str]
|
||||
callbacks: ReadOnly[Sequence[str]]
|
||||
status: Literal["healthy", "unhealthy"]
|
||||
details: str | None
|
||||
|
||||
|
|
|
|||
|
|
@ -904,18 +904,21 @@ def _get_validated_callback_metadata(item: dict, *, source: str) -> AddTeamCallb
|
|||
class KeyAndTeamLoggingSettings:
|
||||
"""
|
||||
Helper class to get the dynamic logging settings for the key and team
|
||||
|
||||
An empty ``logging`` list is the same as no ``logging`` key: both return ``None`` so the
|
||||
caller falls through to the next level. Disabling a callback is ``litellm_disabled_callbacks``.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
|
||||
if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata:
|
||||
return decrypt_callback_vars(user_api_key_dict.metadata).get("logging")
|
||||
return decrypt_callback_vars(user_api_key_dict.metadata).get("logging") or None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
|
||||
if user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata:
|
||||
return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging")
|
||||
return decrypt_callback_vars(user_api_key_dict.team_metadata).get("logging") or None
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -1029,8 +1032,8 @@ def resolve_tenant_otel_destinations(
|
|||
|
||||
Key settings win over team settings outright, the same precedence
|
||||
``_get_dynamic_logging_metadata`` applies, so one caller never exports the same
|
||||
backend to two accounts. An empty key-level list counts as configured, since that
|
||||
is what disabling a key's callbacks writes. Returns empty when OTEL V2 is off, when
|
||||
backend to two accounts. An empty key-level list is unset and falls through to the
|
||||
team, the way the runtime parser reads it. Returns empty when OTEL V2 is off, when
|
||||
neither level named a destination-capable backend, or when the config is
|
||||
incomplete, and the request then keeps the operator's own exporters.
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.constants import (
|
|||
MINIMUM_CUSTOM_KEY_LENGTH,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.models.credentials import CredentialItem
|
||||
|
|
@ -6973,7 +6974,9 @@ async def key_health(
|
|||
Check the health of the key
|
||||
|
||||
Checks:
|
||||
- If key based logging is configured correctly - sends a test log
|
||||
- If the logging that applies to this key (key metadata, team metadata, or
|
||||
`default_team_settings` in the config) is configured correctly - sends a test log
|
||||
and, for gcs_bucket, flushes the queue and reports the upload result
|
||||
|
||||
Usage
|
||||
|
||||
|
|
@ -7015,29 +7018,40 @@ async def key_health(
|
|||
}
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
_get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # the request-time resolver; the health check must report the same callbacks a request would use
|
||||
)
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
try:
|
||||
# Get the key's metadata
|
||||
key_metadata: Final = user_api_key_dict.metadata
|
||||
|
||||
health_status: Final[KeyHealthResponse] = KeyHealthResponse(
|
||||
key="healthy",
|
||||
logging_callbacks=None,
|
||||
)
|
||||
|
||||
# Check if logging is configured in metadata
|
||||
if key_metadata and "logging" in key_metadata:
|
||||
logging_statuses: Final = await test_key_logging(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request=request,
|
||||
key_logging=decrypt_callback_vars(key_metadata)["logging"],
|
||||
_raise_if_key_logging_missing_callback_name(decrypt_callback_vars(key_metadata)["logging"])
|
||||
|
||||
callback_settings: Final = _get_dynamic_logging_metadata(
|
||||
user_api_key_dict=user_api_key_dict, proxy_config=proxy_config
|
||||
)
|
||||
logging_callbacks: Final = (
|
||||
()
|
||||
if callback_settings is None
|
||||
else tuple(
|
||||
dict.fromkeys(
|
||||
(*(callback_settings.success_callback or ()), *(callback_settings.failure_callback or ()))
|
||||
)
|
||||
)
|
||||
health_status["logging_callbacks"] = logging_statuses
|
||||
)
|
||||
if not logging_callbacks:
|
||||
return KeyHealthResponse(key="healthy", logging_callbacks=None)
|
||||
|
||||
# Check if any logging callback is unhealthy
|
||||
if logging_statuses.get("status") == "unhealthy":
|
||||
health_status["key"] = "unhealthy"
|
||||
|
||||
return KeyHealthResponse(**health_status)
|
||||
logging_statuses: Final = await test_key_logging(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request=request,
|
||||
logging_callbacks=logging_callbacks,
|
||||
)
|
||||
return KeyHealthResponse(
|
||||
key="unhealthy" if logging_statuses.get("status") == "unhealthy" else "healthy",
|
||||
logging_callbacks=logging_statuses,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise ProxyException(
|
||||
|
|
@ -7072,31 +7086,41 @@ async def _can_user_query_key_info(
|
|||
return False
|
||||
|
||||
|
||||
def _raise_if_key_logging_missing_callback_name(key_logging: Sequence[Mapping[str, str]]) -> None:
|
||||
if any(callback.get("callback_name") is None for callback in key_logging):
|
||||
raise ValueError("callback_name is required in key_logging")
|
||||
|
||||
|
||||
async def flush_gcs_and_describe_failures(gcs_logger: CustomLogger | None) -> str | None:
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
|
||||
if not isinstance(gcs_logger, GCSBucketLogger):
|
||||
return "gcs_bucket callback was selected but no GCS logger was initialized"
|
||||
flush_result: Final = await gcs_logger.flush_queue_and_report()
|
||||
if flush_result.failed == 0:
|
||||
return None
|
||||
return f"GCS upload failed for {flush_result.failed} event(s), {flush_result.sent} uploaded"
|
||||
|
||||
|
||||
async def test_key_logging(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request: Request,
|
||||
key_logging: Sequence[Mapping[str, str]],
|
||||
logging_callbacks: Sequence[str],
|
||||
) -> LoggingCallbackStatus:
|
||||
"""
|
||||
Test the key-based logging
|
||||
Test the logging callbacks that apply to this key
|
||||
|
||||
- Test that key logging is correctly formatted and all args are passed correctly
|
||||
- Make a mock completion call -> user can check if it's correctly logged
|
||||
- For gcs_bucket, flush the queue and report whether the upload succeeded
|
||||
- Check if any logger.exceptions were triggered -> if they were then returns it to the user client side
|
||||
"""
|
||||
import logging
|
||||
from io import StringIO
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import get_custom_logger_compatible_class
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.proxy_server import general_settings, proxy_config
|
||||
|
||||
logging_callbacks: Final[list[str]] = []
|
||||
for callback in key_logging:
|
||||
if callback.get("callback_name") is not None:
|
||||
logging_callbacks.append(callback["callback_name"])
|
||||
else:
|
||||
raise ValueError("callback_name is required in key_logging")
|
||||
|
||||
log_capture_string: Final = StringIO()
|
||||
ch: Final = logging.StreamHandler(log_capture_string)
|
||||
ch.setLevel(logging.ERROR)
|
||||
|
|
@ -7131,21 +7155,28 @@ async def test_key_logging(
|
|||
|
||||
await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event
|
||||
|
||||
# Check if any logger exceptions were triggered
|
||||
gcs_failure: Final = (
|
||||
await flush_gcs_and_describe_failures(get_custom_logger_compatible_class("gcs_bucket"))
|
||||
if "gcs_bucket" in logging_callbacks
|
||||
else None
|
||||
)
|
||||
|
||||
log_contents: Final = log_capture_string.getvalue()
|
||||
logger.removeHandler(ch)
|
||||
if log_contents:
|
||||
if gcs_failure is not None or log_contents:
|
||||
return LoggingCallbackStatus(
|
||||
callbacks=logging_callbacks,
|
||||
status="unhealthy",
|
||||
details=f"Logger exceptions triggered, system is unhealthy: {log_contents}",
|
||||
)
|
||||
else:
|
||||
return LoggingCallbackStatus(
|
||||
callbacks=logging_callbacks,
|
||||
status="healthy",
|
||||
details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ",
|
||||
details=f"Logger exceptions triggered, system is unhealthy: {gcs_failure or ''} {log_contents}".strip(),
|
||||
)
|
||||
return LoggingCallbackStatus(
|
||||
callbacks=logging_callbacks,
|
||||
status="healthy",
|
||||
details=(
|
||||
"No logger exceptions triggered, system is healthy. "
|
||||
f"Manually check if logs were sent to {', '.join(logging_callbacks)}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_KEY_ALIAS_PATTERN: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.@]{0,253}[a-zA-Z0-9]$")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -33,3 +34,9 @@ class GCSLogQueueItem(TypedDict):
|
|||
payload: StandardLoggingPayload
|
||||
kwargs: dict[str, Any]
|
||||
response_obj: Any | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GCSFlushResult:
|
||||
sent: int
|
||||
failed: int
|
||||
|
|
|
|||
104
tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py
Normal file
104
tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
import asyncio
|
||||
import json
|
||||
from typing import Any, Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.types.integrations.gcs_bucket import GCSFlushResult, GCSLoggingConfig, GCSLogQueueItem
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
class _FakeUploadGCSLogger(GCSBucketLogger):
|
||||
"""Skips GCP auth; an upload raises when it carries any id in `failing_ids`, otherwise it is recorded"""
|
||||
|
||||
def __init__(self, queue_maxsize: int = 0) -> None:
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: GCS logging is premium-gated
|
||||
super().__init__(bucket_name="test-bucket")
|
||||
self.log_queue = asyncio.Queue(maxsize=queue_maxsize)
|
||||
self.failing_ids: frozenset[str] = frozenset()
|
||||
self.uploaded: list[list[str]] = []
|
||||
|
||||
async def enqueue(self, request_id: str) -> None:
|
||||
payload: Final = StandardLoggingPayload(id=request_id) # pyright: ignore[reportCallIssue] # partial payload is enough for queueing
|
||||
await self._enqueue(GCSLogQueueItem(payload=payload, kwargs={}, response_obj={"id": request_id}))
|
||||
|
||||
def queued_ids(self) -> list[str]:
|
||||
return [self.log_queue.get_nowait()["payload"]["id"] for _ in range(self.log_queue.qsize())]
|
||||
|
||||
async def get_gcs_logging_config(self, kwargs: dict[str, Any] | None = None) -> GCSLoggingConfig:
|
||||
return GCSLoggingConfig(bucket_name="test-bucket", vertex_instance=None, path_service_account=None)
|
||||
|
||||
async def construct_request_headers(
|
||||
self, service_account_json: str | None, vertex_instance: VertexBase | None = None
|
||||
) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
async def _log_json_data_on_gcs(
|
||||
self, headers: dict[str, str], bucket_name: str, object_name: str, logging_payload: StandardLoggingPayload | str
|
||||
) -> None:
|
||||
ids: Final = (
|
||||
[json.loads(line)["id"] for line in logging_payload.splitlines()]
|
||||
if isinstance(logging_payload, str)
|
||||
else [logging_payload["id"]]
|
||||
)
|
||||
if self.failing_ids.intersection(ids):
|
||||
raise RuntimeError("storage.googleapis.com returned 404")
|
||||
self.uploaded.append(ids)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_batch_stays_queued_and_is_retried_on_the_next_flush():
|
||||
logger = _FakeUploadGCSLogger()
|
||||
logger.failing_ids = frozenset({"req-1"})
|
||||
await logger.enqueue("req-1")
|
||||
await logger.enqueue("req-2")
|
||||
|
||||
failed_flush = await logger.flush_queue_and_report()
|
||||
|
||||
assert failed_flush == GCSFlushResult(sent=0, failed=2)
|
||||
assert logger.log_queue.qsize() == 2
|
||||
assert logger.uploaded == []
|
||||
|
||||
logger.failing_ids = frozenset()
|
||||
retried_flush = await logger.flush_queue_and_report()
|
||||
|
||||
assert retried_flush == GCSFlushResult(sent=2, failed=0)
|
||||
assert logger.log_queue.qsize() == 0
|
||||
assert logger.uploaded == [["req-1", "req-2"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_individual_mode_requeues_only_the_failed_items():
|
||||
logger = _FakeUploadGCSLogger()
|
||||
logger.use_batched_logging = False
|
||||
logger.failing_ids = frozenset({"req-fail"})
|
||||
await logger.enqueue("req-ok")
|
||||
await logger.enqueue("req-fail")
|
||||
|
||||
result = await logger.flush_queue_and_report()
|
||||
|
||||
assert result == GCSFlushResult(sent=1, failed=1)
|
||||
assert logger.uploaded == [["req-ok"]]
|
||||
assert logger.queued_ids() == ["req-fail"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enqueue_on_a_full_queue_whose_flush_failed_drops_the_oldest_event():
|
||||
logger = _FakeUploadGCSLogger(queue_maxsize=2)
|
||||
logger.failing_ids = frozenset({"req-1", "req-2", "req-3"})
|
||||
await logger.enqueue("req-1")
|
||||
await logger.enqueue("req-2")
|
||||
|
||||
await logger.enqueue("req-3")
|
||||
|
||||
assert logger.queued_ids() == ["req-2", "req-3"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_queue_flush_reports_nothing_sent_or_failed():
|
||||
logger = _FakeUploadGCSLogger()
|
||||
|
||||
assert await logger.flush_queue_and_report() == GCSFlushResult(sent=0, failed=0)
|
||||
|
|
@ -1803,15 +1803,17 @@ class TestTenantConfigAgreement:
|
|||
"callback_vars": {"langfuse_public_key": "pk", "langfuse_secret_key": "sk", "langfuse_host": host, **extra},
|
||||
}
|
||||
|
||||
def test_a_key_that_disabled_its_callbacks_does_not_fall_back_to_the_team(self):
|
||||
"""Disabling a key's callbacks stores an empty list, which the sibling parser
|
||||
reads as 'the key configured none'."""
|
||||
def test_a_key_with_an_empty_logging_list_falls_back_to_the_team(self):
|
||||
"""An empty key-level list is unset, the same way the sibling parser reads it,
|
||||
so the team's destination applies."""
|
||||
auth = UserAPIKeyAuth(
|
||||
metadata={"logging": []},
|
||||
team_metadata={"logging": [self._entry("http://team.local")]},
|
||||
)
|
||||
|
||||
assert resolve_tenant_otel_destinations(auth) == ()
|
||||
destinations = resolve_tenant_otel_destinations(auth)
|
||||
|
||||
assert [d.endpoint for d in destinations] == ["http://team.local/api/public/otel"]
|
||||
|
||||
def test_two_entries_for_one_backend_merge_their_vars_last_wins(self):
|
||||
auth = UserAPIKeyAuth(
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ import inspect
|
|||
|
||||
from litellm.proxy._types import (
|
||||
GenerateKeyRequest,
|
||||
KeyHealthResponse,
|
||||
LoggingCallbackStatus,
|
||||
NewUserRequest,
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
|
|
@ -18186,3 +18188,92 @@ async def test_key_creator_cannot_detach_project_without_admin_access():
|
|||
)
|
||||
assert exc.value.status_code == 403
|
||||
assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)
|
||||
|
||||
|
||||
def _default_team_gcs_proxy_config(team_id: str):
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
pc: Final = ProxyConfig()
|
||||
pc.config = {
|
||||
"litellm_settings": {
|
||||
"default_team_settings": [
|
||||
{"team_id": team_id, "success_callback": ["gcs_bucket"], "failure_callback": ["gcs_bucket"]}
|
||||
]
|
||||
}
|
||||
}
|
||||
return pc
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_health_tests_the_team_callbacks_an_empty_key_logging_list_falls_back_to():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import key_health
|
||||
|
||||
caller: Final = UserAPIKeyAuth(api_key="sk-1", team_id="team-gcs", metadata={"logging": []}, team_metadata={})
|
||||
logging_status: Final = LoggingCallbackStatus(callbacks=("gcs_bucket",), status="unhealthy", details="404")
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: key_health reads the module-level proxy config
|
||||
patch( # test-quality-ok: the mock completion behind test_key_logging needs a running proxy
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.test_key_logging",
|
||||
AsyncMock(return_value=logging_status),
|
||||
) as test_logging,
|
||||
):
|
||||
response = await key_health(request=MagicMock(), user_api_key_dict=caller)
|
||||
|
||||
assert response == KeyHealthResponse(key="unhealthy", logging_callbacks=logging_status)
|
||||
assert test_logging.await_args.kwargs["logging_callbacks"] == ("gcs_bucket",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_health_without_any_effective_callbacks_reports_healthy_and_sends_no_test_log():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import key_health
|
||||
|
||||
caller: Final = UserAPIKeyAuth(api_key="sk-1", team_id="team-plain", metadata={"logging": []}, team_metadata={})
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: key_health reads the module-level proxy config
|
||||
patch( # test-quality-ok: the mock completion behind test_key_logging needs a running proxy
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.test_key_logging", AsyncMock()
|
||||
) as test_logging,
|
||||
):
|
||||
response = await key_health(request=MagicMock(), user_api_key_dict=caller)
|
||||
|
||||
assert response == KeyHealthResponse(key="healthy", logging_callbacks=None)
|
||||
test_logging.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_health_rejects_key_logging_entries_without_a_callback_name():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import key_health
|
||||
|
||||
caller: Final = UserAPIKeyAuth(api_key="sk-1", metadata={"logging": [{"callback_type": "success"}]})
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")): # test-quality-ok: key_health reads the module-level proxy config
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await key_health(request=MagicMock(), user_api_key_dict=caller)
|
||||
|
||||
assert "callback_name is required" in exc.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_gcs_reports_the_failed_upload_count_from_the_registered_logger():
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures
|
||||
from litellm.types.integrations.gcs_bucket import GCSFlushResult
|
||||
|
||||
class _StuckUploadGCSLogger(GCSBucketLogger):
|
||||
def __init__(self) -> None:
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: GCS logging is premium-gated
|
||||
super().__init__(bucket_name="test-bucket")
|
||||
|
||||
async def flush_queue_and_report(self) -> GCSFlushResult:
|
||||
return GCSFlushResult(sent=0, failed=3)
|
||||
|
||||
assert await flush_gcs_and_describe_failures(_StuckUploadGCSLogger()) == "GCS upload failed for 3 event(s), 0 uploaded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_gcs_names_a_missing_logger_when_the_callback_never_initialized():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures
|
||||
|
||||
assert (
|
||||
await flush_gcs_and_describe_failures(None)
|
||||
== "gcs_bucket callback was selected but no GCS logger was initialized"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2170,6 +2170,69 @@ def test_key_dynamic_logging_settings():
|
|||
assert result is None
|
||||
|
||||
|
||||
def test_empty_logging_list_on_key_and_team_is_unset():
|
||||
"""A UI-generated `logging: []` is the same as no logging metadata, not an explicit override"""
|
||||
auth = UserAPIKeyAuth(api_key="test-key", metadata={"logging": []}, team_metadata={"logging": []})
|
||||
|
||||
assert KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(auth) is None
|
||||
assert KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(auth) is None
|
||||
|
||||
|
||||
def test_empty_key_logging_falls_back_to_team_logging():
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_id="team-1",
|
||||
metadata={"logging": []},
|
||||
team_metadata={
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "gcs_bucket",
|
||||
"callback_type": "success_and_failure",
|
||||
"callback_vars": {"gcs_bucket_name": "team-bucket"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
result = _get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=MagicMock())
|
||||
|
||||
assert result is not None
|
||||
assert result.success_callback == ["gcs_bucket"]
|
||||
assert result.failure_callback == ["gcs_bucket"]
|
||||
assert result.callback_vars == {"gcs_bucket_name": "team-bucket"}
|
||||
|
||||
|
||||
def test_empty_key_and_team_logging_falls_back_to_default_team_settings():
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
pc = ProxyConfig()
|
||||
pc.config = {
|
||||
"litellm_settings": {
|
||||
"default_team_settings": [
|
||||
{
|
||||
"team_id": "team-gcs",
|
||||
"success_callback": ["gcs_bucket"],
|
||||
"failure_callback": ["gcs_bucket"],
|
||||
"turn_off_message_logging": True,
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_id="team-gcs",
|
||||
metadata={"logging": []},
|
||||
team_metadata={"logging": []},
|
||||
)
|
||||
|
||||
result = _get_dynamic_logging_metadata(user_api_key_dict=auth, proxy_config=pc)
|
||||
|
||||
assert result is not None
|
||||
assert result.success_callback == ["gcs_bucket"]
|
||||
assert result.failure_callback == ["gcs_bucket"]
|
||||
assert result.callback_vars == {"turn_off_message_logging": "True"}
|
||||
|
||||
|
||||
def test_team_dynamic_logging_settings():
|
||||
"""
|
||||
Test KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings method with arize and langfuse callbacks
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue