wehub-resource-sync
113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
import inspect
|
||
import logging
|
||
|
||
from mlflow.telemetry.events import AutologgingEvent
|
||
from mlflow.telemetry.track import _record_event
|
||
from mlflow.utils.annotations import experimental as experimental
|
||
from mlflow.utils.autologging_utils import autologging_integration, safe_patch
|
||
|
||
FLAVOR_NAME = "agno"
|
||
_logger = logging.getLogger(__name__)
|
||
|
||
|
||
def autolog(*, log_traces: bool = True, disable: bool = False, silent: bool = False) -> None:
|
||
"""
|
||
Enables (or disables) and configures autologging from Agno to MLflow.
|
||
|
||
For Agno V2 (>= 2.0.0), this uses OpenTelemetry instrumentation via OpenInference.
|
||
|
||
Args:
|
||
log_traces: If ``True``, traces are logged for Agno Agents.
|
||
disable: If ``True``, disables Agno autologging.
|
||
silent: If ``True``, suppresses all MLflow event logs and warnings.
|
||
"""
|
||
from mlflow.agno.autolog_v1 import patched_async_class_call, patched_class_call
|
||
from mlflow.agno.autolog_v2 import _is_agno_v2, _setup_otel_instrumentation, _uninstrument_otel
|
||
|
||
# NB: The @autologging_integration annotation is used for adding shared logic. However, one
|
||
# caveat is that the wrapped function is NOT executed when disable=True is passed. This prevents
|
||
# us from running cleaning up logging when autologging is turned off. To workaround this, we
|
||
# annotate _autolog() instead of this entrypoint, and define the cleanup logic outside it.
|
||
# This needs to be called before doing any safe-patching (otherwise safe-patch will be no-op).
|
||
_autolog(log_traces=log_traces, disable=disable, silent=silent)
|
||
|
||
# Check if Agno V2 is installed
|
||
if _is_agno_v2():
|
||
_logger.debug("Detected Agno V2, using OpenTelemetry instrumentation")
|
||
if disable or not log_traces:
|
||
_uninstrument_otel()
|
||
else:
|
||
_setup_otel_instrumentation()
|
||
_record_event(
|
||
AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
|
||
)
|
||
return
|
||
|
||
# For Agno V1, use the existing patching method
|
||
from mlflow.agno.utils import discover_storage_backends, find_model_subclasses
|
||
|
||
class_map = {
|
||
"agno.agent.Agent": ["run", "arun"],
|
||
"agno.team.Team": ["run", "arun"],
|
||
"agno.tools.function.FunctionCall": ["execute", "aexecute"],
|
||
}
|
||
|
||
if storages := discover_storage_backends():
|
||
class_map.update({
|
||
cls.__module__ + "." + cls.__name__: [
|
||
"create",
|
||
"read",
|
||
"upsert",
|
||
"drop",
|
||
"upgrade_schema",
|
||
]
|
||
for cls in storages
|
||
})
|
||
|
||
if models := find_model_subclasses():
|
||
class_map.update({
|
||
# TODO: Support streaming
|
||
cls.__module__ + "." + cls.__name__: ["invoke", "ainvoke"]
|
||
for cls in models
|
||
})
|
||
|
||
for cls_path, methods in class_map.items():
|
||
mod_name, cls_name = cls_path.rsplit(".", 1)
|
||
try:
|
||
module = __import__(mod_name, fromlist=[cls_name])
|
||
cls = getattr(module, cls_name)
|
||
except (ImportError, AttributeError) as exc:
|
||
_logger.debug("Agno autologging: failed to import %s – %s", cls_path, exc)
|
||
continue
|
||
|
||
for method_name in methods:
|
||
try:
|
||
original = getattr(cls, method_name)
|
||
wrapper = (
|
||
patched_async_class_call
|
||
if inspect.iscoroutinefunction(original)
|
||
else patched_class_call
|
||
)
|
||
safe_patch(FLAVOR_NAME, cls, method_name, wrapper)
|
||
except AttributeError as exc:
|
||
_logger.debug(
|
||
"Agno autologging: cannot patch %s.%s – %s", cls_path, method_name, exc
|
||
)
|
||
|
||
_record_event(
|
||
AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
|
||
)
|
||
|
||
|
||
# This is required by mlflow.autolog()
|
||
autolog.integration_name = FLAVOR_NAME
|
||
|
||
|
||
@autologging_integration(FLAVOR_NAME)
|
||
def _autolog(
|
||
log_traces: bool,
|
||
disable: bool = False,
|
||
silent: bool = False,
|
||
):
|
||
pass
|