wehub-resource-sync
126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
"""
|
|
The ``mlflow.crewai`` module provides an API for tracing CrewAI AI agents.
|
|
"""
|
|
|
|
import importlib
|
|
import logging
|
|
|
|
from packaging.version import Version
|
|
|
|
from mlflow.crewai.autolog import (
|
|
patched_class_call,
|
|
patched_native_tool_call,
|
|
patched_standalone_call,
|
|
)
|
|
from mlflow.telemetry.events import AutologgingEvent
|
|
from mlflow.telemetry.track import _record_event
|
|
from mlflow.utils.autologging_utils import autologging_integration, safe_patch
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
FLAVOR_NAME = "crewai"
|
|
|
|
|
|
@autologging_integration(FLAVOR_NAME)
|
|
def autolog(
|
|
log_traces: bool = True,
|
|
disable: bool = False,
|
|
silent: bool = False,
|
|
):
|
|
"""
|
|
Enables (or disables) and configures autologging from CrewAI to MLflow.
|
|
Note that asynchronous APIs and Tool calling are not recorded now.
|
|
|
|
Args:
|
|
log_traces: If ``True``, traces are logged for CrewAI agents.
|
|
If ``False``, no traces are collected during inference. Default to ``True``.
|
|
disable: If ``True``, disables the CrewAI autologging. Default to ``False``.
|
|
silent: If ``True``, suppress all event logs and warnings from MLflow during CrewAI
|
|
autologging. If ``False``, show all events and warnings.
|
|
"""
|
|
# TODO: Handle asynchronous tasks and crew executions
|
|
import crewai
|
|
|
|
CREWAI_VERSION = Version(crewai.__version__)
|
|
|
|
# _create_long_term_memory was replaced by _save_to_memory in crewai 1.10.0
|
|
_memory_method = (
|
|
"_save_to_memory" if CREWAI_VERSION >= Version("1.10.0") else "_create_long_term_memory"
|
|
)
|
|
# crewai 1.14.5 renamed the module and class: base_agent_executor_mixin.CrewAgentExecutorMixin
|
|
# -> base_agent_executor.BaseAgentExecutor
|
|
_executor_path = (
|
|
"crewai.agents.agent_builder.base_agent_executor.BaseAgentExecutor"
|
|
if CREWAI_VERSION >= Version("1.14.5")
|
|
else "crewai.agents.agent_builder.base_agent_executor_mixin.CrewAgentExecutorMixin"
|
|
)
|
|
class_method_map = {
|
|
"crewai.Crew": ["kickoff", "kickoff_for_each", "train"],
|
|
"crewai.Agent": ["execute_task"],
|
|
"crewai.Task": ["execute_sync"],
|
|
"crewai.LLM": ["call"],
|
|
"crewai.Flow": ["kickoff"],
|
|
_executor_path: [_memory_method],
|
|
}
|
|
standalone_method_map = {}
|
|
|
|
if CREWAI_VERSION >= Version("0.83.0"):
|
|
# knowledge and memory are not available before 0.83.0
|
|
# ShortTermMemory/LongTermMemory/EntityMemory were replaced by unified MemoryScope in 1.10.0
|
|
if CREWAI_VERSION < Version("1.10.0"):
|
|
class_method_map.update({
|
|
"crewai.memory.ShortTermMemory": ["save", "search"],
|
|
"crewai.memory.LongTermMemory": ["save", "search"],
|
|
"crewai.memory.EntityMemory": ["save", "search"],
|
|
})
|
|
if CREWAI_VERSION < Version("0.157.0"):
|
|
class_method_map.update({"crewai.memory.UserMemory": ["save", "search"]})
|
|
class_method_map.update({"crewai.Knowledge": ["query"]})
|
|
|
|
# Modern Tool calling support for CrewAI >= 0.114.0
|
|
if CREWAI_VERSION >= Version("0.114.0"):
|
|
standalone_method_map.update({
|
|
"crewai.agents.crew_agent_executor": ["execute_tool_and_check_finality"]
|
|
})
|
|
|
|
# Native function calling support for CrewAI >= 1.9.0
|
|
native_tool_method_map = {}
|
|
if CREWAI_VERSION >= Version("1.9.0"):
|
|
native_tool_method_map["crewai.agents.crew_agent_executor.CrewAgentExecutor"] = [
|
|
"_handle_native_tool_calls"
|
|
]
|
|
|
|
try:
|
|
_apply_patches(standalone_method_map, _import_module, patched_standalone_call)
|
|
_apply_patches(class_method_map, _import_class, patched_class_call)
|
|
_apply_patches(native_tool_method_map, _import_class, patched_native_tool_call)
|
|
except (AttributeError, ModuleNotFoundError) as e:
|
|
_logger.error("An exception happens when applying auto-tracing to crewai. Exception: %s", e)
|
|
|
|
_record_event(
|
|
AutologgingEvent, {"flavor": FLAVOR_NAME, "log_traces": log_traces, "disable": disable}
|
|
)
|
|
|
|
|
|
def _apply_patches(target_map, resolver, patch_fn):
|
|
for target_path, methods in target_map.items():
|
|
target = resolver(target_path)
|
|
for method in methods:
|
|
safe_patch(
|
|
FLAVOR_NAME,
|
|
target,
|
|
method,
|
|
patch_fn,
|
|
)
|
|
|
|
|
|
def _import_module(module_path: str):
|
|
return importlib.import_module(module_path)
|
|
|
|
|
|
def _import_class(class_path: str):
|
|
*module_parts, class_name = class_path.rsplit(".", 1)
|
|
module_path = ".".join(module_parts)
|
|
module = importlib.import_module(module_path)
|
|
return getattr(module, class_name)
|