diff --git a/adalflow/adalflow/tracing/mlflow_integration.py b/adalflow/adalflow/tracing/mlflow_integration.py index f8f9851c6..b86a430ef 100644 --- a/adalflow/adalflow/tracing/mlflow_integration.py +++ b/adalflow/adalflow/tracing/mlflow_integration.py @@ -12,12 +12,21 @@ try: # Do NOT set ADALFLOW_DISABLE_TRACING here - it should be set before imports import mlflow - # BUG: This import is not working for mlflow 3.7.0 - from mlflow.openai._agent_tracer import MlflowOpenAgentTracingProcessor - MLFLOW_AVAILABLE = True + + # Try importing the agent tracer (may not exist in newer MLflow versions) + try: + from mlflow.openai._agent_tracer import MlflowOpenAgentTracingProcessor + MLFLOW_AGENT_TRACER_AVAILABLE = True + except (ImportError, AttributeError): + MLFLOW_AGENT_TRACER_AVAILABLE = False + log.debug("MLflow agent tracer not available in this MLflow version") + MlflowOpenAgentTracingProcessor = None + except ImportError: MLFLOW_AVAILABLE = False + MLFLOW_AGENT_TRACER_AVAILABLE = False + MlflowOpenAgentTracingProcessor = None log.warning("MLflow not available. Install with: pip install mlflow") @@ -69,6 +78,14 @@ def enable_mlflow_local( if not MLFLOW_AVAILABLE: log.error("MLflow is not installed. Cannot enable MLflow tracing.") return False + + if not MLFLOW_AGENT_TRACER_AVAILABLE: + log.error( + "MLflow agent tracer is not available in this MLflow version. " + "This feature requires MLflow < 3.7.0. Please downgrade MLflow or " + "use an alternative tracing method." + ) + return False try: # Set the environment variable BEFORE accessing the provider @@ -148,6 +165,14 @@ def enable_mlflow_local_with_server( if not MLFLOW_AVAILABLE: log.error("MLflow is not installed. Cannot enable MLflow tracing.") return False + + if not MLFLOW_AGENT_TRACER_AVAILABLE: + log.error( + "MLflow agent tracer is not available in this MLflow version. " + "This feature requires MLflow < 3.7.0. Please downgrade MLflow or " + "use an alternative tracing method." + ) + return False try: # Set the environment variable BEFORE accessing the provider