genesis-flow 1.0.0__py3-none-any.whl

mlflow/genai/label_schemas/__init__.py

@@ -0,0 +1,153 @@
+"""
+Databricks Agent Label Schemas Python SDK. For more details see Databricks Agent Evaluation:
+<https://docs.databricks.com/en/generative-ai/agent-evaluation/index.html>
+
+The API docs can be found here:
+<https://api-docs.databricks.com/python/databricks-agents/latest/databricks_agent_eval.html#review-app>
+"""
+
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+from mlflow.genai.label_schemas.label_schemas import (
+    InputCategorical,
+    InputCategoricalList,
+    InputNumeric,
+    InputText,
+    InputTextList,
+    LabelSchema,
+    LabelSchemaType,
+)
+from mlflow.genai.labeling import ReviewApp
+
+if TYPE_CHECKING:
+    from databricks.agents.review_app import ReviewApp
+
+_ERROR_MSG = (
+    "The `databricks-agents` package is required to use `mlflow.genai.label_schemas`. "
+    "Please install it with `pip install databricks-agents`."
+)
+
+EXPECTED_FACTS = "expected_facts"
+GUIDELINES = "guidelines"
+EXPECTED_RESPONSE = "expected_response"
+
+
+def create_label_schema(
+    name: str,
+    *,
+    type: Literal["feedback", "expectation"],
+    title: str,
+    input: Union[
+        InputCategorical,
+        InputCategoricalList,
+        InputText,
+        InputTextList,
+        InputNumeric,
+    ],
+    instruction: Optional[str] = None,
+    enable_comment: bool = False,
+    overwrite: bool = False,
+) -> LabelSchema:
+    """Create a new label schema for the review app.
+
+    A label schema defines the type of input that stakeholders will provide when labeling items
+    in the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        name: The name of the label schema. Must be unique across the review app.
+        type: The type of the label schema. Either "feedback" or "expectation".
+        title: The title of the label schema shown to stakeholders.
+        input: The input type of the label schema.
+        instruction: Optional. The instruction shown to stakeholders.
+        enable_comment: Optional. Whether to enable comments for the label schema.
+        overwrite: Optional. Whether to overwrite the existing label schema with the same name.
+
+    Returns:
+        LabelSchema: The created label schema.
+    """
+    try:
+        from databricks.agents import review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    app = review_app.get_review_app()
+    return app.create_label_schema(
+        name=name,
+        type=type,
+        title=title,
+        input=input._to_databricks_input(),
+        instruction=instruction,
+        enable_comment=enable_comment,
+        overwrite=overwrite,
+    )
+
+
+def get_label_schema(name: str) -> LabelSchema:
+    """Get a label schema from the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        name: The name of the label schema to get.
+
+    Returns:
+        LabelSchema: The label schema.
+    """
+    try:
+        from databricks.agents import review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    app = review_app.get_review_app()
+    label_schema = next(
+        (label_schema for label_schema in app.label_schemas if label_schema.name == name),
+        None,
+    )
+    if label_schema is None:
+        raise ValueError(f"Label schema with name `{name}` not found")
+    return LabelSchema._from_databricks_label_schema(label_schema)
+
+
+def delete_label_schema(name: str) -> "ReviewApp":
+    """Delete a label schema from the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        name: The name of the label schema to delete.
+
+    Returns:
+        ReviewApp: The review app.
+    """
+    try:
+        from databricks.agents import review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    app = review_app.get_review_app()
+    return ReviewApp(app.delete_label_schema(name))
+
+
+__all__ = [
+    "EXPECTED_FACTS",
+    "GUIDELINES",
+    "EXPECTED_RESPONSE",
+    "LabelSchemaType",
+    "LabelSchema",
+    "InputCategorical",
+    "InputCategoricalList",
+    "InputNumeric",
+    "InputText",
+    "InputTextList",
+    "create_label_schema",
+    "get_label_schema",
+    "delete_label_schema",
+]

mlflow/genai/label_schemas/label_schemas.py

@@ -0,0 +1,209 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional, TypeVar, Union
+
+from mlflow.genai.utils.enum_utils import StrEnum
+
+if TYPE_CHECKING:
+    from databricks.agents.review_app import label_schemas as _label_schemas
+
+    _InputCategorical = _label_schemas.InputCategorical
+    _InputCategoricalList = _label_schemas.InputCategoricalList
+    _InputNumeric = _label_schemas.InputNumeric
+    _InputText = _label_schemas.InputText
+    _InputTextList = _label_schemas.InputTextList
+    _LabelSchema = _label_schemas.LabelSchema
+
+DatabricksInputType = TypeVar("DatabricksInputType")
+_InputType = TypeVar("_InputType", bound="InputType")
+
+
+class InputType(ABC):
+    """Base class for all input types."""
+
+    @abstractmethod
+    def _to_databricks_input(self) -> DatabricksInputType:
+        """Convert to the internal Databricks input type."""
+
+    @classmethod
+    @abstractmethod
+    def _from_databricks_input(cls, input_obj: DatabricksInputType) -> _InputType:
+        """Create from the internal Databricks input type."""
+
+
+@dataclass
+class InputCategorical(InputType):
+    """A single-select dropdown for collecting assessments from stakeholders.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    options: list[str]
+    """List of available options for the categorical selection."""
+
+    def _to_databricks_input(self) -> "_InputCategorical":
+        """Convert to the internal Databricks input type."""
+        from databricks.agents.review_app import label_schemas as _label_schemas
+
+        return _label_schemas.InputCategorical(options=self.options)
+
+    @classmethod
+    def _from_databricks_input(cls, input_obj: "_InputCategorical") -> "InputCategorical":
+        """Create from the internal Databricks input type."""
+        return cls(options=input_obj.options)
+
+
+@dataclass
+class InputCategoricalList(InputType):
+    """A multi-select dropdown for collecting assessments from stakeholders.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    options: list[str]
+    """List of available options for the multi-select categorical (dropdown)."""
+
+    def _to_databricks_input(self) -> "_InputCategoricalList":
+        """Convert to the internal Databricks input type."""
+        from databricks.agents.review_app import label_schemas as _label_schemas
+
+        return _label_schemas.InputCategoricalList(options=self.options)
+
+    @classmethod
+    def _from_databricks_input(cls, input_obj: "_InputCategoricalList") -> "InputCategoricalList":
+        """Create from the internal Databricks input type."""
+        return cls(options=input_obj.options)
+
+
+@dataclass
+class InputTextList(InputType):
+    """Like `Text`, but allows multiple entries.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    max_length_each: Optional[int] = None
+    """Maximum character length for each individual text entry. None means no limit."""
+
+    max_count: Optional[int] = None
+    """Maximum number of text entries allowed. None means no limit."""
+
+    def _to_databricks_input(self) -> "_InputTextList":
+        """Convert to the internal Databricks input type."""
+        from databricks.agents.review_app import label_schemas as _label_schemas
+
+        return _label_schemas.InputTextList(
+            max_length_each=self.max_length_each, max_count=self.max_count
+        )
+
+    @classmethod
+    def _from_databricks_input(cls, input_obj: "_InputTextList") -> "InputTextList":
+        """Create from the internal Databricks input type."""
+        return cls(max_length_each=input_obj.max_length_each, max_count=input_obj.max_count)
+
+
+@dataclass
+class InputText(InputType):
+    """A free-form text box for collecting assessments from stakeholders.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    max_length: Optional[int] = None
+    """Maximum character length for the text input. None means no limit."""
+
+    def _to_databricks_input(self) -> "_InputText":
+        """Convert to the internal Databricks input type."""
+        from databricks.agents.review_app import label_schemas as _label_schemas
+
+        return _label_schemas.InputText(max_length=self.max_length)
+
+    @classmethod
+    def _from_databricks_input(cls, input_obj: "_InputText") -> "InputText":
+        """Create from the internal Databricks input type."""
+        return cls(max_length=input_obj.max_length)
+
+
+@dataclass
+class InputNumeric(InputType):
+    """A numeric input for collecting assessments from stakeholders.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    min_value: Optional[float] = None
+    """Minimum allowed numeric value. None means no minimum limit."""
+
+    max_value: Optional[float] = None
+    """Maximum allowed numeric value. None means no maximum limit."""
+
+    def _to_databricks_input(self) -> "_InputNumeric":
+        """Convert to the internal Databricks input type."""
+        from databricks.agents.review_app import label_schemas as _label_schemas
+
+        return _label_schemas.InputNumeric(min_value=self.min_value, max_value=self.max_value)
+
+    @classmethod
+    def _from_databricks_input(cls, input_obj: "_InputNumeric") -> "InputNumeric":
+        """Create from the internal Databricks input type."""
+        return cls(min_value=input_obj.min_value, max_value=input_obj.max_value)
+
+
+class LabelSchemaType(StrEnum):
+    """Type of label schema."""
+
+    FEEDBACK = "feedback"
+    EXPECTATION = "expectation"
+
+
+@dataclass(frozen=True)
+class LabelSchema:
+    """A label schema for collecting input from stakeholders.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    name: str
+    """Unique name identifier for the label schema."""
+
+    type: LabelSchemaType
+    """Type of the label schema, either 'feedback' or 'expectation'."""
+
+    title: str
+    """Display title shown to stakeholders in the labeling review UI."""
+
+    input: Union[InputCategorical, InputCategoricalList, InputText, InputTextList, InputNumeric]
+    """
+    Input type specification that defines how stakeholders will provide their assessment
+    (e.g., dropdown, text box, numeric input)
+    """
+    instruction: Optional[str] = None
+    """Optional detailed instructions shown to stakeholders for guidance."""
+
+    enable_comment: bool = False
+    """Whether to enable additional comment functionality for reviewers."""
+
+    @classmethod
+    def _from_databricks_label_schema(cls, schema: "_LabelSchema") -> "LabelSchema":
+        """Convert from the internal Databricks label schema type."""
+
+        return cls(
+            name=schema.name,
+            type=schema.type,
+            title=schema.title,
+            input=schema.input._from_databricks_input(),
+            instruction=schema.instruction,
+            enable_comment=schema.enable_comment,
+        )

mlflow/genai/labeling/__init__.py

@@ -0,0 +1,159 @@
+"""
+Databricks Agent Labeling Python SDK. For more details see Databricks Agent Evaluation:
+<https://docs.databricks.com/en/generative-ai/agent-evaluation/index.html>
+
+The API docs can be found here:
+<https://api-docs.databricks.com/python/databricks-agents/latest/databricks_agent_eval.html#review-app>
+"""
+
+from typing import Any, Optional
+
+from mlflow.genai.labeling.labeling import Agent, LabelingSession, ReviewApp
+
+_ERROR_MSG = (
+    "The `databricks-agents` package is required to use `mlflow.genai.labeling`. "
+    "Please install it with `pip install databricks-agents`."
+)
+
+
+def get_review_app(experiment_id: Optional[str] = None) -> "ReviewApp":
+    """Gets or creates (if it doesn't exist) the review app for the given experiment ID.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        experiment_id: Optional. The experiment ID for which to get the review app. If not provided,
+            the experiment ID is inferred from the current active environment.
+
+    Returns:
+        ReviewApp: The review app.
+    """
+    try:
+        from databricks.agents.review_app import get_review_app as _get_review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    return ReviewApp(_get_review_app(experiment_id))
+
+
+def create_labeling_session(
+    name: str,
+    *,
+    assigned_users: Optional[list[str]] = None,
+    agent: Optional[str] = None,
+    label_schemas: Optional[list[str]] = None,
+    enable_multi_turn_chat: bool = False,
+    custom_inputs: Optional[dict[str, Any]] = None,
+) -> LabelingSession:
+    """Create a new labeling session in the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        name: The name of the labeling session.
+        assigned_users: The users that will be assigned to label items in the session.
+        agent: The agent to be used to generate responses for the items in the session.
+        label_schemas: The label schemas to be used in the session.
+        enable_multi_turn_chat: Whether to enable multi-turn chat labeling for the session.
+        custom_inputs: Optional. Custom inputs to be used in the session.
+
+    Returns:
+        LabelingSession: The created labeling session.
+    """
+    try:
+        from databricks.agents.review_app import get_review_app as _get_review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    return LabelingSession(
+        _get_review_app().create_labeling_session(
+            name=name,
+            assigned_users=assigned_users or [],
+            agent=agent,
+            label_schemas=label_schemas or [],
+            enable_multi_turn_chat=enable_multi_turn_chat,
+            custom_inputs=custom_inputs,
+        )
+    )
+
+
+def get_labeling_sessions() -> list[LabelingSession]:
+    """Get all labeling sessions from the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Returns:
+        list[LabelingSession]: The list of labeling sessions.
+    """
+    try:
+        from databricks.agents.review_app import get_review_app as _get_review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    return [LabelingSession(session) for session in _get_review_app().get_labeling_sessions()]
+
+
+def get_labeling_session(run_id: str) -> LabelingSession:
+    """Get a labeling session from the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        run_id: The mlflow run ID of the labeling session to get.
+
+    Returns:
+        LabelingSession: The labeling session.
+    """
+    labeling_sessions = get_labeling_sessions()
+    labeling_session = next(
+        (
+            labeling_session
+            for labeling_session in labeling_sessions
+            if labeling_session.mlflow_run_id == run_id
+        ),
+        None,
+    )
+    if labeling_session is None:
+        raise ValueError(f"Labeling session with run_id `{run_id}` not found")
+    return LabelingSession(labeling_session)
+
+
+def delete_labeling_session(labeling_session: LabelingSession) -> "ReviewApp":
+    """Delete a labeling session from the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+
+    Args:
+        labeling_session: The labeling session to delete.
+
+    Returns:
+        ReviewApp: The review app.
+    """
+    try:
+        from databricks.agents.review_app import get_review_app as _get_review_app
+    except ImportError:
+        raise ImportError(_ERROR_MSG) from None
+
+    return ReviewApp(_get_review_app().delete_labeling_session(labeling_session._session))
+
+
+__all__ = [
+    "Agent",
+    "LabelingSession",
+    "ReviewApp",
+    "get_review_app",
+    "create_labeling_session",
+    "get_labeling_sessions",
+    "get_labeling_session",
+    "delete_labeling_session",
+]