genesis-flow 1.0.0__py3-none-any.whl

mlflow/genai/labeling/labeling.py

@@ -0,0 +1,250 @@
+from typing import TYPE_CHECKING, Any, Iterable, Optional, Union
+
+from mlflow.entities import Trace
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from databricks.agents.review_app import (
+        LabelingSession as _LabelingSession,
+    )
+    from databricks.agents.review_app import (
+        LabelSchema as _LabelSchema,
+    )
+    from databricks.agents.review_app import (
+        ReviewApp as _ReviewApp,
+    )
+    from databricks.agents.review_app.labeling import Agent as _Agent
+
+
+class Agent:
+    """The agent configuration, used for generating responses in the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    def __init__(self, agent: "_Agent"):
+        self._agent = agent
+
+    @property
+    def agent_name(self) -> str:
+        """The name of the agent."""
+        return self._agent.agent_name
+
+    @property
+    def model_serving_endpoint(self) -> str:
+        """The model serving endpoint used by the agent."""
+        return self._agent.model_serving_endpoint
+
+
+class LabelingSession:
+    """A session for labeling items in the review app.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    def __init__(self, session: "_LabelingSession"):
+        self._session = session
+
+    @property
+    def name(self) -> str:
+        """The name of the labeling session."""
+        return self._session.name
+
+    @property
+    def assigned_users(self) -> list[str]:
+        """The users assigned to label items in the session."""
+        return self._session.assigned_users
+
+    @property
+    def agent(self) -> Optional[str]:
+        """The agent used to generate responses for the items in the session."""
+        return self._session.agent
+
+    @property
+    def label_schemas(self) -> list[str]:
+        """The label schemas used in the session."""
+        return self._session.label_schemas
+
+    @property
+    def labeling_session_id(self) -> str:
+        """The unique identifier of the labeling session."""
+        return self._session.labeling_session_id
+
+    @property
+    def mlflow_run_id(self) -> str:
+        """The MLflow run ID associated with the session."""
+        return self._session.mlflow_run_id
+
+    @property
+    def review_app_id(self) -> str:
+        """The review app ID associated with the session."""
+        return self._session.review_app_id
+
+    @property
+    def experiment_id(self) -> str:
+        """The experiment ID associated with the session."""
+        return self._session.experiment_id
+
+    @property
+    def url(self) -> str:
+        """The URL of the labeling session in the review app."""
+        return self._session.url
+
+    @property
+    def enable_multi_turn_chat(self) -> bool:
+        """Whether multi-turn chat is enabled for the session."""
+        return self._session.enable_multi_turn_chat
+
+    @property
+    def custom_inputs(self) -> Optional[dict[str, Any]]:
+        """Custom inputs used in the session."""
+        return self._session.custom_inputs
+
+    def add_dataset(
+        self, dataset_name: str, record_ids: Optional[list[str]] = None
+    ) -> "LabelingSession":
+        """Add a dataset to the labeling session.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            dataset_name: The name of the dataset.
+            record_ids: Optional. The individual record ids to be added to the session. If not
+                provided, all records in the dataset will be added.
+
+        Returns:
+            LabelingSession: The updated labeling session.
+        """
+        return LabelingSession(self._session.add_dataset(dataset_name, record_ids))
+
+    def add_traces(
+        self,
+        traces: Union[Iterable[Trace], Iterable[str], "pd.DataFrame"],
+    ) -> "LabelingSession":
+        """Add traces to the labeling session.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            traces: Can be either:
+                a) a pandas DataFrame with a 'trace' column. The 'trace' column should contain
+                either `mlflow.entities.Trace` objects or their json string representations.
+                b) an iterable of `mlflow.entities.Trace` objects.
+                c) an iterable of json string representations of `mlflow.entities.Trace` objects.
+
+        Returns:
+            LabelingSession: The updated labeling session.
+        """
+        return LabelingSession(self._session.add_traces(traces))
+
+    def sync(self, to_dataset: str) -> None:
+        """Sync the traces and expectations from the labeling session to a dataset.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            to_dataset: The name of the dataset to sync traces and expectations to.
+        """
+        self._session.sync_expectations(to_dataset)
+
+    def set_assigned_users(self, assigned_users: list[str]) -> "LabelingSession":
+        """Set the assigned users for the labeling session.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            assigned_users: The list of users to assign to the session.
+
+        Returns:
+            LabelingSession: The updated labeling session.
+        """
+        return LabelingSession(self._session.set_assigned_users(assigned_users))
+
+
+class ReviewApp:
+    """A review app is used to collect feedback from stakeholders for a given experiment.
+
+    .. note::
+        This functionality is only available in Databricks. Please run
+        `pip install mlflow[databricks]` to use it.
+    """
+
+    def __init__(self, app: "_ReviewApp"):
+        self._app = app
+
+    @property
+    def review_app_id(self) -> str:
+        """The ID of the review app."""
+        return self._app.review_app_id
+
+    @property
+    def experiment_id(self) -> str:
+        """The ID of the experiment."""
+        return self._app.experiment_id
+
+    @property
+    def url(self) -> str:
+        """The URL of the review app for stakeholders to provide feedback."""
+        return self._app.url
+
+    @property
+    def agents(self) -> list[Agent]:
+        """The agents to be used to generate responses."""
+        return [Agent(agent) for agent in self._app.agents]
+
+    @property
+    def label_schemas(self) -> list["_LabelSchema"]:
+        """The label schemas to be used in the review app."""
+        return self._app.label_schemas
+
+    def add_agent(
+        self, *, agent_name: str, model_serving_endpoint: str, overwrite: bool = False
+    ) -> "ReviewApp":
+        """Add an agent to the review app to be used to generate responses.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            agent_name: The name of the agent.
+            model_serving_endpoint: The model serving endpoint to be used by the agent.
+            overwrite: Whether to overwrite an existing agent with the same name.
+
+        Returns:
+            ReviewApp: The updated review app.
+        """
+        return ReviewApp(
+            self._app.add_agent(
+                agent_name=agent_name,
+                model_serving_endpoint=model_serving_endpoint,
+                overwrite=overwrite,
+            )
+        )
+
+    def remove_agent(self, agent_name: str) -> "ReviewApp":
+        """Remove an agent from the review app.
+
+        .. note::
+            This functionality is only available in Databricks. Please run
+            `pip install mlflow[databricks]` to use it.
+
+        Args:
+            agent_name: The name of the agent to remove.
+
+        Returns:
+            ReviewApp: The updated review app.
+        """
+        return ReviewApp(self._app.remove_agent(agent_name))

mlflow/genai/optimize/__init__.py

@@ -0,0 +1,13 @@
+from mlflow.genai.optimize.base import optimize_prompt
+from mlflow.genai.optimize.types import (
+    LLMParams,
+    OptimizerConfig,
+    PromptOptimizationResult,
+)
+
+__all__ = [
+    "optimize_prompt",
+    "OptimizerConfig",
+    "LLMParams",
+    "PromptOptimizationResult",
+]

mlflow/genai/optimize/base.py

@@ -0,0 +1,198 @@
+import inspect
+import logging
+from contextlib import contextmanager
+from dataclasses import asdict
+from typing import TYPE_CHECKING, Optional, Union
+
+from mlflow.entities.model_registry import PromptVersion
+from mlflow.exceptions import MlflowException
+from mlflow.genai.evaluation.utils import (
+    _convert_eval_set_to_df,
+)
+from mlflow.genai.optimize.optimizers import _BaseOptimizer, _DSPyMIPROv2Optimizer
+from mlflow.genai.optimize.types import (
+    OBJECTIVE_FN,
+    LLMParams,
+    OptimizerConfig,
+    PromptOptimizationResult,
+)
+from mlflow.genai.prompts import load_prompt
+from mlflow.genai.scorers import Scorer
+from mlflow.tracking.fluent import log_params, log_table, start_run
+from mlflow.utils.annotations import experimental
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from genai.evaluation.utils import EvaluationDatasetTypes
+
+_ALGORITHMS = {"DSPy/MIPROv2": _DSPyMIPROv2Optimizer}
+
+_logger = logging.getLogger(__name__)
+
+
+@experimental(version="3.0.0")
+def optimize_prompt(
+    *,
+    target_llm_params: LLMParams,
+    prompt: Union[str, PromptVersion],
+    train_data: "EvaluationDatasetTypes",
+    scorers: list[Scorer],
+    objective: Optional[OBJECTIVE_FN] = None,
+    eval_data: Optional["EvaluationDatasetTypes"] = None,
+    optimizer_config: Optional[OptimizerConfig] = None,
+) -> PromptOptimizationResult:
+    """
+    Optimize a LLM prompt using the given dataset and evaluation metrics.
+    The optimized prompt template is automatically registered as a new version of the
+    original prompt and included in the result.
+    Currently, this API only supports DSPy's MIPROv2 optimizer.
+
+    Args:
+        target_llm_params: Parameters for the the LLM that prompt is optimized for.
+            The model name must be specified in the format `<provider>/<model>`.
+        prompt: The URI or Prompt object of the MLflow prompt to optimize.
+            The optimized prompt is registered as a new version of the prompt.
+        train_data: Training dataset used for optimization.
+            The data must be one of the following formats:
+
+            * An EvaluationDataset entity
+            * Pandas DataFrame
+            * Spark DataFrame
+            * List of dictionaries
+
+            The dataset must include the following columns:
+
+            - inputs: A column containing single inputs in dict format.
+              Each input should contain keys matching the variables in the prompt template.
+            - expectations: A column containing a dictionary
+              of ground truths for individual output fields.
+
+        scorers: List of scorers that evaluate the inputs, outputs and expectations.
+            Note: Trace input is not supported for optimization. Use inputs, outputs and
+            expectations for optimization. Also, pass the `objective` argument
+            when using scorers with string or :class:`~mlflow.entities.Feedback` type outputs.
+        objective: A callable that computes the overall performance metric from individual
+            assessments. Takes a dict mapping assessment names to assessment scores and
+            returns a float value (greater is better).
+        eval_data: Evaluation dataset with the same format as train_data. If not provided,
+            train_data will be automatically split into training and evaluation sets.
+        optimizer_config: Configuration parameters for the optimizer.
+
+    Returns:
+        PromptOptimizationResult: The optimization result including the optimized prompt.
+
+    Example:
+
+        .. code-block:: python
+
+            import os
+            import mlflow
+            from typing import Any
+            from mlflow.genai.scorers import scorer
+            from mlflow.genai.optimize import OptimizerConfig, LLMParams
+
+            os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"
+
+
+            @scorer
+            def exact_match(expectations: dict[str, Any], outputs: dict[str, Any]) -> bool:
+                return expectations == outputs
+
+
+            prompt = mlflow.genai.register_prompt(
+                name="qa",
+                template="Answer the following question: {{question}}",
+            )
+
+            result = mlflow.genai.optimize_prompt(
+                target_llm_params=LLMParams(model_name="openai/gpt-4.1-nano"),
+                train_data=[
+                    {"inputs": {"question": f"{i}+1"}, "expectations": {"answer": f"{i + 1}"}}
+                    for i in range(100)
+                ],
+                scorers=[exact_match],
+                prompt=prompt.uri,
+                optimizer_config=OptimizerConfig(num_instruction_candidates=5),
+            )
+
+            print(result.prompt.template)
+    """
+    if optimizer_config is None:
+        optimizer_config = OptimizerConfig()
+    optimzer = _select_optimizer(optimizer_config)
+    _validate_scorers(scorers)
+
+    train_data = _convert_eval_set_to_df(train_data)
+    if eval_data is not None:
+        eval_data = _convert_eval_set_to_df(eval_data)
+
+    if isinstance(prompt, str):
+        prompt: PromptVersion = load_prompt(prompt)
+
+    with _maybe_start_autolog(optimizer_config, train_data, eval_data, prompt, target_llm_params):
+        optimized_prompt = optimzer.optimize(
+            prompt=prompt,
+            target_llm_params=target_llm_params,
+            train_data=train_data,
+            scorers=scorers,
+            objective=objective,
+            eval_data=eval_data,
+        )
+
+    return PromptOptimizationResult(prompt=optimized_prompt)
+
+
+def _select_optimizer(optimizer_config: OptimizerConfig) -> _BaseOptimizer:
+    if optimizer_config.algorithm not in _ALGORITHMS:
+        raise ValueError(
+            f"Unsupported algorithm: '{optimizer_config.algorithm}'. "
+            f"Available algorithms: {list(_ALGORITHMS.keys())}. "
+            "Please choose from the supported algorithms above."
+        )
+
+    return _ALGORITHMS[optimizer_config.algorithm](optimizer_config)
+
+
+def _validate_scorers(scorers: list[Scorer]) -> None:
+    for scorer in scorers:
+        if not isinstance(scorer, Scorer):
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid scorer: {scorer}. Please use @scorer decorator "
+                "on your function or inherit from Scorer class."
+            )
+
+        signature = inspect.signature(scorer)
+        if "trace" in signature.parameters:
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid scorer parameter: {scorer} contains 'trace' parameter. "
+                "Optimization scorers can only use: inputs, outputs, or expectations. "
+                "Remove 'trace' parameter from your scorer function."
+            )
+
+
+@contextmanager
+def _maybe_start_autolog(
+    optimizer_config: OptimizerConfig,
+    train_data: "pd.DataFrame",
+    eval_data: Optional["pd.DataFrame"],
+    prompt: PromptVersion,
+    target_llm_params: LLMParams,
+):
+    if optimizer_config.autolog:
+        with start_run() as run:
+            _logger.info(
+                f"🚀 MLflow Run `{run.info.run_id}` started for prompt optimization! "
+                "Watch the run to track the optimization progress."
+            )
+            log_table(train_data, "train_data.json")
+            if eval_data is not None:
+                log_table(eval_data, "eval_data.json")
+            params = {
+                "prompt_uri": prompt.uri,
+                **{f"target_llm_params.{k}": v for k, v in asdict(target_llm_params).items()},
+                **{f"optimizer_config.{k}": v for k, v in asdict(optimizer_config).items()},
+            }
+            log_params(params)
+            yield
+    else:
+        yield

mlflow/genai/optimize/optimizers/__init__.py

@@ -0,0 +1,4 @@
+from mlflow.genai.optimize.optimizers.base_optimizer import _BaseOptimizer
+from mlflow.genai.optimize.optimizers.dspy_mipro_optimizer import _DSPyMIPROv2Optimizer
+
+__all__ = ["_BaseOptimizer", "_DSPyMIPROv2Optimizer"]

mlflow/genai/optimize/optimizers/base_optimizer.py

@@ -0,0 +1,38 @@
+import abc
+from typing import TYPE_CHECKING, Optional
+
+from mlflow.entities.model_registry import Prompt
+from mlflow.genai.optimize.types import OBJECTIVE_FN, LLMParams, OptimizerConfig
+from mlflow.genai.scorers import Scorer
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+class _BaseOptimizer(abc.ABC):
+    def __init__(self, optimizer_config: OptimizerConfig):
+        self.optimizer_config = optimizer_config
+
+    @abc.abstractmethod
+    def optimize(
+        self,
+        prompt: Prompt,
+        target_llm_params: LLMParams,
+        train_data: "pd.DataFrame",
+        scorers: list[Scorer],
+        objective: Optional[OBJECTIVE_FN] = None,
+        eval_data: Optional["pd.DataFrame"] = None,
+    ) -> Prompt:
+        """Optimize the given prompt using the specified configuration.
+
+        Args:
+            prompt: The prompt to optimize.
+            target_llm_params: Parameters for the agent LLM.
+            train_data: Training dataset for optimization.
+            scorers: List of scorers to evaluate the optimization.
+            objective: Optional function to compute overall performance metric.
+            eval_data: Optional evaluation dataset.
+
+        Returns:
+            The optimized prompt registered in the prompt registry as a new version.
+        """