From 0c4aae034758a6fe928ccbcaafdef3036bf45757 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 15 Oct 2025 22:07:16 -0400 Subject: [PATCH] fix: add arize from ui --- litellm/integrations/arize/arize.py | 34 ++- .../health_endpoints/_health_endpoints.py | 33 +-- .../arize/test_arize_health_check.py | 183 +++++++++++++ .../src/components/callback_info_helpers.tsx | 252 +++++++++--------- .../src/components/settings.tsx | 180 ++++++++----- 5 files changed, 482 insertions(+), 200 deletions(-) create mode 100644 tests/test_litellm/integrations/arize/test_arize_health_check.py diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 1d78e4cc69c..06e05f1271d 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -103,7 +103,39 @@ class ArizeLogger(OpenTelemetry): ): """Arize is used mainly for LLM I/O tracing, sending Proxy Server Request adds bloat to arize logs""" pass - + + async def async_health_check(self): + """ + Performs a health check for Arize integration. + + Returns: + dict: Health check result with status and message + """ + try: + config = self.get_arize_config() + + if not config.space_key: + return { + "status": "unhealthy", + "error_message": "ARIZE_SPACE_KEY environment variable not set" + } + + if not config.api_key: + return { + "status": "unhealthy", + "error_message": "ARIZE_API_KEY environment variable not set" + } + + return { + "status": "healthy", + "message": "Arize credentials are configured properly" + } + + except Exception as e: + return { + "status": "unhealthy", + "error_message": f"Arize health check failed: {str(e)}" + } def construct_dynamic_otel_headers( self, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 883bff3185f..35a80e8b7fa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -34,7 +34,7 @@ from litellm.proxy.health_check import ( #### Health ENDPOINTS #### router = APIRouter() - +services = Union[Literal["slack_budget_alerts", "langfuse", "slack", "openmeter", "webhook", "email", "braintrust", "datadog", "generic_api", "arize"], str] @router.get( "/test", @@ -64,20 +64,7 @@ async def test_endpoint(request: Request): ) async def health_services_endpoint( # noqa: PLR0915 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - service: Union[ - Literal[ - "slack_budget_alerts", - "langfuse", - "slack", - "openmeter", - "webhook", - "email", - "braintrust", - "datadog", - "generic_api", - ], - str, - ] = fastapi.Query(description="Specify the service being hit."), + service: services = fastapi.Query(description="Specify the service being hit."), ): """ Use this admin-only endpoint to check if the service is healthy. @@ -113,11 +100,12 @@ async def health_services_endpoint( # noqa: PLR0915 "langsmith", "datadog", "generic_api", + "arize", ]: raise HTTPException( status_code=400, detail={ - "error": f"Service must be in list. Service={service}. List={['slack_budget_alerts']}" + "error": f"Service must be in list. Service={service} not in {services}" }, ) @@ -150,6 +138,19 @@ async def health_services_endpoint( # noqa: PLR0915 else "Datadog is healthy" ), } + elif service == "arize": + from litellm.integrations.arize.arize import ArizeLogger + + arize_logger = ArizeLogger() + response = await arize_logger.async_health_check() + return { + "status": response["status"], + "message": ( + response["error_message"] + if response["status"] == "unhealthy" + else "Arize is healthy" + ), + } elif service == "langfuse": from litellm.integrations.langfuse.langfuse import LangFuseLogger diff --git a/tests/test_litellm/integrations/arize/test_arize_health_check.py b/tests/test_litellm/integrations/arize/test_arize_health_check.py new file mode 100644 index 00000000000..91d0b42d48d --- /dev/null +++ b/tests/test_litellm/integrations/arize/test_arize_health_check.py @@ -0,0 +1,183 @@ +""" +Test Arize health check functionality and proxy integration. +""" +import json +import os +import sys +from unittest.mock import patch, MagicMock + +# Adds the grandparent directory to sys.path to allow importing project modules +sys.path.insert(0, os.path.abspath("../..")) + +import asyncio +import pytest + +import litellm +from litellm.integrations.arize.arize import ArizeLogger +from litellm.types.utils import StandardCallbackDynamicParams + + +class TestArizeHealthCheck: + """Test Arize health check functionality.""" + + @pytest.mark.asyncio + async def test_arize_health_check_with_credentials(self): + """Test Arize health check returns healthy when credentials are available.""" + + with patch.dict(os.environ, { + "ARIZE_SPACE_KEY": "test-space-key", + "ARIZE_API_KEY": "test-api-key", + "ARIZE_ENDPOINT": "https://otlp.arize.com/v1" + }): + arize_logger = ArizeLogger() + response = await arize_logger.async_health_check() + + assert response["status"] == "healthy" + assert "configured properly" in response["message"] + + @pytest.mark.asyncio + async def test_arize_health_check_missing_space_key(self): + """Test Arize health check returns unhealthy when space key is missing.""" + + with patch.dict(os.environ, { + "ARIZE_API_KEY": "test-api-key" + }, clear=True): + arize_logger = ArizeLogger() + response = await arize_logger.async_health_check() + + assert response["status"] == "unhealthy" + assert "ARIZE_SPACE_KEY" in response["error_message"] + + @pytest.mark.asyncio + async def test_arize_health_check_missing_api_key(self): + """Test Arize health check returns unhealthy when API key is missing.""" + + with patch.dict(os.environ, { + "ARIZE_SPACE_KEY": "test-space-key" + }, clear=True): + arize_logger = ArizeLogger() + response = await arize_logger.async_health_check() + + assert response["status"] == "unhealthy" + assert "ARIZE_API_KEY" in response["error_message"] + + @pytest.mark.asyncio + async def test_arize_health_check_missing_both_keys(self): + """Test Arize health check when both keys are missing.""" + + with patch.dict(os.environ, {}, clear=True): + arize_logger = ArizeLogger() + response = await arize_logger.async_health_check() + + assert response["status"] == "unhealthy" + assert "ARIZE_SPACE_KEY" in response["error_message"] + + +class TestArizeIntegrationWithProxy: + """Test Arize integration with LiteLLM completion requests.""" + + @pytest.mark.asyncio + async def test_arize_logging_with_completion(self): + """Test that Arize logging works with actual completion requests.""" + + with patch.dict(os.environ, { + "ARIZE_SPACE_KEY": "test-space-key", + "ARIZE_API_KEY": "test-api-key", + "ARIZE_ENDPOINT": "https://otlp.arize.com/v1" + }): + # Create ArizeLogger instance + arize_logger = ArizeLogger() + + # Store original callbacks + original_callbacks = litellm.success_callback.copy() if litellm.success_callback else [] + + try: + # Add ArizeLogger to callbacks + litellm.success_callback = [arize_logger] + + # Make completion request + response = await litellm.acompletion( + model="openai/litellm-mock-response-model", + messages=[{"role": "user", "content": "Test message for Arize health check"}], + mock_response="This is a test response that validates Arize integration.", + user="test-arize-health" + ) + + # Verify response is valid + assert response is not None + print(f"Response type: {type(response)}") + print("✅ Arize completion request completed successfully") + + # Give time for async logging + await asyncio.sleep(0.1) + + print("✅ Arize completion logging test successful") + + finally: + # Restore original callbacks + litellm.success_callback = original_callbacks + + def test_arize_get_config(self): + """Test ArizeLogger.get_arize_config() method.""" + + with patch.dict(os.environ, { + "ARIZE_SPACE_KEY": "test-space-123", + "ARIZE_API_KEY": "test-api-456", + "ARIZE_ENDPOINT": "https://custom.arize.com/v1" + }): + config = ArizeLogger.get_arize_config() + + assert config.space_key == "test-space-123" + assert config.api_key == "test-api-456" + assert config.endpoint == "https://custom.arize.com/v1" + assert config.protocol == "otlp_grpc" + + def test_arize_get_config_defaults(self): + """Test ArizeLogger.get_arize_config() with default endpoint.""" + + with patch.dict(os.environ, { + "ARIZE_SPACE_KEY": "test-space-default", + "ARIZE_API_KEY": "test-api-default" + }, clear=True): + config = ArizeLogger.get_arize_config() + + assert config.space_key == "test-space-default" + assert config.api_key == "test-api-default" + assert config.endpoint == "https://otlp.arize.com/v1" # Default endpoint + assert config.protocol == "otlp_grpc" # Default protocol + + def test_arize_construct_dynamic_headers(self): + """Test dynamic OTEL headers construction for team/key logging.""" + + arize_logger = ArizeLogger() + + dynamic_params = StandardCallbackDynamicParams( + arize_space_key="dynamic-space-123", + arize_api_key="dynamic-api-456" + ) + + headers = arize_logger.construct_dynamic_otel_headers(dynamic_params) + + assert headers is not None + assert headers["arize-space-id"] == "dynamic-space-123" + assert headers["api_key"] == "dynamic-api-456" + + def test_arize_construct_dynamic_headers_space_id_fallback(self): + """Test dynamic headers with arize_space_id parameter (fallback).""" + + arize_logger = ArizeLogger() + + dynamic_params = StandardCallbackDynamicParams( + arize_space_id="fallback-space-789", # Using space_id instead of space_key + arize_api_key="fallback-api-999" + ) + + headers = arize_logger.construct_dynamic_otel_headers(dynamic_params) + + assert headers is not None + assert headers["arize-space-id"] == "fallback-space-789" + assert headers["api_key"] == "fallback-api-999" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index 66ef4d87385..5c13d2a095a 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -1,137 +1,151 @@ -export enum Callbacks { - Braintrust = "Braintrust", - CustomCallbackAPI = "Custom Callback API", - Datadog = "Datadog", - Langfuse = "Langfuse", - LangfuseOtel = "LangfuseOtel", - LangSmith = "LangSmith", - Lago = "Lago", - OpenMeter = "OpenMeter", - OTel = "Open Telemetry", - S3 = "S3", - Arize = "Arize", -} - -export const callback_map: Record = { - Braintrust: "braintrust", - CustomCallbackAPI: "custom_callback_api", - Datadog: "datadog", - Langfuse: "langfuse", - LangfuseOtel: "langfuse_otel", - LangSmith: "langsmith", - Lago: "lago", - OpenMeter: "openmeter", - OTel: "otel", - S3: "s3", - Arize: "arize", -} - -// Reverse mapping from internal values to display names -export const reverse_callback_map: Record = Object.fromEntries( - Object.entries(callback_map).map(([key, value]) => [value, key]) -); - -// Utility function to convert internal callback values to display names -export const mapInternalToDisplayNames = (internalValues: string[]): string[] => { - return internalValues.map(value => reverse_callback_map[value] || value); -}; - -// Utility function to convert display names to internal callback values -export const mapDisplayToInternalNames = (displayValues: string[]): string[] => { - return displayValues.map(value => callback_map[value] || value); -}; - const asset_logos_folder = '/ui/assets/logos/'; -interface CallbackInfo { - logo: string; +interface CallbackConfig { + id: string; // Internal callback name (e.g., "arize", "custom_callback_api") + displayName: string; // User-facing name (e.g., "Arize", "Custom Callback API") + logo: string; // Logo path supports_key_team_logging: boolean; dynamic_params: Record; - description: string | null; + description: string; } -export const callbackInfo: Record = { - [Callbacks.Langfuse]: { +// Single source of truth for ALL callback configurations +export const CALLBACK_CONFIGS: CallbackConfig[] = [ + { + id: "arize", + displayName: "Arize", + logo: `${asset_logos_folder}arize.png`, + supports_key_team_logging: true, + dynamic_params: { + "arize_api_key": "password", + "arize_space_key": "password", + }, + description: "Arize Logging Integration" + }, + { + id: "braintrust", + displayName: "Braintrust", + logo: `${asset_logos_folder}braintrust.png`, + supports_key_team_logging: false, + dynamic_params: { + "braintrust_api_key": "password", + "braintrust_project_name": "text" + }, + description: "Braintrust Logging Integration" + }, + { + id: "custom_callback_api", + displayName: "Custom Callback API", + logo: `${asset_logos_folder}custom.svg`, + supports_key_team_logging: true, + dynamic_params: { + "custom_callback_api_url": "text", + "custom_callback_api_headers": "text" + }, + description: "Custom Callback API Logging Integration" + }, + { + id: "datadog", + displayName: "Datadog", + logo: `${asset_logos_folder}datadog.png`, + supports_key_team_logging: false, + dynamic_params: { + "dd_api_key": "password", + "dd_site": "text" + }, + description: "Datadog Logging Integration" + }, + { + id: "lago", + displayName: "Lago", + logo: `${asset_logos_folder}lago.svg`, + supports_key_team_logging: false, + dynamic_params: { + "lago_api_url": "text", + "lago_api_key": "password" + }, + description: "Lago Billing Logging Integration" + }, + { + id: "langfuse", + displayName: "Langfuse", logo: `${asset_logos_folder}langfuse.png`, supports_key_team_logging: true, dynamic_params: { - "langfuse_public_key": "text", - "langfuse_secret_key": "password", - "langfuse_host": "text" + "langfuse_public_key": "text", + "langfuse_secret_key": "password", + "langfuse_host": "text" }, description: "Langfuse v2 Logging Integration" + }, + { + id: "langfuse_otel", + displayName: "Langfuse OTEL", + logo: `${asset_logos_folder}langfuse.png`, + supports_key_team_logging: true, + dynamic_params: { + "langfuse_public_key": "text", + "langfuse_secret_key": "password", + "langfuse_host": "text" }, - [Callbacks.LangfuseOtel]: { - logo: `${asset_logos_folder}langfuse.png`, - supports_key_team_logging: true, - dynamic_params: { - "langfuse_public_key": "text", - "langfuse_secret_key": "password", - "langfuse_host": "text" - }, - description: "Langfuse v3 OTEL Logging Integration" + description: "Langfuse v3 OTEL Logging Integration" + }, + { + id: "langsmith", + displayName: "LangSmith", + logo: `${asset_logos_folder}langsmith.png`, + supports_key_team_logging: true, + dynamic_params: { + "langsmith_api_key": "password", + "langsmith_project": "text", + "langsmith_base_url": "text", + "langsmith_sampling_rate": "number" }, - [Callbacks.Arize]: { - logo: `${asset_logos_folder}arize.png`, - supports_key_team_logging: true, - dynamic_params: { - "arize_api_key": "password", - "arize_space_id": "text", - }, - description: "Arize Logging Integration" + description: "Langsmith Logging Integration" + }, + { + id: "openmeter", + displayName: "OpenMeter", + logo: `${asset_logos_folder}openmeter.png`, + supports_key_team_logging: false, + dynamic_params: { + "openmeter_api_key": "password", + "openmeter_base_url": "text" }, - [Callbacks.LangSmith]: { - logo: `${asset_logos_folder}langsmith.png`, - supports_key_team_logging: true, - dynamic_params: { - "langsmith_api_key": "password", - "langsmith_project": "text", - "langsmith_base_url": "text", - "langsmith_sampling_rate": "number" - }, - description: "Langsmith Logging Integration" + description: "OpenMeter Logging Integration" + }, + { + id: "otel", + displayName: "Open Telemetry", + logo: `${asset_logos_folder}otel.png`, + supports_key_team_logging: false, + dynamic_params: { + "otel_endpoint": "text", + "otel_headers": "text" }, - [Callbacks.Braintrust]: { - logo: `${asset_logos_folder}braintrust.png`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "Braintrust Logging Integration" + description: "OpenTelemetry Logging Integration" + }, + { + id: "s3", + displayName: "S3", + logo: `${asset_logos_folder}aws.svg`, + supports_key_team_logging: false, + dynamic_params: { + "s3_bucket_name": "text", + "aws_access_key_id": "password", + "aws_secret_access_key": "password", + "aws_region": "text" }, - [Callbacks.CustomCallbackAPI]: { - logo: `${asset_logos_folder}custom.svg`, - supports_key_team_logging: true, - dynamic_params: {}, - description: "Custom Callback API Logging Integration" - }, - [Callbacks.Datadog]: { - logo: `${asset_logos_folder}datadog.png`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "Datadog Logging Integration" - }, - [Callbacks.Lago]: { - logo: `${asset_logos_folder}lago.svg`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "Lago Billing Logging Integration" - }, - [Callbacks.OpenMeter]: { - logo: `${asset_logos_folder}openmeter.png`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "OpenMeter Logging Integration" - }, - [Callbacks.OTel]: { - logo: `${asset_logos_folder}otel.png`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "OpenTelemetry Logging Integration" - }, - [Callbacks.S3]: { - logo: `${asset_logos_folder}aws.svg`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "S3 Bucket (AWS) Logging Integration" - } + description: "S3 Bucket (AWS) Logging Integration" + } +]; + +// Utility functions for easy access +export const getCallbackById = (id: string): CallbackConfig | undefined => { + return CALLBACK_CONFIGS.find(callback => callback.id === id); +}; + +export const getCallbackByDisplayName = (displayName: string): CallbackConfig | undefined => { + return CALLBACK_CONFIGS.find(callback => callback.displayName === displayName); }; \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 55997ce95cd..752bce3b5a8 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -45,9 +45,8 @@ import { import AlertingSettings from "./alerting/alerting_settings"; import FormItem from "antd/es/form/FormItem"; import { - callback_map, - callbackInfo, - Callbacks, + CALLBACK_CONFIGS, + getCallbackById, } from "./callback_info_helpers"; import { parseErrorMessage } from "./shared/errorUtils"; interface SettingsPageProps { @@ -239,18 +238,21 @@ const Settings: React.FC = ({ } }; - const handleSelectedCallbackChange = ( - callbackObject: genericCallbackParams - ) => { - setSelectedCallback(callbackObject.litellm_callback_name); - - if (callbackObject && callbackObject.litellm_callback_params) { - setSelectedCallbackParams(callbackObject.litellm_callback_params); + const handleSelectedCallbackChange = (callbackName: string) => { + setSelectedCallback(callbackName); + + // Get the callback configuration using the new clean structure + const callbackConfig = getCallbackById(callbackName); + + // Get the parameters from the callback configuration + if (callbackConfig?.dynamic_params) { + const params = Object.keys(callbackConfig.dynamic_params); + setSelectedCallbackParams(params); } else { setSelectedCallbackParams([]); } }; - + const handleSaveAlerts = async () => { if (!accessToken) { return; @@ -639,74 +641,124 @@ const Settings: React.FC = ({ wrapperCol={{ span: 16 }} labelAlign="left" > - <> - {selectedCallbackParams && - selectedCallbackParams.map((param) => ( - - - - ))} + {selectedCallbackParams && selectedCallbackParams.length > 0 && ( +
+ {selectedCallbackParams.map((param) => { + // Get the callback configuration to look up parameter types + const callbackConfig = getCallbackById(selectedCallback || ''); + const paramType = callbackConfig?.dynamic_params[param] || "text"; + + const fieldLabel = param.replace(/_/g, " ").replace(/\b\w/g, l => l.toUpperCase()); + + return ( + + {fieldLabel} + * + + } + name={param} + key={param} + className="mb-4" + rules={[ + { + required: true, + message: `Please enter the ${fieldLabel.toLowerCase()}`, + }, + ]} + > + {paramType === "password" ? ( + + ) : paramType === "number" ? ( + + ) : ( + + )} + + ); + })} +
+ )} -
- Save +
+ + + Add Callback +
-