mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
9b53ee878f
* perf(ci): distribute unit cases and reuse security sweep connections * revert(ci): reject first speed trial after new failures * perf(ci): avoid repeated test discovery and parsing * perf(ci): isolate parallel security workers * revert(ci): restore measured security execution isolation * perf(tests): reuse source scans and deduplicate cache cleanup * perf(ci): restore parallel security with safe log batching * test(ci): isolate fixture state and schema comparisons * test(ci): make schema and scheduler fixtures deterministic * test(ci): drain queued logging before fixture teardown * test(ci): preserve logging delivery without idle worker restarts * test(ci): keep guardrail sinks available during parallel bursts * test(ci): isolate full proxy lifecycle from gateway imports * test(ci): preserve first failure responses and sweep evidence * test(ci): restore documentation URL settings between cases * test(ci): isolate open MCP tools in private gateways * test(ci): group private MCP cases and handle exited workers * test(ci): retain proxy exits and MCP failure details * test(ci): repair merge interactions with main Key the schema shard's full clone on artifact-name, since main removed the unit-flag input. Give each owned Redis its own log file again. Hand the realtime transcription tests main added a logging object whose dispatch handlers are awaitable, so the per-test logging drain no longer waits on a MagicMock task forever * perf(ci): shard owned integration cases and close lifecycle leaks * fix(ci): qualify shard execution and harden interrupted cleanup * fix(ci): bound worker supervision and own fault traffic * perf(ci): reuse unparametrized fixtures across bounded case groups * fix(ci): preserve boot fault allowances and complete exit records * test(ci): tolerate startup retries and retain readiness details * perf(ci): balance integration shards by recorded duration * fix(ci): load migration database helpers only for workers
469 lines
18 KiB
Python
469 lines
18 KiB
Python
"""
|
|
Callback to log events to a Generic API Endpoint
|
|
|
|
- Creates a StandardLoggingPayload
|
|
- Adds to batch queue
|
|
- Flushes based on CustomBatchLogger settings
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
import traceback
|
|
from collections.abc import Mapping
|
|
from typing import Final, Literal
|
|
|
|
import httpx
|
|
|
|
import litellm
|
|
from litellm._logging import verbose_logger
|
|
from litellm._uuid import uuid
|
|
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
|
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
|
from litellm.llms.custom_httpx.http_handler import (
|
|
AsyncHTTPHandler,
|
|
get_async_httpx_client,
|
|
httpxSpecialProvider,
|
|
)
|
|
from litellm.types.utils import StandardLoggingPayload
|
|
|
|
API_EVENT_TYPES = Literal["llm_api_success", "llm_api_failure"]
|
|
LOG_FORMAT_TYPES = Literal["json_array", "ndjson", "single"]
|
|
|
|
|
|
def load_compatible_callbacks() -> dict:
|
|
"""
|
|
Load the generic_api_compatible_callbacks.json file
|
|
|
|
Returns:
|
|
Dict: Dictionary of compatible callbacks configuration
|
|
"""
|
|
try:
|
|
json_path: Final = os.path.join(os.path.dirname(__file__), "generic_api_compatible_callbacks.json")
|
|
with open(json_path, "r") as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
verbose_logger.warning("Error loading generic_api_compatible_callbacks.json: %s", e)
|
|
return {}
|
|
|
|
|
|
def is_callback_compatible(callback_name: str) -> bool:
|
|
"""
|
|
Check if a callback_name exists in the compatible callbacks list
|
|
|
|
Args:
|
|
callback_name: Name of the callback to check
|
|
|
|
Returns:
|
|
bool: True if callback_name exists in the compatible callbacks, False otherwise
|
|
"""
|
|
compatible_callbacks: Final = load_compatible_callbacks()
|
|
return callback_name in compatible_callbacks
|
|
|
|
|
|
def get_callback_config(callback_name: str) -> dict | None:
|
|
"""
|
|
Get the configuration for a specific callback
|
|
|
|
Args:
|
|
callback_name: Name of the callback to get config for
|
|
|
|
Returns:
|
|
Optional[Dict]: Configuration dict for the callback, or None if not found
|
|
"""
|
|
compatible_callbacks: Final = load_compatible_callbacks()
|
|
return compatible_callbacks.get(callback_name)
|
|
|
|
|
|
def substitute_env_variables(value: str) -> str:
|
|
"""
|
|
Replace {{environment_variables.VAR_NAME}} patterns with actual environment variable values
|
|
|
|
Args:
|
|
value: String that may contain {{environment_variables.VAR_NAME}} patterns
|
|
|
|
Returns:
|
|
str: String with environment variables substituted
|
|
"""
|
|
pattern: Final = r"\{\{environment_variables\.([A-Z_]+)\}\}"
|
|
|
|
def replace_env_var(match):
|
|
env_var_name: Final = match.group(1)
|
|
return os.getenv(env_var_name, "")
|
|
|
|
return re.sub(pattern, replace_env_var, value)
|
|
|
|
|
|
class GenericAPILogger(CustomBatchLogger):
|
|
def __init__(
|
|
self,
|
|
endpoint: str | None = None,
|
|
headers: dict | None = None,
|
|
event_types: list[API_EVENT_TYPES] | None = None,
|
|
callback_name: str | None = None,
|
|
log_format: LOG_FORMAT_TYPES | None = None,
|
|
max_retries: int = 0,
|
|
retry_delay: float = 1.0,
|
|
timeout: float | httpx.Timeout | None = None,
|
|
async_httpx_client: AsyncHTTPHandler | None = None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Initialize the GenericAPILogger
|
|
|
|
Args:
|
|
endpoint: Optional[str] = None,
|
|
headers: Optional[dict] = None,
|
|
event_types: Optional[List[API_EVENT_TYPES]] = None,
|
|
callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json
|
|
log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single"
|
|
max_retries: Number of retry attempts after the initial request fails. Defaults to 0.
|
|
retry_delay: Initial retry delay in seconds. Retries use exponential backoff.
|
|
timeout: Optional timeout to use for Generic API callback requests.
|
|
"""
|
|
#########################################################
|
|
# Check if callback_name is provided and load config
|
|
#########################################################
|
|
if callback_name:
|
|
if is_callback_compatible(callback_name):
|
|
verbose_logger.debug("Loading configuration for callback: %s", callback_name)
|
|
callback_config: Final = get_callback_config(callback_name)
|
|
|
|
# Use config from JSON if not explicitly provided
|
|
if callback_config:
|
|
if endpoint is None and "endpoint" in callback_config:
|
|
endpoint = substitute_env_variables(callback_config["endpoint"])
|
|
|
|
if "headers" in callback_config:
|
|
headers = headers or {}
|
|
for key, value in callback_config["headers"].items():
|
|
if key not in headers:
|
|
headers[key] = substitute_env_variables(value)
|
|
|
|
if event_types is None and "event_types" in callback_config:
|
|
event_types = callback_config["event_types"]
|
|
|
|
if log_format is None and "log_format" in callback_config:
|
|
log_format = callback_config["log_format"]
|
|
else:
|
|
verbose_logger.warning(
|
|
"callback_name '%s' not found in generic_api_compatible_callbacks.json", callback_name
|
|
)
|
|
|
|
#########################################################
|
|
# Init httpx client
|
|
#########################################################
|
|
self.async_httpx_client = (
|
|
async_httpx_client
|
|
if async_httpx_client is not None
|
|
else get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
|
)
|
|
endpoint = endpoint or os.getenv("GENERIC_LOGGER_ENDPOINT")
|
|
if endpoint is None:
|
|
raise ValueError(
|
|
"endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables"
|
|
)
|
|
|
|
self.headers: dict[str, str] = self._get_headers(headers)
|
|
self.endpoint: str = endpoint
|
|
self.event_types: list[API_EVENT_TYPES] | None = event_types
|
|
self.callback_name: str | None = callback_name
|
|
self.max_retries = max(0, int(max_retries or 0))
|
|
retry_delay_value: Final = 0.0 if retry_delay is None else retry_delay
|
|
self.retry_delay = max(0.0, float(retry_delay_value))
|
|
self.timeout = timeout
|
|
|
|
# Validate and store log_format
|
|
if log_format is not None and log_format not in [
|
|
"json_array",
|
|
"ndjson",
|
|
"single",
|
|
]:
|
|
raise ValueError(f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'")
|
|
self.log_format: LOG_FORMAT_TYPES = log_format or "json_array"
|
|
|
|
verbose_logger.debug(
|
|
"in init GenericAPILogger, callback_name: %s, endpoint %s, headers %s, event_types: %s, log_format: %s",
|
|
self.callback_name,
|
|
self.endpoint,
|
|
self.headers,
|
|
self.event_types,
|
|
self.log_format,
|
|
)
|
|
|
|
#########################################################
|
|
# Init variables for batch flushing logs
|
|
#########################################################
|
|
self.flush_lock = asyncio.Lock()
|
|
super().__init__(**kwargs, flush_lock=self.flush_lock)
|
|
asyncio.create_task(self.periodic_flush())
|
|
self.log_queue: list[Mapping[str, object] | StandardLoggingPayload] = []
|
|
|
|
def _get_headers(self, headers: dict | None = None):
|
|
"""
|
|
Get headers for the Generic API Logger
|
|
|
|
Returns:
|
|
Dict: Headers for the Generic API Logger
|
|
|
|
Args:
|
|
headers: Optional[dict] = None
|
|
"""
|
|
# Process headers from different sources
|
|
headers_dict: Final = {
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
# 1. First check for headers from env var
|
|
env_headers: Final = os.getenv("GENERIC_LOGGER_HEADERS")
|
|
if env_headers:
|
|
try:
|
|
# Parse headers in format "key1=value1,key2=value2" or "key1=value1"
|
|
header_items: Final = env_headers.split(",")
|
|
for item in header_items:
|
|
if "=" in item:
|
|
key, value = item.split("=", 1)
|
|
headers_dict[key.strip()] = value.strip()
|
|
except Exception as e:
|
|
verbose_logger.warning("Error parsing headers from environment variables: %s", e)
|
|
|
|
# 2. Update with litellm generic headers if available
|
|
if litellm.generic_logger_headers:
|
|
headers_dict.update(litellm.generic_logger_headers)
|
|
|
|
# 3. Override with directly provided headers if any
|
|
if headers:
|
|
headers_dict.update(headers)
|
|
|
|
return headers_dict
|
|
|
|
def _should_retry_exception(self, exception: Exception) -> bool:
|
|
if isinstance(exception, (litellm.Timeout, httpx.TransportError)):
|
|
return True
|
|
|
|
if isinstance(exception, httpx.HTTPStatusError):
|
|
return exception.response.status_code >= 500
|
|
|
|
return False
|
|
|
|
async def _sleep_before_retry(self, attempt: int) -> None:
|
|
if self.retry_delay <= 0:
|
|
return
|
|
|
|
delay: Final = self.retry_delay * (2**attempt)
|
|
await asyncio.sleep(delay)
|
|
|
|
async def _post_with_retries(self, data: str) -> httpx.Response:
|
|
total_attempts: Final = self.max_retries + 1
|
|
for attempt in range(total_attempts):
|
|
try:
|
|
return await self.async_httpx_client.post(
|
|
url=self.endpoint,
|
|
headers=self.headers,
|
|
data=data,
|
|
timeout=self.timeout,
|
|
)
|
|
except Exception as e:
|
|
is_last_attempt = attempt == self.max_retries
|
|
should_retry = self._should_retry_exception(e)
|
|
if is_last_attempt or not should_retry:
|
|
raise
|
|
|
|
verbose_logger.warning(
|
|
"Generic API Logger - retrying request to %s after error: %s (attempt %s/%s)",
|
|
self.endpoint,
|
|
str(e),
|
|
attempt + 1,
|
|
total_attempts,
|
|
)
|
|
await self._sleep_before_retry(attempt)
|
|
|
|
raise RuntimeError("Generic API Logger retry loop exited unexpectedly")
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
"""
|
|
Async Log success events to Generic API Endpoint
|
|
|
|
- Creates a StandardLoggingPayload
|
|
- Adds to batch queue
|
|
- Flushes based on CustomBatchLogger settings
|
|
|
|
Raises:
|
|
Raises a NON Blocking verbose_logger.exception if an error occurs
|
|
"""
|
|
|
|
if self.event_types is not None and "llm_api_success" not in self.event_types:
|
|
return
|
|
|
|
try:
|
|
verbose_logger.debug("Generic API Logger - Enters logging function for model %s", kwargs)
|
|
standard_logging_payload: Final = kwargs.get("standard_logging_object", None)
|
|
|
|
# Backwards compatibility with old logging payload
|
|
if litellm.generic_api_use_v1 is True:
|
|
payload: Final = self._get_v1_logging_payload(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
self.log_queue.append(payload)
|
|
else:
|
|
# New logging payload, StandardLoggingPayload
|
|
self.log_queue.append(standard_logging_payload)
|
|
|
|
if len(self.log_queue) >= self.batch_size:
|
|
await self.async_send_batch()
|
|
|
|
except Exception as e:
|
|
verbose_logger.exception("Generic API Logger Error - %s\n%s", e, traceback.format_exc())
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
"""
|
|
Async Log failure events to Generic API Endpoint
|
|
|
|
- Creates a StandardLoggingPayload
|
|
- Adds to batch queue
|
|
"""
|
|
if self.event_types is not None and "llm_api_failure" not in self.event_types:
|
|
return
|
|
|
|
try:
|
|
verbose_logger.debug("Generic API Logger - Enters logging function for model %s", kwargs)
|
|
standard_logging_payload: Final = kwargs.get("standard_logging_object", None)
|
|
|
|
if litellm.generic_api_use_v1 is True:
|
|
payload: Final = self._get_v1_logging_payload(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
self.log_queue.append(payload)
|
|
else:
|
|
self.log_queue.append(standard_logging_payload)
|
|
|
|
if len(self.log_queue) >= self.batch_size:
|
|
await self.async_send_batch()
|
|
|
|
except Exception as e:
|
|
verbose_logger.exception("Generic API Logger Error - %s\n%s", e, traceback.format_exc())
|
|
|
|
async def flush_queue(self) -> None:
|
|
await self.async_send_batch()
|
|
|
|
async def async_send_batch(self) -> None:
|
|
async with self.flush_lock:
|
|
await self._send_batch()
|
|
|
|
async def _send_batch(self) -> None:
|
|
"""
|
|
Sends the batch of messages to Generic API Endpoint
|
|
|
|
Supports three formats:
|
|
- json_array: Sends all logs as a JSON array (default)
|
|
- ndjson: Sends logs as newline-delimited JSON
|
|
- single: Sends each log as individual HTTP request in parallel
|
|
"""
|
|
batch: Final = tuple(self.log_queue)
|
|
if not batch:
|
|
return
|
|
try:
|
|
verbose_logger.debug(
|
|
"Generic API Logger - about to flush %s events in '%s' format", len(batch), self.log_format
|
|
)
|
|
|
|
if self.log_format == "single":
|
|
# Send each log as individual HTTP request in parallel
|
|
tasks: Final = []
|
|
for log_entry in batch:
|
|
task = self._post_with_retries(data=safe_dumps(log_entry))
|
|
tasks.append(task)
|
|
|
|
# Execute all requests in parallel
|
|
responses: Final = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Log results
|
|
for idx, result in enumerate(responses):
|
|
if isinstance(result, Exception):
|
|
verbose_logger.exception("Generic API Logger - Error sending log %s: %s", idx, result)
|
|
else:
|
|
# result is a Response object
|
|
verbose_logger.debug(
|
|
"Generic API Logger - sent log %s, status: %s",
|
|
idx,
|
|
result.status_code,
|
|
)
|
|
else:
|
|
# Format the payload based on log_format
|
|
if self.log_format == "json_array":
|
|
data = safe_dumps(batch)
|
|
elif self.log_format == "ndjson":
|
|
data = "\n".join(safe_dumps(log) for log in batch)
|
|
else:
|
|
raise ValueError(f"Unknown log_format: {self.log_format}")
|
|
|
|
# Make POST request
|
|
response: Final = await self._post_with_retries(data=data)
|
|
|
|
verbose_logger.debug(
|
|
"Generic API Logger - sent batch to %s, status: %s, format: %s",
|
|
self.endpoint,
|
|
response.status_code,
|
|
self.log_format,
|
|
)
|
|
|
|
except Exception as e:
|
|
verbose_logger.exception("Generic API Logger Error sending batch - %s\n%s", e, traceback.format_exc())
|
|
finally:
|
|
del self.log_queue[: len(batch)]
|
|
self.last_flush_time = time.time()
|
|
|
|
def _get_v1_logging_payload(self, kwargs, response_obj, start_time, end_time) -> dict:
|
|
"""
|
|
Maintained for backwards compatibility with old logging payload
|
|
|
|
Returns a dict of the payload to send to the Generic API Endpoint
|
|
"""
|
|
verbose_logger.debug("GenericAPILogger Logging - Enters logging function for model %s", kwargs)
|
|
|
|
# construct payload to send custom logger
|
|
# follows the same params as langfuse.py
|
|
litellm_params: Final = kwargs.get("litellm_params", {})
|
|
metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None
|
|
messages: Final = kwargs.get("messages")
|
|
cost: Final = kwargs.get("response_cost", 0.0)
|
|
optional_params: Final = kwargs.get("optional_params", {})
|
|
call_type: Final = kwargs.get("call_type", "litellm.completion")
|
|
cache_hit: Final = kwargs.get("cache_hit", False)
|
|
usage: Final = response_obj["usage"]
|
|
id: Final = response_obj.get("id", str(uuid.uuid4()))
|
|
|
|
# Build the initial payload
|
|
payload: Final = {
|
|
"id": id,
|
|
"call_type": call_type,
|
|
"cache_hit": cache_hit,
|
|
"startTime": start_time,
|
|
"endTime": end_time,
|
|
"model": kwargs.get("model", ""),
|
|
"user": kwargs.get("user", ""),
|
|
"modelParameters": optional_params,
|
|
"messages": messages,
|
|
"response": response_obj,
|
|
"usage": usage,
|
|
"metadata": metadata,
|
|
"cost": cost,
|
|
}
|
|
|
|
# Ensure everything in the payload is converted to str
|
|
for key, value in payload.items():
|
|
try:
|
|
payload[key] = str(value)
|
|
except Exception:
|
|
# non blocking if it can't cast to a str
|
|
pass
|
|
|
|
return payload
|