scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

[Feat] Add Arize Team Based Logging (#12264)

Browse source 
* working KeyAndTeamLoggingSettings

* add arize_space_id to StandardCallbackDynamicParams

* add construct_dynamic_arize_headers

* add construct_dynamic_arize_headers to

* update otel arize logging

* fix imports

* test_construct_dynamic_arize_headers

* dynamic key/team settings

* test_get_dynamic_logging_metadata_with_arize_team_logging
This commit is contained in:
Ishaan Jaff 2025-07-02 17:15:35 -07:00 • committed by GitHub
parent 5a852ca647
commit eef0623443
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 233 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/integrations/arize/arize.py
View file

@ -12,6 +12,7 @@ from litellm.integrations.arize import _utils
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.types.integrations.arize import ArizeConfig
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import StandardCallbackDynamicParams
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -102,3 +103,30 @@ class ArizeLogger(OpenTelemetry):
):
"""Arize is used mainly for LLM I/O tracing, sending Proxy Server Request adds bloat to arize logs"""
pass
@staticmethod
def construct_dynamic_arize_headers(
standard_callback_dynamic_params: StandardCallbackDynamicParams
):
"""
Construct dynamic Arize headers from standard callback dynamic params
Returns:
dict: A dictionary of dynamic Arize headers
"""
dynamic_headers = {}
if standard_callback_dynamic_params.get("arize_space_key"):
dynamic_headers["space_key"] = standard_callback_dynamic_params.get(
"arize_space_key"
)
if standard_callback_dynamic_params.get("arize_api_key"):
dynamic_headers["api_key"] = standard_callback_dynamic_params.get(
"arize_api_key"
)
if standard_callback_dynamic_params.get("arize_space_id"):
dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get(
"arize_space_id"
)
return dynamic_headers

15
litellm/integrations/opentelemetry.py
View file

@ -429,6 +429,8 @@ class OpenTelemetry(CustomLogger):
"""
from opentelemetry import trace
from litellm.integrations.arize.arize import ArizeLogger
standard_callback_dynamic_params: Optional[
StandardCallbackDynamicParams
] = kwargs.get("standard_callback_dynamic_params")
@ -439,14 +441,7 @@ class OpenTelemetry(CustomLogger):
dynamic_headers = {}
# Handle Arize headers
if standard_callback_dynamic_params.get("arize_space_key"):
dynamic_headers["space_key"] = standard_callback_dynamic_params.get(
"arize_space_key"
)
if standard_callback_dynamic_params.get("arize_api_key"):
dynamic_headers["api_key"] = standard_callback_dynamic_params.get(
"arize_api_key"
)
dynamic_headers = ArizeLogger.construct_dynamic_arize_headers(standard_callback_dynamic_params=standard_callback_dynamic_params)
# Only create a span processor if we have headers to use
if len(dynamic_headers) > 0:
@ -579,7 +574,9 @@ class OpenTelemetry(CustomLogger):
)
return
elif self.callback_name == "langfuse_otel":
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
from litellm.integrations.langfuse.langfuse_otel import (
LangfuseOtelLogger,
)
LangfuseOtelLogger.set_langfuse_otel_attributes(
span, kwargs, response_obj

42
litellm/proxy/litellm_pre_call_utils.py
View file

@ -141,20 +141,49 @@ def convert_key_logging_metadata_to_callback(
return team_callback_settings_obj
class KeyAndTeamLoggingSettings:
"""
Helper class to get the dynamic logging settings for the key and team
"""
@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 user_api_key_dict.metadata["logging"]
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 user_api_key_dict.team_metadata["logging"]
return None
def _get_dynamic_logging_metadata(
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
) -> Optional[TeamCallbackMetadata]:
callback_settings_obj: Optional[TeamCallbackMetadata] = None
if (
user_api_key_dict.metadata is not None
and "logging" in user_api_key_dict.metadata
):
for item in user_api_key_dict.metadata["logging"]:
key_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
team_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
#########################################################################################
# Key-based callbacks
#########################################################################################
if key_dynamic_logging_settings is not None:
for item in key_dynamic_logging_settings:
callback_settings_obj = convert_key_logging_metadata_to_callback(
data=AddTeamCallback(**item),
team_callback_settings_obj=callback_settings_obj,
)
#########################################################################################
# Team-based callbacks
#########################################################################################
elif team_dynamic_logging_settings is not None:
for item in team_dynamic_logging_settings:
callback_settings_obj = convert_key_logging_metadata_to_callback(
data=AddTeamCallback(**item),
team_callback_settings_obj=callback_settings_obj,
)
#########################################################################################
# Deprecated format - maintained for backwards compatibility
#########################################################################################
elif (
user_api_key_dict.team_metadata is not None
and "callback_settings" in user_api_key_dict.team_metadata
@ -174,6 +203,9 @@ def _get_dynamic_logging_metadata(
verbose_proxy_logger.debug(
"Team callback settings activated: %s", callback_settings_obj
)
#########################################################################################
# Enter here when configured on the config.yaml file.
#########################################################################################
elif user_api_key_dict.team_id is not None:
callback_settings_obj = (
LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config(

39
tests/test_litellm/integrations/arize/test_arize_utils.py
View file

@ -232,3 +232,42 @@ async def test_arize_dynamic_params():
test_arize_logger.standard_callback_dynamic_params.get("arize_space_key")
== "test_space_key_dynamic"
)
def test_construct_dynamic_arize_headers():
"""
Test the construct_dynamic_arize_headers method with various input scenarios.
Ensures that dynamic Arize headers are properly constructed from callback parameters.
"""
# Test with all parameters present
dynamic_params_full = {
"arize_space_key": "test_space_key",
"arize_api_key": "test_api_key",
"arize_space_id": "test_space_id"
}
headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_full)
expected_headers = {
"space_key": "test_space_key",
"api_key": "test_api_key",
"arize-space-id": "test_space_id"
}
assert headers == expected_headers
# Test with only space_id
dynamic_params_space_id_only = {
"arize_space_id": "test_space_id"
}
headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_space_id_only)
expected_headers = {
"arize-space-id": "test_space_id"
}
assert headers == expected_headers
# Test with empty parameters dict
dynamic_params_empty = {}
headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_empty)
assert headers == {}

124
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -8,9 +8,11 @@ from unittest.mock import MagicMock, patch
import pytest
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import (
KeyAndTeamLoggingSettings,
LiteLLMProxyRequestSetup,
_get_dynamic_logging_metadata,
_get_enforced_params,
add_litellm_data_to_request,
check_if_token_is_service_account,
@ -216,3 +218,123 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart():
"jobID:214590dsff09fds",
"taskName:run_page_classification",
]
def test_key_dynamic_logging_settings():
"""
Test KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings method with arize and langfuse callbacks
"""
# Test with arize logging
key_with_arize = UserAPIKeyAuth(
api_key="test-key",
metadata={
"logging": [
{"callback_name": "arize", "callback_type": "success"}
]
},
team_metadata={}
)
result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize)
assert result == [{"callback_name": "arize", "callback_type": "success"}]
# Test with langfuse logging
key_with_langfuse = UserAPIKeyAuth(
api_key="test-key",
metadata={
"logging": [
{"callback_name": "langfuse", "callback_type": "success"}
]
},
team_metadata={}
)
result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse)
assert result == [{"callback_name": "langfuse", "callback_type": "success"}]
# Test with no logging metadata
key_without_logging = UserAPIKeyAuth(
api_key="test-key",
metadata={},
team_metadata={}
)
result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_without_logging)
assert result is None
def test_team_dynamic_logging_settings():
"""
Test KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings method with arize and langfuse callbacks
"""
# Test with arize team logging
key_with_team_arize = UserAPIKeyAuth(
api_key="test-key",
metadata={},
team_metadata={
"logging": [
{"callback_name": "arize", "callback_type": "failure"}
]
}
)
result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_arize)
assert result == [{"callback_name": "arize", "callback_type": "failure"}]
# Test with langfuse team logging
key_with_team_langfuse = UserAPIKeyAuth(
api_key="test-key",
metadata={},
team_metadata={
"logging": [
{"callback_name": "langfuse", "callback_type": "success"}
]
}
)
result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_langfuse)
assert result == [{"callback_name": "langfuse", "callback_type": "success"}]
# Test with no team logging metadata
key_without_team_logging = UserAPIKeyAuth(
api_key="test-key",
metadata={},
team_metadata={}
)
result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_without_team_logging)
assert result is None
def test_get_dynamic_logging_metadata_with_arize_team_logging():
"""
Test _get_dynamic_logging_metadata function with arize team logging and dynamic parameters
"""
# Setup user with arize team logging including callback_vars
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
metadata={},
team_metadata={
"logging": [
{
"callback_name": "arize",
"callback_type": "success",
"callback_vars": {
"arize_api_key": "test_arize_api_key",
"arize_space_id": "test_arize_space_id"
}
}
]
}
)
# Mock proxy_config (not used in this test path since we have team dynamic logging)
mock_proxy_config = MagicMock()
# Call the function
result = _get_dynamic_logging_metadata(
user_api_key_dict=user_api_key_dict,
proxy_config=mock_proxy_config
)
# Verify the result
assert result is not None
assert isinstance(result, TeamCallbackMetadata)
assert result.success_callback == ["arize"]
assert result.callback_vars is not None
assert result.callback_vars["arize_api_key"] == "test_arize_api_key"
assert result.callback_vars["arize_space_id"] == "test_arize_space_id"