1  from mlflow.telemetry.events import AutologgingEvent
 2  from mlflow.telemetry.track import _record_event
 3  from mlflow.utils.autologging_utils import autologging_integration
 4  
 5  FLAVOR_NAME = "ag2"
 6  
 7  
 8  def autolog(
 9      log_traces: bool = True,
10      disable: bool = False,
11      silent: bool = False,
12  ):
13      """
14      Enables (or disables) and configures autologging from ag2 to MLflow. Currently, MLflow
15      only supports tracing for ag2 agents.
16  
17      Args:
18          log_traces: If ``True``, traces are logged for AG2 agents by using runtime logging.
19              If ``False``, no traces are collected during inference. Default to ``True``.
20          disable: If ``True``, disables the AG2 autologging. Default to ``False``.
21          silent: If ``True``, suppress all event logs and warnings from MLflow during AG2
22              autologging. If ``False``, show all events and warnings.
23      """
24      from autogen import runtime_logging
25  
26      from mlflow.ag2.ag2_logger import MlflowAg2Logger
27  
28      # NB: The @autologging_integration annotation is used for adding shared logic. However, one
29      # caveat is that the wrapped function is NOT executed when disable=True is passed. This prevents
30      # us from running cleaning up logging when autologging is turned off. To workaround this, we
31      # annotate _autolog() instead of this entrypoint, and define the cleanup logic outside it.
32      # TODO: since this implementation is inconsistent, explore a universal way to solve the issue.
33      if log_traces and not disable:
34          runtime_logging.start(logger=MlflowAg2Logger())
35      else:
36          runtime_logging.stop()
37  
38      _autolog(log_traces=log_traces, disable=disable, silent=silent)
39  
40  
41  # This is required by mlflow.autolog()
42  autolog.integration_name = FLAVOR_NAME
43  
44  
45  @autologging_integration(FLAVOR_NAME)
46  def _autolog(
47      log_traces: bool = True,
48      disable: bool = False,
49      silent: bool = False,
50  ):
51      """
52      This is a dummy function only for the purpose of adding the autologging_integration annotation.
53      We cannot add the annotation directly to the autolog() function above due to the reason
54      mentioned in the comment above. Note that this function MUST declare the same signature as the
55      autolog(), otherwise the annotation will not work properly.
56      """
57      _record_event(
58          AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
59      )