genesis-flow 1.0.0__py3-none-any.whl

mlflow/client.py

@@ -0,0 +1,12 @@
+"""
+The ``mlflow.client`` module provides a Python CRUD interface to MLflow Experiments, Runs,
+Model Versions, and Registered Models. This is a lower level API that directly translates to MLflow
+`REST API <../rest-api.html>`_ calls.
+For a higher level API for managing an "active run", use the :py:mod:`mlflow` module.
+"""
+
+from mlflow.tracking.client import MlflowClient
+
+__all__ = [
+    "MlflowClient",
+]

mlflow/config/__init__.py

@@ -0,0 +1,56 @@
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_ASYNC_LOGGING,
+)
+from mlflow.system_metrics import (
+    disable_system_metrics_logging,
+    enable_system_metrics_logging,
+    set_system_metrics_node_id,
+    set_system_metrics_samples_before_logging,
+    set_system_metrics_sampling_interval,
+)
+from mlflow.tracking import (
+    get_registry_uri,
+    get_tracking_uri,
+    is_tracking_uri_set,
+    set_registry_uri,
+    set_tracking_uri,
+)
+
+
+def enable_async_logging(enable=True):
+    """Enable or disable async logging globally.
+
+    Args:
+        enable: bool, if True, enable async logging. If False, disable async logging.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        mlflow.config.enable_async_logging(True)
+
+        with mlflow.start_run():
+            mlflow.log_param("a", 1)  # This will be logged asynchronously
+
+        mlflow.config.enable_async_logging(False)
+        with mlflow.start_run():
+            mlflow.log_param("a", 1)  # This will be logged synchronously
+    """
+
+    MLFLOW_ENABLE_ASYNC_LOGGING.set(enable)
+
+
+__all__ = [
+    "enable_system_metrics_logging",
+    "disable_system_metrics_logging",
+    "enable_async_logging",
+    "get_registry_uri",
+    "get_tracking_uri",
+    "is_tracking_uri_set",
+    "set_registry_uri",
+    "set_system_metrics_sampling_interval",
+    "set_system_metrics_samples_before_logging",
+    "set_system_metrics_node_id",
+    "set_tracking_uri",
+]

mlflow/crewai/__init__.py

@@ -0,0 +1,79 @@
+"""
+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,
+)
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+_logger = logging.getLogger(__name__)
+
+FLAVOR_NAME = "crewai"
+
+
+@experimental(version="2.19.0")
+@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
+    # TODO: Tool calling is not supported since the interface of tool in CrewAI is
+    # changing drastically. Add patching once it's stabilized
+    import crewai
+
+    class_method_map = {
+        "crewai.Crew": ["kickoff", "kickoff_for_each", "train"],
+        "crewai.Agent": ["execute_task"],
+        "crewai.Task": ["execute_sync"],
+        "crewai.LLM": ["call"],
+        "crewai.Flow": ["kickoff"],
+        "crewai.agents.agent_builder.base_agent_executor_mixin.CrewAgentExecutorMixin": [
+            "_create_long_term_memory"
+        ],
+    }
+    if Version(crewai.__version__) >= Version("0.83.0"):
+        # knowledge and memory are not available before 0.83.0
+        class_method_map.update(
+            {
+                "crewai.memory.ShortTermMemory": ["save", "search"],
+                "crewai.memory.LongTermMemory": ["save", "search"],
+                "crewai.memory.UserMemory": ["save", "search"],
+                "crewai.memory.EntityMemory": ["save", "search"],
+                "crewai.Knowledge": ["query"],
+            }
+        )
+    try:
+        for class_path, methods in class_method_map.items():
+            *module_parts, class_name = class_path.rsplit(".", 1)
+            module_path = ".".join(module_parts)
+            module = importlib.import_module(module_path)
+            cls = getattr(module, class_name)
+            for method in methods:
+                safe_patch(
+                    FLAVOR_NAME,
+                    cls,
+                    method,
+                    patched_class_call,
+                )
+    except (AttributeError, ModuleNotFoundError) as e:
+        _logger.error("An exception happens when applying auto-tracing to crewai. Exception: %s", e)

mlflow/crewai/autolog.py

@@ -0,0 +1,253 @@
+import inspect
+import json
+import logging
+import warnings
+
+from packaging.version import Version
+
+import mlflow
+from mlflow.crewai.chat import set_span_chat_attributes
+from mlflow.entities import SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.tracing.utils import TraceJSONEncoder
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    config = AutoLoggingConfig.init(flavor_name=mlflow.gemini.FLAVOR_NAME)
+
+    if config.log_traces:
+        fullname = f"{self.__class__.__name__}.{original.__name__}"
+        span_type = _get_span_type(self)
+        with mlflow.start_span(name=fullname, span_type=span_type) as span:
+            inputs = _construct_full_inputs(original, self, *args, **kwargs)
+            span.set_inputs(inputs)
+            _set_span_attributes(span=span, instance=self)
+            result = original(self, *args, **kwargs)
+
+            if span_type == SpanType.LLM:
+                set_span_chat_attributes(
+                    span=span, messages=inputs.get("messages", []), output=result
+                )
+            # Need to convert the response of generate_content for better visualization
+            outputs = result.__dict__ if hasattr(result, "__dict__") else result
+            span.set_outputs(outputs)
+
+            return result
+
+
+def _get_span_type(instance) -> str:
+    import crewai
+    from crewai import LLM, Agent, Crew, Task
+    from crewai.flow.flow import Flow
+
+    try:
+        if isinstance(instance, (Flow, Crew, Task)):
+            return SpanType.CHAIN
+        elif isinstance(instance, Agent):
+            return SpanType.AGENT
+        elif isinstance(instance, LLM):
+            return SpanType.LLM
+        elif isinstance(instance, Flow):
+            return SpanType.CHAIN
+        elif isinstance(
+            instance, crewai.agents.agent_builder.base_agent_executor_mixin.CrewAgentExecutorMixin
+        ):
+            return SpanType.RETRIEVER
+
+        # Knowledge and Memory are not available before 0.83.0
+        if Version(crewai.__version__) >= Version("0.83.0"):
+            if isinstance(
+                instance,
+                (
+                    crewai.memory.ShortTermMemory,
+                    crewai.memory.LongTermMemory,
+                    crewai.memory.UserMemory,
+                    crewai.memory.EntityMemory,
+                    crewai.Knowledge,
+                ),
+            ):
+                return SpanType.RETRIEVER
+    except AttributeError as e:
+        _logger.warn("An exception happens when resolving the span type. Exception: %s", e)
+
+    return SpanType.UNKNOWN
+
+
+def _is_serializable(value):
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            # There is type mismatch in some crewai class, suppress warning here
+            json.dumps(value, cls=TraceJSONEncoder, ensure_ascii=False)
+        return True
+    except (TypeError, ValueError):
+        return False
+
+
+def _construct_full_inputs(func, *args, **kwargs):
+    signature = inspect.signature(func)
+    # This does not create copy. So values should not be mutated directly
+    arguments = signature.bind_partial(*args, **kwargs).arguments
+
+    if "self" in arguments:
+        arguments.pop("self")
+
+    # Avoid non serializable objects and circular references
+    return {
+        k: v.__dict__ if hasattr(v, "__dict__") else v
+        for k, v in arguments.items()
+        if v is not None and _is_serializable(v)
+    }
+
+
+def _set_span_attributes(span: LiveSpan, instance):
+    # Crewai is available only python >=3.10, so importing libraries inside methods.
+    try:
+        import crewai
+        from crewai import LLM, Agent, Crew, Task
+        from crewai.flow.flow import Flow
+
+        ## Memory class does not have helpful attributes
+        if isinstance(instance, Crew):
+            for key, value in instance.__dict__.items():
+                if value is not None:
+                    if key == "tasks":
+                        value = _parse_tasks(value)
+                    elif key == "agents":
+                        value = _parse_agents(value)
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif isinstance(instance, Agent):
+            agent = _get_agent_attributes(instance)
+            for key, value in agent.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif isinstance(instance, Task):
+            task = _get_task_attributes(instance)
+            for key, value in task.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif isinstance(instance, LLM):
+            llm = _get_llm_attributes(instance)
+            for key, value in llm.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif isinstance(instance, Flow):
+            for key, value in instance.__dict__.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif Version(crewai.__version__) >= Version("0.83.0"):
+            if isinstance(instance, crewai.Knowledge):
+                for key, value in instance.__dict__.items():
+                    if value is not None and key != "storage":
+                        span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+    except AttributeError as e:
+        _logger.warn("An exception happens when saving span attributes. Exception: %s", e)
+
+
+def _get_agent_attributes(instance):
+    agent = {}
+    for key, value in instance.__dict__.items():
+        if key == "tools":
+            value = _parse_tools(value)
+        if value is None:
+            continue
+        agent[key] = str(value)
+
+    return agent
+
+
+def _get_task_attributes(instance):
+    task = {}
+    for key, value in instance.__dict__.items():
+        if value is None:
+            continue
+        if key == "tools":
+            value = _parse_tools(value)
+            task[key] = value
+        elif key == "agent":
+            task[key] = value.role
+        else:
+            task[key] = str(value)
+    return task
+
+
+def _get_llm_attributes(instance):
+    llm = {}
+    for key, value in instance.__dict__.items():
+        if value is None:
+            continue
+        elif key in ["callbacks", "api_key"]:
+            # Skip callbacks until how they should be logged are decided
+            continue
+        else:
+            llm[key] = str(value)
+    return llm
+
+
+def _parse_agents(agents):
+    attributes = []
+    for agent in agents:
+        model = None
+        if agent.llm is not None:
+            if hasattr(agent.llm, "model"):
+                model = agent.llm.model
+            elif hasattr(agent.llm, "model_name"):
+                model = agent.llm.model_name
+        attributes.append(
+            {
+                "id": str(agent.id),
+                "role": agent.role,
+                "goal": agent.goal,
+                "backstory": agent.backstory,
+                "cache": agent.cache,
+                "config": agent.config,
+                "verbose": agent.verbose,
+                "allow_delegation": agent.allow_delegation,
+                "tools": agent.tools,
+                "max_iter": agent.max_iter,
+                "llm": str(model if model is not None else ""),
+            }
+        )
+    return attributes
+
+
+def _parse_tasks(tasks):
+    return [
+        {
+            "agent": task.agent.role,
+            "description": task.description,
+            "async_execution": task.async_execution,
+            "expected_output": task.expected_output,
+            "human_input": task.human_input,
+            "tools": task.tools,
+            "output_file": task.output_file,
+        }
+        for task in tasks
+    ]
+
+
+def _parse_tools(tools):
+    result = []
+    for tool in tools:
+        res = {}
+        if hasattr(tool, "name") and tool.name is not None:
+            res["name"] = tool.name
+        if hasattr(tool, "description") and tool.description is not None:
+            res["description"] = tool.description
+        if res:
+            result.append(
+                {
+                    "type": "function",
+                    "function": res,
+                }
+            )
+    return result

mlflow/crewai/chat.py

@@ -0,0 +1,29 @@
+import logging
+from typing import Any
+
+from mlflow.entities.span import LiveSpan
+from mlflow.exceptions import MlflowException
+from mlflow.tracing import set_span_chat_messages
+
+_logger = logging.getLogger(__name__)
+
+
+def set_span_chat_attributes(span: LiveSpan, messages: list[dict[str, Any]], output: Any):
+    """
+    This method logs chat messages to a span when LLM class is called. Since the CrewAI library
+    includes tool definition in the text content, mlflow.chat.tools attribute is not recorded.
+
+    Args:
+        span: A span object to record the chat messages.
+        messages: Input messages for the LLM call.
+        output: Response from LLM. Though the signature is str, tool response might be returned.
+    """
+    output_message = {"role": "assistant", "content": str(output)}
+
+    try:
+        set_span_chat_messages(span, [*messages, output_message])
+    except MlflowException:
+        _logger.debug(
+            "Failed to set chat messages on span",
+            exc_info=True,
+        )

mlflow/data/__init__.py

@@ -0,0 +1,75 @@
+import sys
+from contextlib import suppress
+from typing import Union
+
+from mlflow.data import dataset_registry
+from mlflow.data import sources as mlflow_data_sources
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json, get_registered_sources
+from mlflow.entities import Dataset as DatasetEntity
+from mlflow.entities import DatasetInput
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+with suppress(ImportError):
+    # Suppressing ImportError to pass mlflow-skinny testing.
+    from mlflow.data import meta_dataset  # noqa: F401
+
+
+def get_source(dataset: Union[DatasetEntity, DatasetInput, Dataset]) -> DatasetSource:
+    """Obtains the source of the specified dataset or dataset input.
+
+    Args:
+        dataset:
+            An instance of :py:class:`mlflow.data.dataset.Dataset <mlflow.data.dataset.Dataset>`,
+            :py:class:`mlflow.entities.Dataset`, or :py:class:`mlflow.entities.DatasetInput`.
+
+    Returns:
+        An instance of :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+
+    """
+    if isinstance(dataset, DatasetInput):
+        dataset: DatasetEntity = dataset.dataset
+
+    if isinstance(dataset, DatasetEntity):
+        dataset_source: DatasetSource = get_dataset_source_from_json(
+            source_json=dataset.source,
+            source_type=dataset.source_type,
+        )
+    elif isinstance(dataset, Dataset):
+        dataset_source: DatasetSource = dataset.source
+    else:
+        raise MlflowException(
+            f"Unrecognized dataset type {type(dataset)}. Expected one of: "
+            f"`mlflow.data.dataset.Dataset`,"
+            f" `mlflow.entities.Dataset`, `mlflow.entities.DatasetInput`.",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    return dataset_source
+
+
+__all__ = ["get_source"]
+
+
+def _define_dataset_constructors_in_current_module():
+    data_module = sys.modules[__name__]
+    for (
+        constructor_name,
+        constructor_fn,
+    ) in dataset_registry.get_registered_constructors().items():
+        setattr(data_module, constructor_name, constructor_fn)
+        __all__.append(constructor_name)
+
+
+_define_dataset_constructors_in_current_module()
+
+
+def _define_dataset_sources_in_sources_module():
+    for source in get_registered_sources():
+        setattr(mlflow_data_sources, source.__name__, source)
+        mlflow_data_sources.__all__.append(source.__name__)
+
+
+_define_dataset_sources_in_sources_module()

mlflow/data/artifact_dataset_sources.py

@@ -0,0 +1,170 @@
+import re
+import warnings
+from pathlib import Path
+from typing import Any, TypeVar
+from urllib.parse import urlparse
+
+from mlflow.artifacts import download_artifacts
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repository_registry import get_registered_artifact_repositories
+from mlflow.utils.uri import is_local_uri
+
+
+def register_artifact_dataset_sources():
+    from mlflow.data.dataset_source_registry import register_dataset_source
+
+    registered_source_schemes = set()
+    artifact_schemes_to_exclude = [
+        "http",
+        "https",
+        "runs",
+        "models",
+        "mlflow-artifacts",
+        # DBFS supports two access patterns: dbfs:/ (URI) and /dbfs (FUSE).
+        # The DBFS artifact repository online supports dbfs:/ (URI). To ensure
+        # a consistent dictionary representation of DBFS datasets across the URI and
+        # FUSE representations, we exclude dbfs from the set of dataset sources
+        # that are autogenerated using artifact repositories and instead define
+        # a separate DBFSDatasetSource elsewhere
+        "dbfs",
+    ]
+    schemes_to_artifact_repos = get_registered_artifact_repositories()
+    for scheme, artifact_repo in schemes_to_artifact_repos.items():
+        if scheme in artifact_schemes_to_exclude or scheme in registered_source_schemes:
+            continue
+
+        if "ArtifactRepository" in artifact_repo.__name__:
+            # Artifact repository name is something like "LocalArtifactRepository",
+            # "S3ArtifactRepository", etc. To preserve capitalization, strip ArtifactRepository
+            # and replace it with ArtifactDatasetSource
+            dataset_source_name = artifact_repo.__name__.replace(
+                "ArtifactRepository", "ArtifactDatasetSource"
+            )
+        else:
+            # Artifact repository name has some other form, e.g. "dbfs_artifact_repo_factory".
+            # In this case, generate the name by capitalizing the first letter of the scheme and
+            # appending ArtifactRepository
+            scheme = str(scheme)
+
+            def camelcase_scheme(scheme):
+                parts = re.split(r"[-_]", scheme)
+                return "".join([part.capitalize() for part in parts])
+
+            source_name_prefix = camelcase_scheme(scheme)
+            dataset_source_name = source_name_prefix + "ArtifactDatasetSource"
+
+        try:
+            registered_source_schemes.add(scheme)
+            dataset_source = _create_dataset_source_for_artifact_repo(
+                scheme=scheme, dataset_source_name=dataset_source_name
+            )
+            register_dataset_source(dataset_source)
+        except Exception as e:
+            warnings.warn(
+                f"Failed to register a dataset source for URIs with scheme '{scheme}': {e}",
+                stacklevel=2,
+            )
+
+
+def _create_dataset_source_for_artifact_repo(scheme: str, dataset_source_name: str):
+    from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+
+    if scheme in ["", "file"]:
+        source_type = "local"
+        class_docstring = "Represents the source of a dataset stored on the local filesystem."
+    else:
+        source_type = scheme
+        class_docstring = (
+            f"Represents a filesystem-based or blob-storage-based dataset source identified by a"
+            f" URI with scheme '{scheme}'."
+        )
+
+    DatasetForArtifactRepoSourceType = TypeVar(dataset_source_name)
+
+    class ArtifactRepoSource(FileSystemDatasetSource):
+        def __init__(self, uri: str):
+            self._uri = uri
+
+        @property
+        def uri(self):
+            """
+            The URI with scheme '{scheme}' referring to the dataset source filesystem location.
+
+            Returns
+                The URI with scheme '{scheme}' referring to the dataset source filesystem
+                location.
+            """
+            return self._uri
+
+        @staticmethod
+        def _get_source_type() -> str:
+            return source_type
+
+        def load(self, dst_path=None) -> str:
+            """
+            Downloads the dataset source to the local filesystem.
+
+            Args:
+                dst_path: Path of the local filesystem destination directory to which to download
+                    the dataset source. If the directory does not exist, it is created. If
+                    unspecified, the dataset source is downloaded to a new uniquely-named
+                    directory on the local filesystem, unless the dataset source already
+                    exists on the local filesystem, in which case its local path is
+                    returned directly.
+
+            Returns:
+                The path to the downloaded dataset source on the local filesystem.
+            """
+            return download_artifacts(artifact_uri=self.uri, dst_path=dst_path)
+
+        @staticmethod
+        def _can_resolve(raw_source: Any):
+            is_local_source_type = ArtifactRepoSource._get_source_type() == "local"
+
+            if not isinstance(raw_source, str) and (
+                not isinstance(raw_source, Path) and is_local_source_type
+            ):
+                return False
+
+            try:
+                if is_local_source_type:
+                    return is_local_uri(str(raw_source), is_tracking_or_registry_uri=False)
+                else:
+                    parsed_source = urlparse(str(raw_source))
+                    return parsed_source.scheme == scheme
+            except Exception:
+                return False
+
+        @classmethod
+        def _resolve(cls, raw_source: Any) -> DatasetForArtifactRepoSourceType:
+            return cls(str(raw_source))
+
+        def to_dict(self) -> dict[Any, Any]:
+            """
+            Returns:
+                A JSON-compatible dictionary representation of the {dataset_source_name}.
+            """
+            return {
+                "uri": self.uri,
+            }
+
+        @classmethod
+        def from_dict(cls, source_dict: dict[Any, Any]) -> DatasetForArtifactRepoSourceType:
+            uri = source_dict.get("uri")
+            if uri is None:
+                raise MlflowException(
+                    f'Failed to parse {dataset_source_name}. Missing expected key: "uri"',
+                    INVALID_PARAMETER_VALUE,
+                )
+
+            return cls(uri=uri)
+
+    ArtifactRepoSource.__name__ = dataset_source_name
+    ArtifactRepoSource.__qualname__ = dataset_source_name
+    ArtifactRepoSource.__doc__ = class_docstring
+    ArtifactRepoSource.to_dict.__doc__ = ArtifactRepoSource.to_dict.__doc__.format(
+        dataset_source_name=dataset_source_name
+    )
+    ArtifactRepoSource.uri.__doc__ = ArtifactRepoSource.uri.__doc__.format(scheme=scheme)
+    return ArtifactRepoSource