From 063db7ba9151d15b36c6ff27ccd9af088756a13f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 18 Mar 2025 15:51:58 -0700 Subject: [PATCH] fix arize logging --- litellm/litellm_core_utils/litellm_logging.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0945c45491..dded8232f0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -29,6 +29,7 @@ from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.cost_calculator import _select_model_name_for_cost_calc +from litellm.integrations.arize.arize import ArizeLogger from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.mlflow import MlflowLogger @@ -2658,11 +2659,11 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 ) for callback in _in_memory_loggers: if ( - isinstance(callback, OpenTelemetry) + isinstance(callback, ArizeLogger) and callback.callback_name == "arize" ): return callback # type: ignore - _otel_logger = OpenTelemetry(config=otel_config, callback_name="arize") + _otel_logger = ArizeLogger(config=otel_config, callback_name="arize") _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore elif logging_integration == "arize_phoenix": @@ -2897,15 +2898,13 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 if isinstance(callback, OpenTelemetry): return callback elif logging_integration == "arize": - from litellm.integrations.opentelemetry import OpenTelemetry - if "ARIZE_SPACE_KEY" not in os.environ: raise ValueError("ARIZE_SPACE_KEY not found in environment variables") if "ARIZE_API_KEY" not in os.environ: raise ValueError("ARIZE_API_KEY not found in environment variables") for callback in _in_memory_loggers: if ( - isinstance(callback, OpenTelemetry) + isinstance(callback, ArizeLogger) and callback.callback_name == "arize" ): return callback