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 e5d83d60195..9964b6c94cc 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 5b411894a81..d11f6a58368 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -1,136 +1,150 @@ -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", -} +const asset_logos_folder = '/ui/assets/logos/'; -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]: { - 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 v2 Logging Integration", - }, - [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", - }, - [Callbacks.Arize]: { +// 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_id: "text", + "arize_api_key": "password", + "arize_space_key": "password", }, - description: "Arize Logging Integration", + description: "Arize Logging Integration" }, - [Callbacks.LangSmith]: { + { + 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" + }, + 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" + }, + 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", + "langsmith_api_key": "password", + "langsmith_project": "text", + "langsmith_base_url": "text", + "langsmith_sampling_rate": "number" }, - description: "Langsmith Logging Integration", + description: "Langsmith Logging Integration" }, - [Callbacks.Braintrust]: { - logo: `${asset_logos_folder}braintrust.png`, - supports_key_team_logging: false, - dynamic_params: {}, - description: "Braintrust Logging Integration", - }, - [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]: { + { + id: "openmeter", + displayName: "OpenMeter", logo: `${asset_logos_folder}openmeter.png`, supports_key_team_logging: false, - dynamic_params: {}, - description: "OpenMeter Logging Integration", + dynamic_params: { + "openmeter_api_key": "password", + "openmeter_base_url": "text" + }, + description: "OpenMeter Logging Integration" }, - [Callbacks.OTel]: { + { + id: "otel", + displayName: "Open Telemetry", logo: `${asset_logos_folder}otel.png`, supports_key_team_logging: false, - dynamic_params: {}, - description: "OpenTelemetry Logging Integration", + dynamic_params: { + "otel_endpoint": "text", + "otel_headers": "text" + }, + description: "OpenTelemetry Logging Integration" }, - [Callbacks.S3]: { + { + id: "s3", + displayName: "S3", logo: `${asset_logos_folder}aws.svg`, supports_key_team_logging: false, - dynamic_params: {}, - description: "S3 Bucket (AWS) Logging Integration", - }, + dynamic_params: { + "s3_bucket_name": "text", + "aws_access_key_id": "password", + "aws_secret_access_key": "password", + "aws_region": "text" + }, + 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); }; diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 4724476c4e9..727c8b9b511 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -32,7 +32,10 @@ const { Title, Paragraph } = Typography; import { getCallbacksCall, setCallbacksCall, serviceHealthCheck, deleteCallback } from "./networking"; import AlertingSettings from "./alerting/alerting_settings"; import FormItem from "antd/es/form/FormItem"; -import { callback_map, callbackInfo, Callbacks } from "./callback_info_helpers"; +import { + CALLBACK_CONFIGS, + getCallbackById, +} from "./callback_info_helpers"; import { parseErrorMessage } from "./shared/errorUtils"; interface SettingsPageProps { accessToken: string | null; @@ -211,16 +214,21 @@ const Settings: React.FC = ({ accessToken, userRole, userID, } }; - 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; @@ -586,67 +594,124 @@ const Settings: React.FC = ({ accessToken, userRole, userID, 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 +
-