genesis-flow 1.0.0__py3-none-any.whl

mlflow/smolagents/__init__.py

@@ -0,0 +1,66 @@
+"""
+The ``mlflow.smolagents`` module provides an API for tracing Smolagents AI agents.
+"""
+
+import logging
+
+from mlflow.smolagents.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 = "smolagents"
+
+
+@experimental(version="3.0.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Smolagents to MLflow.
+    Note that asynchronous APIs and Tool calling are not recorded now.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Smolagents agents.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Smolagents autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Smolagents
+            autologging. If ``False``, show all events and warnings.
+    """
+    import smolagents
+    from smolagents import models
+
+    class_method_map = {
+        "MultiStepAgent": ["run"],
+        "CodeAgent": ["step"],
+        "ToolCallingAgent": ["step"],
+        "Tool": ["__call__"],
+    }
+
+    try:
+        for _, attr in vars(smolagents).items():
+            if isinstance(attr, type) and issubclass(attr, models.Model):
+                class_method_map.setdefault(attr.__name__, []).append("__call__")
+    except Exception as e:
+        _logger.warn("the error happens while registering models to class_method_map: %s", e)
+
+    try:
+        for class_name, methods in class_method_map.items():
+            cls = getattr(smolagents, 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 smolagents. Exception: %s", e
+        )

mlflow/smolagents/autolog.py

@@ -0,0 +1,139 @@
+import inspect
+import logging
+
+import mlflow
+from mlflow.entities import SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.smolagents.chat import set_span_chat_attributes
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    try:
+        config = AutoLoggingConfig.init(flavor_name=mlflow.smolagents.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)
+
+                # Need to convert the response of smolagents API for better visualization
+                outputs = result.__dict__ if hasattr(result, "__dict__") else result
+                if span_type == SpanType.CHAT_MODEL:
+                    set_span_chat_attributes(
+                        span=span, messages=inputs.get("messages", []), output=outputs
+                    )
+                span.set_outputs(outputs)
+                return result
+    except Exception as e:
+        _logger.error("the error occurred while patching")
+        raise e
+
+
+def _get_span_type(instance) -> str:
+    from smolagents import CodeAgent, MultiStepAgent, Tool, ToolCallingAgent, models
+
+    if isinstance(instance, (MultiStepAgent, CodeAgent, ToolCallingAgent)):
+        return SpanType.AGENT
+    elif isinstance(instance, Tool):
+        return SpanType.TOOL
+    elif isinstance(instance, models.Model):
+        return SpanType.CHAT_MODEL
+
+    return SpanType.UNKNOWN
+
+
+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
+    }
+
+
+def _set_span_attributes(span: LiveSpan, instance):
+    # Smolagents is available only python >= 3.10, so importing libraries inside methods.
+    try:
+        from smolagents import CodeAgent, MultiStepAgent, Tool, ToolCallingAgent, models
+
+        if isinstance(instance, (MultiStepAgent, CodeAgent, ToolCallingAgent)):
+            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, Tool):
+            tool = _get_tool_attributes(instance)
+            for key, value in tool.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+        elif issubclass(type(instance), models.Model):
+            model = _get_model_attributes(instance)
+            for key, value in model.items():
+                if value is not None:
+                    span.set_attribute(key, str(value) if isinstance(value, list) else value)
+
+    except Exception 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 _inner_get_tool_attributes(tool_dict):
+    res = {}
+    if hasattr(tool_dict, "name") and tool_dict.name is not None:
+        res["name"] = tool_dict.name
+    if hasattr(tool_dict, "description") and tool_dict.description is not None:
+        res["description"] = tool_dict.description
+    result = {}
+    if res:
+        result["type"] = "function"
+        result["function"] = res
+    return result
+
+
+def _get_tool_attributes(instance):
+    instance_dict = instance.__dict__
+    return _inner_get_tool_attributes(instance_dict)
+
+
+def _parse_tools(tools):
+    result = []
+    for tool in tools:
+        res = _inner_get_tool_attributes(tool)
+        result.append(res)
+    return result
+
+
+def _get_model_attributes(instance):
+    model = {}
+    for key, value in instance.__dict__.items():
+        if value is None or key == "api_key":
+            continue
+        model[key] = str(value)
+    return model

mlflow/smolagents/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.
+
+    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.
+    """
+    content = str(output.get("content", "")) if isinstance(output, dict) else str(output)
+    output_message = {"role": "assistant", "content": content}
+
+    try:
+        set_span_chat_messages(span, [*messages, output_message])
+    except MlflowException:
+        _logger.debug(
+            "Failed to set chat messages on span",
+            exc_info=True,
+        )

mlflow/store/__init__.py

@@ -0,0 +1,10 @@
+from mlflow.store import _unity_catalog  # noqa: F401
+from mlflow.store.artifact import artifact_repo
+from mlflow.store.tracking import abstract_store
+
+__all__ = [
+    # tracking server meta-data stores
+    "abstract_store",
+    # artifact repository stores
+    "artifact_repo",
+]

mlflow/store/_unity_catalog/__init__.py

@@ -0,0 +1 @@
+from mlflow.store._unity_catalog import registry  # noqa: F401

mlflow/store/_unity_catalog/lineage/__init__.py

@@ -0,0 +1 @@
+from mlflow.store._unity_catalog import lineage  # noqa: F401

mlflow/store/_unity_catalog/lineage/constants.py

@@ -0,0 +1,2 @@
+_DATABRICKS_ORG_ID_HEADER = "x-databricks-org-id"
+_DATABRICKS_LINEAGE_ID_HEADER = "X-Databricks-Lineage-Identifier"

mlflow/store/_unity_catalog/registry/__init__.py

@@ -0,0 +1,6 @@
+from mlflow.store._unity_catalog.registry import (
+    rest_store as rest_store,
+)
+from mlflow.store._unity_catalog.registry import (
+    uc_oss_rest_store as uc_oss_rest_store,
+)

mlflow/store/_unity_catalog/registry/prompt_info.py

@@ -0,0 +1,75 @@
+"""
+Internal PromptInfo entity for Unity Catalog prompt operations.
+
+This is an implementation detail for the Unity Catalog store and should not be
+considered part of the public MLflow API.
+"""
+
+from typing import Optional
+
+
+class PromptInfo:
+    """
+    Internal entity for prompt information from Unity Catalog. This represents
+    prompt metadata without version-specific details like template.
+
+    This maps to the Unity Catalog PromptInfo protobuf message.
+
+    Note: This is an internal implementation detail and not part of the public API.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        creation_timestamp: Optional[int] = None,
+        tags: Optional[dict[str, str]] = None,
+    ):
+        """
+        Construct a PromptInfo entity.
+
+        Args:
+            name: Name of the prompt.
+            description: Description of the prompt.
+            creation_timestamp: Timestamp when the prompt was created.
+            tags: Prompt-level metadata as key-value pairs.
+        """
+        self._name = name
+        self._description = description
+        self._creation_timestamp = creation_timestamp
+        self._tags = tags or {}
+
+    @property
+    def name(self) -> str:
+        """The name of the prompt."""
+        return self._name
+
+    @property
+    def description(self) -> Optional[str]:
+        """The description of the prompt."""
+        return self._description
+
+    @property
+    def creation_timestamp(self) -> Optional[int]:
+        """The creation timestamp of the prompt."""
+        return self._creation_timestamp
+
+    @property
+    def tags(self) -> dict[str, str]:
+        """Prompt-level metadata as key-value pairs."""
+        return self._tags.copy()
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, PromptInfo):
+            return False
+        return (
+            self.name == other.name
+            and self.description == other.description
+            and self.creation_timestamp == other.creation_timestamp
+            and self.tags == other.tags
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"<PromptInfo: name='{self.name}', description='{self.description}', tags={self.tags}>"
+        )