genesis-flow 1.0.0__py3-none-any.whl

mlflow/artifacts/__init__.py

@@ -0,0 +1,268 @@
+"""
+APIs for interacting with artifacts in MLflow
+"""
+
+import json
+import pathlib
+import posixpath
+import tempfile
+from typing import Any, Optional
+
+from mlflow.entities.file_info import FileInfo
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE
+from mlflow.tracking import _get_store
+from mlflow.tracking.artifact_utils import (
+    _download_artifact_from_uri,
+    _get_root_uri_and_artifact_path,
+    add_databricks_profile_info_to_artifact_uri,
+    get_artifact_repository,
+)
+
+
+def download_artifacts(
+    artifact_uri: Optional[str] = None,
+    run_id: Optional[str] = None,
+    artifact_path: Optional[str] = None,
+    dst_path: Optional[str] = None,
+    tracking_uri: Optional[str] = None,
+) -> str:
+    """Download an artifact file or directory to a local directory.
+
+    Args:
+        artifact_uri: URI pointing to the artifacts. Supported formats include:
+
+            * ``runs:/<run_id>/<artifact_path>``
+              Example: ``runs:/500cf58bee2b40a4a82861cc31a617b1/my_model.pkl``
+
+            * ``models:/<model_name>/<stage>``
+              Example: ``models:/my_model/Production``
+
+            * ``models:/<model_name>/<version>/path/to/model``
+              Example: ``models:/my_model/2/path/to/model``
+
+            * ``models:/<model_name>@<alias>/path/to/model``
+              Example: ``models:/my_model@staging/path/to/model``
+
+            * Cloud storage URIs: ``s3://<bucket>/<path>`` or ``gs://<bucket>/<path>``
+
+            * Tracking server artifact URIs: ``http://<host>/mlartifacts`` or
+              ``mlflow-artifacts://<host>/mlartifacts``
+
+            Exactly one of ``artifact_uri`` or ``run_id`` must be specified.
+        run_id: ID of the MLflow Run containing the artifacts. Exactly one of ``run_id`` or
+            ``artifact_uri`` must be specified.
+        artifact_path: (For use with ``run_id``) If specified, a path relative to the MLflow
+            Run's root directory containing the artifacts to download.
+        dst_path: Path of the local filesystem destination directory to which to download the
+            specified artifacts. If the directory does not exist, it is created. If
+            unspecified, the artifacts are downloaded to a new uniquely-named directory on
+            the local filesystem, unless the artifacts already exist on the local
+            filesystem, in which case their local path is returned directly.
+        tracking_uri: The tracking URI to be used when downloading artifacts.
+
+    Returns:
+        The location of the artifact file or directory on the local filesystem.
+    """
+    if (run_id, artifact_uri).count(None) != 1:
+        raise MlflowException(
+            message="Exactly one of `run_id` or `artifact_uri` must be specified",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+    elif artifact_uri is not None and artifact_path is not None:
+        raise MlflowException(
+            message="`artifact_path` cannot be specified if `artifact_uri` is specified",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    if dst_path is not None:
+        pathlib.Path(dst_path).mkdir(exist_ok=True, parents=True)
+
+    if artifact_uri is not None:
+        return _download_artifact_from_uri(
+            artifact_uri, output_path=dst_path, tracking_uri=tracking_uri
+        )
+
+    # Use `runs:/<run_id>/<artifact_path>` to download both run and model (if exists) artifacts
+    if run_id and artifact_path:
+        return _download_artifact_from_uri(
+            f"runs:/{posixpath.join(run_id, artifact_path)}",
+            output_path=dst_path,
+            tracking_uri=tracking_uri,
+        )
+
+    artifact_path = artifact_path if artifact_path is not None else ""
+
+    store = _get_store(store_uri=tracking_uri)
+    artifact_uri = store.get_run(run_id).info.artifact_uri
+    artifact_repo = get_artifact_repository(
+        add_databricks_profile_info_to_artifact_uri(artifact_uri, tracking_uri),
+        tracking_uri=tracking_uri,
+    )
+    return artifact_repo.download_artifacts(artifact_path, dst_path=dst_path)
+
+
+def list_artifacts(
+    artifact_uri: Optional[str] = None,
+    run_id: Optional[str] = None,
+    artifact_path: Optional[str] = None,
+    tracking_uri: Optional[str] = None,
+) -> list[FileInfo]:
+    """List artifacts at the specified URI.
+
+    Args:
+        artifact_uri: URI pointing to the artifacts, such as
+            ``"runs:/500cf58bee2b40a4a82861cc31a617b1/my_model.pkl"``,
+            ``"models:/my_model/Production"``, or ``"s3://my_bucket/my/file.txt"``.
+            Exactly one of ``artifact_uri`` or ``run_id`` must be specified.
+        run_id: ID of the MLflow Run containing the artifacts. Exactly one of ``run_id`` or
+            ``artifact_uri`` must be specified.
+        artifact_path: (For use with ``run_id``) If specified, a path relative to the MLflow
+            Run's root directory containing the artifacts to list.
+        tracking_uri: The tracking URI to be used when list artifacts.
+
+    Returns:
+        List of artifacts as FileInfo listed directly under path.
+    """
+    if (run_id, artifact_uri).count(None) != 1:
+        raise MlflowException.invalid_parameter_value(
+            message="Exactly one of `run_id` or `artifact_uri` must be specified",
+        )
+    elif artifact_uri is not None and artifact_path is not None:
+        raise MlflowException.invalid_parameter_value(
+            message="`artifact_path` cannot be specified if `artifact_uri` is specified",
+        )
+
+    if artifact_uri is not None:
+        root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
+        return get_artifact_repository(
+            artifact_uri=root_uri, tracking_uri=tracking_uri
+        ).list_artifacts(artifact_path)
+
+    # Use `runs:/<run_id>/<artifact_path>` to list both run and model (if exists) artifacts
+    if run_id and artifact_path:
+        return get_artifact_repository(
+            artifact_uri=f"runs:/{run_id}", tracking_uri=tracking_uri
+        ).list_artifacts(artifact_path)
+
+    store = _get_store(store_uri=tracking_uri)
+    artifact_uri = store.get_run(run_id).info.artifact_uri
+    artifact_repo = get_artifact_repository(
+        add_databricks_profile_info_to_artifact_uri(artifact_uri, tracking_uri),
+        tracking_uri=tracking_uri,
+    )
+    return artifact_repo.list_artifacts(artifact_path)
+
+
+def load_text(artifact_uri: str) -> str:
+    """Loads the artifact contents as a string.
+
+    Args:
+        artifact_uri: Artifact location.
+
+    Returns:
+        The contents of the artifact as a string.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        with mlflow.start_run() as run:
+            artifact_uri = run.info.artifact_uri
+            mlflow.log_text("This is a sentence", "file.txt")
+            file_content = mlflow.artifacts.load_text(artifact_uri + "/file.txt")
+            print(file_content)
+
+    .. code-block:: text
+        :caption: Output
+
+        This is a sentence
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        with open(local_artifact) as local_artifact_fd:
+            try:
+                return str(local_artifact_fd.read())
+            except Exception:
+                raise MlflowException("Unable to form a str object from file content", BAD_REQUEST)
+
+
+def load_dict(artifact_uri: str) -> dict[str, Any]:
+    """Loads the artifact contents as a dictionary.
+
+    Args:
+        artifact_uri: artifact location.
+
+    Returns:
+        A dictionary.
+
+    .. code-block:: python
+      :caption: Example
+
+      import mlflow
+
+      with mlflow.start_run() as run:
+          artifact_uri = run.info.artifact_uri
+          mlflow.log_dict({"mlflow-version": "0.28", "n_cores": "10"}, "config.json")
+          config_json = mlflow.artifacts.load_dict(artifact_uri + "/config.json")
+          print(config_json)
+
+    .. code-block:: text
+      :caption: Output
+
+      {'mlflow-version': '0.28', 'n_cores': '10'}
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        with open(local_artifact) as local_artifact_fd:
+            try:
+                return json.load(local_artifact_fd)
+            except json.JSONDecodeError:
+                raise MlflowException("Unable to form a JSON object from file content", BAD_REQUEST)
+
+
+def load_image(artifact_uri: str):
+    """Loads artifact contents as a ``PIL.Image.Image`` object
+
+    Args:
+        artifact_uri: Artifact location.
+
+    Returns:
+        A PIL.Image object.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        from PIL import Image
+
+        with mlflow.start_run() as run:
+            image = Image.new("RGB", (100, 100))
+            artifact_uri = run.info.artifact_uri
+            mlflow.log_image(image, "image.png")
+            image = mlflow.artifacts.load_image(artifact_uri + "/image.png")
+            print(image)
+
+    .. code-block:: text
+        :caption: Output
+
+        <PIL.PngImagePlugin.PngImageFile image mode=RGB size=100x100 at 0x11D2FA3D0>
+    """
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ImportError(
+            "`load_image` requires Pillow. Please install it via: pip install Pillow"
+        ) from exc
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        local_artifact = download_artifacts(artifact_uri, dst_path=tmpdir)
+        try:
+            image_obj = Image.open(local_artifact)
+            image_obj.load()
+            return image_obj
+        except Exception:
+            raise MlflowException(
+                "Unable to form a PIL Image object from file content", BAD_REQUEST
+            )

mlflow/autogen/__init__.py

@@ -0,0 +1,144 @@
+import logging
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+import mlflow
+from mlflow.autogen.chat import (
+    convert_assistant_message_to_chat_message,
+    log_chat_messages,
+    log_tools,
+)
+from mlflow.entities import SpanType
+from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
+from mlflow.tracing.utils import construct_full_inputs, set_span_chat_messages
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import (
+    autologging_integration,
+    get_autologging_config,
+    safe_patch,
+)
+
+_logger = logging.getLogger(__name__)
+FLAVOR_NAME = "autogen"
+
+
+@experimental(version="2.16.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging for AutoGen flavor.
+    Due to its patch design, this method needs to be called after importing AutoGen classes.
+
+    Args:
+        log_traces: If ``True``, traces are logged for AutoGen models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the AutoGen autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during AutoGen
+            autologging. If ``False``, show all events and warnings.
+
+    Example:
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        from autogen_agentchat.agents import AssistantAgent
+        from autogen_ext.models.openai import OpenAIChatCompletionClient
+
+        mlflow.autogen.autolog()
+        agent = AssistantAgent("assistant", OpenAIChatCompletionClient(model="gpt-4o-mini"))
+        result = await agent.run(task="Say 'Hello World!'")
+        print(result)
+    """
+    from autogen_agentchat.agents import BaseChatAgent
+    from autogen_core.models import ChatCompletionClient
+
+    async def patched_completion(original, self, *args, **kwargs):
+        if not get_autologging_config(FLAVOR_NAME, "log_traces"):
+            return await original(self, *args, **kwargs)
+        else:
+            with mlflow.start_span(original.__name__, span_type=SpanType.LLM) as span:
+                inputs = construct_full_inputs(original, self, *args, **kwargs)
+                span.set_inputs(
+                    {key: _convert_value_to_dict(value) for key, value in inputs.items()}
+                )
+
+                if tools := inputs.get("tools"):
+                    log_tools(span, tools)
+
+                if messages := inputs.get("messages"):
+                    log_chat_messages(span, messages)
+
+                outputs = await original(self, *args, **kwargs)
+
+                if content := getattr(outputs, "content", None):
+                    if chat_message := convert_assistant_message_to_chat_message(content):
+                        set_span_chat_messages(span, [chat_message], append=True)
+
+                if usage := _parse_usage(outputs):
+                    span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage)
+
+                span.set_outputs(_convert_value_to_dict(outputs))
+
+                return outputs
+
+    async def patched_agent(original, self, *args, **kwargs):
+        if not get_autologging_config(FLAVOR_NAME, "log_traces"):
+            return await original(self, *args, **kwargs)
+        else:
+            with mlflow.start_span(original.__name__, span_type=SpanType.AGENT) as span:
+                inputs = construct_full_inputs(original, self, *args, **kwargs)
+                span.set_inputs(
+                    {key: _convert_value_to_dict(value) for key, value in inputs.items()}
+                )
+
+                if tools := getattr(self, "_tools", None):
+                    log_tools(span, tools)
+
+                outputs = await original(self, *args, **kwargs)
+
+                span.set_outputs(_convert_value_to_dict(outputs))
+
+                return outputs
+
+    for cls in BaseChatAgent.__subclasses__():
+        safe_patch(FLAVOR_NAME, cls, "run", patched_agent)
+        safe_patch(FLAVOR_NAME, cls, "on_messages", patched_agent)
+
+    for cls in _get_all_subclasses(ChatCompletionClient):
+        safe_patch(FLAVOR_NAME, cls, "create", patched_completion)
+
+
+def _convert_value_to_dict(value):
+    # BaseChatMessage does not contain content and type attributes
+    return value.model_dump(serialize_as_any=True) if isinstance(value, BaseModel) else value
+
+
+def _get_all_subclasses(cls):
+    """Get all subclasses recursively"""
+    all_subclasses = []
+
+    for subclass in cls.__subclasses__():
+        all_subclasses.append(subclass)
+        all_subclasses.extend(_get_all_subclasses(subclass))
+
+    return all_subclasses
+
+
+def _parse_usage(output: Any) -> Optional[dict[str, int]]:
+    try:
+        usage = getattr(output, "usage", None)
+        if usage:
+            return {
+                TokenUsageKey.INPUT_TOKENS: usage.prompt_tokens,
+                TokenUsageKey.OUTPUT_TOKENS: usage.completion_tokens,
+                TokenUsageKey.TOTAL_TOKENS: usage.prompt_tokens + usage.completion_tokens,
+            }
+    except Exception as e:
+        _logger.debug(f"Failed to parse token usage from output: {e}")
+    return None

mlflow/autogen/chat.py

@@ -0,0 +1,142 @@
+import logging
+from typing import TYPE_CHECKING, Optional, Union
+
+from opentelemetry.sdk.trace import Span
+
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    TextContentPart,
+    ToolCall,
+)
+
+if TYPE_CHECKING:
+    from autogen_core import FunctionCall
+    from autogen_core.models import LLMMessage
+    from autogen_core.tools import BaseTool, ToolSchema
+
+_logger = logging.getLogger(__name__)
+
+
+def log_tools(span: Span, tools: list[Union["BaseTool", "ToolSchema"]]):
+    """
+    Log Autogen tool definitions into the passed in span.
+
+    Ref: https://microsoft.github.io/autogen/stable/user-guide/core-user-guide/components/tools.html
+
+    Args:
+        span: The span to log the tools into.
+        tools: A list of Autogen BaseTool.
+    """
+    from autogen_core.tools import BaseTool
+
+    try:
+        tools = [
+            ChatTool(
+                type="function",
+                function=tool.schema if isinstance(tool, BaseTool) else tool,
+            )
+            for tool in tools
+        ]
+        set_span_chat_tools(span, tools)
+    except Exception:
+        _logger.debug(f"Failed to log tools to Span {span}.", exc_info=True)
+
+
+def convert_assistant_message_to_chat_message(
+    content: Union[str, list["FunctionCall"]],
+) -> Optional[ChatMessage]:
+    """
+    Convert an Autogen assistant message to a ChatMessage.
+    The content of assistant message is a str or a list of tool calls.
+
+    Args:
+        content: The Autogen assistant message content to convert.
+
+    Returns:
+        A ChatMessage object
+    """
+    from autogen_core import FunctionCall
+
+    if isinstance(content, str):
+        return ChatMessage(role="assistant", content=content)
+    elif isinstance(content, list) and all(isinstance(f, FunctionCall) for f in content):
+        return ChatMessage(
+            role="assistant",
+            tool_calls=[
+                ToolCall(
+                    id=f.id, type="function", function=Function(name=f.name, arguments=f.arguments)
+                )
+                for f in content
+            ],
+        )
+    else:
+        _logger.debug(f"Unsupported message type: {type(content)}. Skipping conversion.")
+
+
+def log_chat_messages(span: Span, messages: list["LLMMessage"]):
+    """
+    Log Autogen chat messages into the passed in span.
+
+    Args:
+        span: The span to log the tools into.
+        messages: A list of Autogen chat messages.
+    """
+    from autogen_core.models import (
+        AssistantMessage,
+        FunctionExecutionResultMessage,
+        SystemMessage,
+        UserMessage,
+    )
+
+    chat_messages = []
+    for message in messages:
+        if isinstance(message, SystemMessage):
+            chat_messages.append(
+                ChatMessage(
+                    role="system",
+                    content=message.content,
+                )
+            )
+        elif isinstance(message, UserMessage):
+            content = message.content
+            if isinstance(content, list):
+                parts = []
+                for part in content:
+                    if isinstance(part, str):
+                        parts.append(
+                            TextContentPart(
+                                type="text",
+                                text=part,
+                            )
+                        )
+                    else:
+                        # The content type of UserMessage is text or image
+                        parts.append(part.to_openai_format())
+                content = parts
+            chat_messages.append(
+                ChatMessage(
+                    role="user",
+                    content=content,
+                )
+            )
+        elif isinstance(message, AssistantMessage):
+            content = message.content
+            if chat_message := convert_assistant_message_to_chat_message(content):
+                chat_messages.append(chat_message)
+        elif isinstance(message, FunctionExecutionResultMessage):
+            chat_messages.append(
+                ChatMessage(
+                    role="user",
+                    content=message.model_dump(),
+                )
+            )
+        else:
+            _logger.debug(f"Unsupported message type: {type(message)}. Skipping logging.")
+
+    try:
+        set_span_chat_messages(span, chat_messages)
+    except Exception:
+        _logger.debug(f"Failed to log chat messages to Span {span}.", exc_info=True)

mlflow/azure/__init__.py

@@ -0,0 +1,26 @@
+"""Azure authentication support for MLflow."""
+
+from mlflow.azure.config import AzureAuthConfig, AuthMethod
+from mlflow.azure.auth_handler import AzureAuthHandler
+from mlflow.azure.connection_factory import ConnectionFactory
+from mlflow.azure.stores import create_store, get_azure_tracking_store, test_azure_connection
+from mlflow.azure.exceptions import (
+    AzureAuthError,
+    TokenAcquisitionError,
+    ConnectionError,
+    ConfigurationError,
+)
+
+__all__ = [
+    "AzureAuthConfig",
+    "AuthMethod", 
+    "AzureAuthHandler",
+    "ConnectionFactory",
+    "create_store",
+    "get_azure_tracking_store",
+    "test_azure_connection",
+    "AzureAuthError",
+    "TokenAcquisitionError",
+    "ConnectionError",
+    "ConfigurationError",
+]