mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge pull request #16335 from BerriAI/litellm_ui_callback_fix
[Feature] Show all callbacks on UI
This commit is contained in:
yuneng-jiang
GitHub
commit
62045477ba
7 changed files with 313 additions and 60 deletions
|
|
@ -80,6 +80,44 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
self.turn_off_message_logging = turn_off_message_logging
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def get_callback_env_vars(callback_name: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
Return the environment variables associated with a given callback
|
||||
name as defined in the proxy callback registry.
|
||||
|
||||
Args:
|
||||
callback_name: The name of the callback to look up.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of required environment variable names.
|
||||
"""
|
||||
if callback_name is None:
|
||||
return []
|
||||
|
||||
normalized_name = callback_name.lower()
|
||||
|
||||
alias_map = {
|
||||
"langfuse_otel": "langfuse",
|
||||
}
|
||||
lookup_name = alias_map.get(normalized_name, normalized_name)
|
||||
|
||||
try:
|
||||
from litellm.proxy._types import AllCallbacks
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
callbacks = AllCallbacks()
|
||||
callback_info = getattr(callbacks, lookup_name, None)
|
||||
if callback_info is None:
|
||||
return []
|
||||
|
||||
params = getattr(callback_info, "litellm_callback_params", None)
|
||||
if not params:
|
||||
return []
|
||||
|
||||
return list(params)
|
||||
|
||||
def log_pre_api_call(self, model, messages, kwargs):
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -2614,6 +2614,14 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
|
|||
ui_callback_name="Lago Billing",
|
||||
)
|
||||
|
||||
traceloop: CallbackOnUI = CallbackOnUI(
|
||||
litellm_callback_name="traceloop",
|
||||
litellm_callback_params=[
|
||||
"TRACELOOP_API_KEY",
|
||||
],
|
||||
ui_callback_name="Traceloop",
|
||||
)
|
||||
|
||||
|
||||
class SpendLogsMetadata(TypedDict):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,8 +3,12 @@ from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Optional
|
|||
import litellm
|
||||
from litellm import get_secret
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
|
||||
from litellm.proxy.types_utils.utils import get_instance_fn
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingPayload,
|
||||
|
|
@ -420,6 +424,27 @@ def get_metadata_variable_name_from_kwargs(
|
|||
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
|
||||
|
||||
def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict:
|
||||
"""Process a single callback and return its data with environment variables"""
|
||||
env_vars = CustomLogger.get_callback_env_vars(_callback)
|
||||
|
||||
env_vars_dict: dict[str, str | None] = {}
|
||||
for _var in env_vars:
|
||||
env_variable = environment_variables.get(_var, None)
|
||||
if env_variable is None:
|
||||
env_vars_dict[_var] = None
|
||||
else:
|
||||
# decode + decrypt the value
|
||||
decrypted_value = decrypt_value_helper(
|
||||
value=env_variable, key=_var
|
||||
)
|
||||
env_vars_dict[_var] = decrypted_value
|
||||
|
||||
return {
|
||||
"name": _callback,
|
||||
"variables": env_vars_dict,
|
||||
"type": callback_type
|
||||
}
|
||||
def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]:
|
||||
if callbacks is None:
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from litellm.types.utils import (
|
|||
TokenCountResponse,
|
||||
)
|
||||
from litellm.utils import load_credentials_from_list
|
||||
from litellm.proxy.common_utils.callback_utils import process_callback
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from aiohttp import ClientSession
|
||||
|
|
@ -168,6 +169,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
_get_parent_otel_span_from_kwargs,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
|
|
@ -9475,8 +9477,10 @@ async def get_config(): # noqa: PLR0915
|
|||
_general_settings = config_data.get("general_settings", {})
|
||||
environment_variables = config_data.get("environment_variables", {})
|
||||
|
||||
# check if "langfuse" in litellm_settings
|
||||
_success_callbacks = _litellm_settings.get("success_callback", [])
|
||||
_failure_callbacks = _litellm_settings.get("failure_callback", [])
|
||||
_success_and_failure_callbacks = _litellm_settings.get("callbacks", [])
|
||||
|
||||
_data_to_return = []
|
||||
"""
|
||||
[
|
||||
|
|
@ -9487,70 +9491,20 @@ async def get_config(): # noqa: PLR0915
|
|||
"LANGFUSE_SECRET_KEY": "value",
|
||||
"LANGFUSE_HOST": "value"
|
||||
},
|
||||
"type": "success"
|
||||
}
|
||||
]
|
||||
|
||||
"""
|
||||
|
||||
for _callback in _success_callbacks:
|
||||
if _callback != "langfuse":
|
||||
if _callback == "openmeter":
|
||||
env_vars = [
|
||||
"OPENMETER_API_KEY",
|
||||
]
|
||||
elif _callback == "braintrust":
|
||||
env_vars = [
|
||||
"BRAINTRUST_API_KEY",
|
||||
"BRAINTRUST_API_BASE",
|
||||
]
|
||||
elif _callback == "traceloop":
|
||||
env_vars = ["TRACELOOP_API_KEY"]
|
||||
elif _callback == "custom_callback_api":
|
||||
env_vars = ["GENERIC_LOGGER_ENDPOINT"]
|
||||
elif _callback == "otel":
|
||||
env_vars = ["OTEL_EXPORTER", "OTEL_ENDPOINT", "OTEL_HEADERS"]
|
||||
elif _callback == "langsmith":
|
||||
env_vars = [
|
||||
"LANGSMITH_API_KEY",
|
||||
"LANGSMITH_PROJECT",
|
||||
"LANGSMITH_DEFAULT_RUN_NAME",
|
||||
]
|
||||
else:
|
||||
env_vars = []
|
||||
|
||||
env_vars_dict = {}
|
||||
for _var in env_vars:
|
||||
env_variable = environment_variables.get(_var, None)
|
||||
if env_variable is None:
|
||||
env_vars_dict[_var] = None
|
||||
else:
|
||||
# decode + decrypt the value
|
||||
decrypted_value = decrypt_value_helper(
|
||||
value=env_variable, key=_var
|
||||
)
|
||||
env_vars_dict[_var] = decrypted_value
|
||||
|
||||
_data_to_return.append({"name": _callback, "variables": env_vars_dict})
|
||||
elif _callback == "langfuse":
|
||||
_langfuse_vars = [
|
||||
"LANGFUSE_PUBLIC_KEY",
|
||||
"LANGFUSE_SECRET_KEY",
|
||||
"LANGFUSE_HOST",
|
||||
]
|
||||
_langfuse_env_vars = {}
|
||||
for _var in _langfuse_vars:
|
||||
env_variable = environment_variables.get(_var, None)
|
||||
if env_variable is None:
|
||||
_langfuse_env_vars[_var] = None
|
||||
else:
|
||||
# decode + decrypt the value
|
||||
decrypted_value = decrypt_value_helper(
|
||||
value=env_variable, key=_var
|
||||
)
|
||||
_langfuse_env_vars[_var] = decrypted_value
|
||||
|
||||
_data_to_return.append(
|
||||
{"name": _callback, "variables": _langfuse_env_vars}
|
||||
)
|
||||
_data_to_return.append(process_callback(_callback, "success", environment_variables))
|
||||
|
||||
for _callback in _failure_callbacks:
|
||||
_data_to_return.append(process_callback(_callback, "failure", environment_variables))
|
||||
|
||||
for _callback in _success_and_failure_callbacks:
|
||||
_data_to_return.append(process_callback(_callback, "success_and_failure", environment_variables))
|
||||
|
||||
# Check if slack alerting is on
|
||||
_alerting = _general_settings.get("alerting", [])
|
||||
|
|
|
|||
|
|
@ -103,6 +103,24 @@ class TmpFunction:
|
|||
)
|
||||
|
||||
|
||||
def test_get_callback_env_vars():
|
||||
env_vars = CustomLogger.get_callback_env_vars("langfuse")
|
||||
assert env_vars == [
|
||||
"LANGFUSE_PUBLIC_KEY",
|
||||
"LANGFUSE_SECRET_KEY",
|
||||
"LANGFUSE_HOST",
|
||||
]
|
||||
|
||||
alias_env_vars = CustomLogger.get_callback_env_vars("langfuse_otel")
|
||||
assert alias_env_vars == env_vars
|
||||
|
||||
missing_env_vars = CustomLogger.get_callback_env_vars("does_not_exist")
|
||||
assert missing_env_vars == []
|
||||
|
||||
none_env_vars = CustomLogger.get_callback_env_vars(None)
|
||||
assert none_env_vars == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chat_openai_stream():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -2464,6 +2464,169 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content
|
|||
assert mock_logger.info.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_config_callbacks_with_all_types(client_no_auth):
|
||||
"""
|
||||
Test that /get/config/callbacks returns all three callback types:
|
||||
- success_callback with type="success"
|
||||
- failure_callback with type="failure"
|
||||
- callbacks (success_and_failure) with type="success_and_failure"
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
# Create a mock config with all three callback types
|
||||
mock_config_data = {
|
||||
"litellm_settings": {
|
||||
"success_callback": ["langfuse", "braintrust"],
|
||||
"failure_callback": ["sentry"],
|
||||
"callbacks": ["otel", "langsmith"]
|
||||
},
|
||||
"environment_variables": {
|
||||
"LANGFUSE_PUBLIC_KEY": "test-public-key",
|
||||
"LANGFUSE_SECRET_KEY": "test-secret-key",
|
||||
"LANGFUSE_HOST": "https://test.langfuse.com",
|
||||
"BRAINTRUST_API_KEY": "test-braintrust-key",
|
||||
"OTEL_EXPORTER": "otlp",
|
||||
"OTEL_ENDPOINT": "http://localhost:4317",
|
||||
"LANGSMITH_API_KEY": "test-langsmith-key",
|
||||
},
|
||||
"general_settings": {}
|
||||
}
|
||||
|
||||
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
||||
|
||||
with patch.object(
|
||||
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
|
||||
), patch(
|
||||
"litellm.proxy.common_utils.callback_utils.decrypt_value_helper",
|
||||
side_effect=lambda value, key=None: value
|
||||
):
|
||||
response = client_no_auth.get("/get/config/callbacks")
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
# Verify response structure
|
||||
assert "status" in result
|
||||
assert result["status"] == "success"
|
||||
assert "callbacks" in result
|
||||
|
||||
callbacks = result["callbacks"]
|
||||
|
||||
# Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure)
|
||||
assert len(callbacks) == 5
|
||||
|
||||
# Group callbacks by type
|
||||
success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"]
|
||||
failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"]
|
||||
success_and_failure_callbacks = [cb for cb in callbacks if cb.get("type") == "success_and_failure"]
|
||||
|
||||
# Verify all callbacks have required fields
|
||||
for callback in callbacks:
|
||||
assert "name" in callback
|
||||
assert "variables" in callback
|
||||
assert "type" in callback
|
||||
assert callback["type"] in ["success", "failure", "success_and_failure"]
|
||||
|
||||
# Verify success callbacks
|
||||
assert len(success_callbacks) == 2
|
||||
success_names = [cb["name"] for cb in success_callbacks]
|
||||
assert "langfuse" in success_names
|
||||
assert "braintrust" in success_names
|
||||
|
||||
# Verify failure callbacks
|
||||
assert len(failure_callbacks) == 1
|
||||
assert failure_callbacks[0]["name"] == "sentry"
|
||||
|
||||
# Verify success_and_failure callbacks
|
||||
assert len(success_and_failure_callbacks) == 2
|
||||
success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks]
|
||||
assert "otel" in success_and_failure_names
|
||||
assert "langsmith" in success_and_failure_names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_config_callbacks_environment_variables(client_no_auth):
|
||||
"""
|
||||
Test that /get/config/callbacks correctly includes environment variables
|
||||
for each callback type with proper decryption.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
# Create a mock config with callbacks and their env vars
|
||||
mock_config_data = {
|
||||
"litellm_settings": {
|
||||
"success_callback": ["langfuse"],
|
||||
"failure_callback": [],
|
||||
"callbacks": ["otel"]
|
||||
},
|
||||
"environment_variables": {
|
||||
"LANGFUSE_PUBLIC_KEY": "encrypted-public-key",
|
||||
"LANGFUSE_SECRET_KEY": "encrypted-secret-key",
|
||||
"LANGFUSE_HOST": "https://cloud.langfuse.com",
|
||||
"OTEL_EXPORTER": "otlp",
|
||||
"OTEL_ENDPOINT": "http://localhost:4317",
|
||||
"OTEL_HEADERS": "key=value",
|
||||
},
|
||||
"general_settings": {}
|
||||
}
|
||||
|
||||
# Mock decrypt to prepend "decrypted-" to values
|
||||
def mock_decrypt(value, key=None):
|
||||
if value and isinstance(value, str) and "encrypted" in value:
|
||||
return f"decrypted-{value}"
|
||||
return value
|
||||
|
||||
proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")
|
||||
|
||||
with patch.object(
|
||||
proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
|
||||
), patch(
|
||||
"litellm.proxy.common_utils.callback_utils.decrypt_value_helper",
|
||||
side_effect=mock_decrypt
|
||||
):
|
||||
response = client_no_auth.get("/get/config/callbacks")
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
callbacks = result["callbacks"]
|
||||
|
||||
# Find langfuse callback (success type)
|
||||
langfuse_callback = next(
|
||||
(cb for cb in callbacks if cb["name"] == "langfuse"), None
|
||||
)
|
||||
assert langfuse_callback is not None
|
||||
assert langfuse_callback["type"] == "success"
|
||||
assert "variables" in langfuse_callback
|
||||
|
||||
# Verify langfuse env vars are present and decrypted
|
||||
langfuse_vars = langfuse_callback["variables"]
|
||||
assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars
|
||||
assert langfuse_vars["LANGFUSE_PUBLIC_KEY"] == "decrypted-encrypted-public-key"
|
||||
assert "LANGFUSE_SECRET_KEY" in langfuse_vars
|
||||
assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "decrypted-encrypted-secret-key"
|
||||
assert "LANGFUSE_HOST" in langfuse_vars
|
||||
assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com"
|
||||
|
||||
# Find otel callback (success_and_failure type)
|
||||
otel_callback = next(
|
||||
(cb for cb in callbacks if cb["name"] == "otel"), None
|
||||
)
|
||||
assert otel_callback is not None
|
||||
assert otel_callback["type"] == "success_and_failure"
|
||||
assert "variables" in otel_callback
|
||||
|
||||
# Verify otel env vars are present
|
||||
otel_vars = otel_callback["variables"]
|
||||
assert "OTEL_EXPORTER" in otel_vars
|
||||
assert otel_vars["OTEL_EXPORTER"] == "otlp"
|
||||
assert "OTEL_ENDPOINT" in otel_vars
|
||||
assert otel_vars["OTEL_ENDPOINT"] == "http://localhost:4317"
|
||||
assert "OTEL_HEADERS" in otel_vars
|
||||
assert otel_vars["OTEL_HEADERS"] == "key=value"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_success_callback_normalization():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
normalize_callback_names,
|
||||
)
|
||||
|
||||
from unittest.mock import patch
|
||||
from litellm.proxy.common_utils.callback_utils import process_callback
|
||||
|
||||
|
||||
def test_get_remaining_tokens_and_requests_from_request_data():
|
||||
model_group = "openrouter/google/gemini-2.0-flash-001"
|
||||
|
|
@ -30,6 +33,50 @@ def test_get_remaining_tokens_and_requests_from_request_data():
|
|||
}
|
||||
|
||||
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars",
|
||||
return_value=["API_KEY", "MISSING_VAR"],
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.callback_utils.decrypt_value_helper",
|
||||
side_effect=lambda value, key: f"decrypted-{key}",
|
||||
)
|
||||
def test_process_callback_with_env_vars(mock_decrypt, mock_get_env_vars):
|
||||
environment_variables = {
|
||||
"API_KEY": "ENC_VALUE",
|
||||
"UNUSED": "SHOULD_BE_IGNORED",
|
||||
}
|
||||
|
||||
result = process_callback(
|
||||
_callback="my_callback",
|
||||
callback_type="input",
|
||||
environment_variables=environment_variables,
|
||||
)
|
||||
|
||||
assert result["name"] == "my_callback"
|
||||
assert result["type"] == "input"
|
||||
assert result["variables"] == {
|
||||
"API_KEY": "decrypted-API_KEY",
|
||||
"MISSING_VAR": None,
|
||||
}
|
||||
|
||||
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars",
|
||||
return_value=[],
|
||||
)
|
||||
def test_process_callback_with_no_required_env_vars(mock_get_env_vars):
|
||||
result = process_callback(
|
||||
_callback="another_callback",
|
||||
callback_type="output",
|
||||
environment_variables={"SHOULD_NOT_BE_USED": "VALUE"},
|
||||
)
|
||||
|
||||
assert result["name"] == "another_callback"
|
||||
assert result["type"] == "output"
|
||||
assert result["variables"] == {}
|
||||
|
||||
|
||||
def test_normalize_callback_names_none_returns_empty_list():
|
||||
assert normalize_callback_names(None) == []
|
||||
assert normalize_callback_names([]) == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue