genesis-flow 1.0.0__py3-none-any.whl

mlflow/genai/optimize/optimizers/dspy_mipro_optimizer.py

@@ -0,0 +1,221 @@
+import contextlib
+import io
+import logging
+import math
+import os
+from typing import TYPE_CHECKING, Optional
+
+from mlflow.entities.model_registry import Prompt
+from mlflow.exceptions import MlflowException
+from mlflow.genai.optimize.optimizers.dspy_optimizer import _DSPyOptimizer
+from mlflow.genai.optimize.types import OBJECTIVE_FN, LLMParams
+from mlflow.genai.prompts import register_prompt
+from mlflow.genai.scorers import Scorer
+from mlflow.tracking._model_registry.fluent import active_run
+from mlflow.tracking.fluent import log_metric, log_param
+
+if TYPE_CHECKING:
+    import dspy
+    import pandas as pd
+
+
+_logger = logging.getLogger(__name__)
+
+
+class _DSPyMIPROv2Optimizer(_DSPyOptimizer):
+    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:
+        import dspy
+
+        from mlflow.genai.optimize.optimizers.utils.dspy_mipro_callback import _DSPyMIPROv2Callback
+        from mlflow.genai.optimize.optimizers.utils.dspy_mipro_utils import format_optimized_prompt
+
+        _logger.info(
+            f"🎯 Starting prompt optimization for: {prompt.uri}\n"
+            f"⏱️ This may take several minutes or longer depending on dataset size...\n"
+            f"📊 Training with {len(train_data)} examples."
+        )
+
+        input_fields = self._get_input_fields(train_data)
+        self._validate_input_fields(input_fields, prompt)
+        output_fields = self._get_output_fields(train_data)
+
+        teacher_settings = {}
+        if self.optimizer_config.optimizer_llm:
+            teacher_lm = dspy.LM(
+                model=self.optimizer_config.optimizer_llm.model_name,
+                temperature=self.optimizer_config.optimizer_llm.temperature,
+                api_base=self.optimizer_config.optimizer_llm.base_uri,
+            )
+            teacher_settings["lm"] = teacher_lm
+
+        lm = dspy.LM(
+            model=target_llm_params.model_name,
+            temperature=target_llm_params.temperature,
+            api_base=target_llm_params.base_uri,
+        )
+
+        instructions = self._extract_instructions(prompt.template, teacher_settings.get("lm", lm))
+
+        signature = dspy.make_signature(
+            {
+                **{key: (_type, dspy.InputField()) for key, _type in input_fields.items()},
+                **{key: (_type, dspy.OutputField()) for key, _type in output_fields.items()},
+            },
+            instructions,
+        )
+
+        # Define main student program
+        program = dspy.Predict(signature)
+
+        train_data = self._convert_to_dspy_dataset(train_data)
+        eval_data = self._convert_to_dspy_dataset(eval_data) if eval_data is not None else None
+
+        num_candidates = self.optimizer_config.num_instruction_candidates
+        optimizer = dspy.MIPROv2(
+            metric=self._convert_to_dspy_metric(input_fields, output_fields, scorers, objective),
+            max_bootstrapped_demos=self.optimizer_config.max_few_show_examples,
+            num_candidates=num_candidates,
+            num_threads=self.optimizer_config.num_threads,
+            teacher_settings=teacher_settings,
+            auto=None,
+        )
+
+        adapter = dspy.JSONAdapter()
+        callbacks = (
+            [
+                _DSPyMIPROv2Callback(prompt.name, input_fields),
+            ]
+            if self.optimizer_config.autolog
+            else []
+        )
+        with dspy.context(lm=lm, adapter=adapter, callbacks=callbacks):
+            with self._maybe_suppress_stdout_stderr():
+                optimized_program = optimizer.compile(
+                    program,
+                    trainset=train_data,
+                    valset=eval_data,
+                    num_trials=self._get_num_trials(num_candidates),
+                    minibatch_size=self._get_minibatch_size(train_data, eval_data),
+                    requires_permission_to_run=False,
+                )
+
+            template = format_optimized_prompt(
+                program=optimized_program,
+                input_fields=input_fields,
+            )
+
+        self._display_optimization_result(optimized_program)
+        final_score = getattr(optimized_program, "score", None)
+        optimized_prompt = register_prompt(
+            name=prompt.name,
+            template=template,
+            tags={
+                "overall_eval_score": str(final_score),
+            },
+        )
+
+        if self.optimizer_config.autolog:
+            self._log_optimization_result(final_score, optimized_prompt)
+
+        return optimized_prompt
+
+    def _get_num_trials(self, num_candidates: int) -> int:
+        # MAX(2*log(num_candidates), 3/2*num_candidates)
+        return int(max(2 * math.log2(num_candidates), 1.5 * num_candidates))
+
+    def _get_minibatch_size(
+        self,
+        train_data: "pd.DataFrame",
+        eval_data: Optional["pd.DataFrame"] = None,
+    ) -> int:
+        # The default minibatch size is 35 in MIPROv2.
+        if eval_data is not None:
+            return min(35, len(eval_data) // 2)
+        return min(35, len(train_data) // 2)
+
+    def _validate_input_fields(self, input_fields: dict[str, type], prompt: Prompt) -> None:
+        if missing_fields := set(prompt.variables) - set(input_fields.keys()):
+            raise MlflowException(
+                f"Validation failed. Missing prompt variables in dataset: {missing_fields}. "
+                "Please ensure your dataset contains columns for all prompt variables."
+            )
+
+    def _extract_instructions(self, template: str, lm: "dspy.LM") -> str:
+        import dspy
+
+        extractor = dspy.Predict(
+            dspy.make_signature(
+                {
+                    "prompt": (str, dspy.InputField()),
+                    "instruction": (str, dspy.OutputField()),
+                },
+                "Extract the core instructions from the prompt "
+                "to use as the system message for the LLM.",
+            )
+        )
+
+        with dspy.context(lm=lm):
+            return extractor(prompt=template).instruction
+
+    @contextlib.contextmanager
+    def _maybe_suppress_stdout_stderr(self):
+        """Context manager for redirecting stdout/stderr based on verbose setting.
+        If verbose is False, redirects output to devnull or StringIO.
+        If verbose is True, doesn't redirect output.
+        """
+        if not self.optimizer_config.verbose:
+            try:
+                output_sink = open(os.devnull, "w")  # noqa: SIM115
+            except (OSError, IOError):
+                output_sink = io.StringIO()
+
+            with output_sink:
+                with (
+                    contextlib.redirect_stdout(output_sink),
+                    contextlib.redirect_stderr(output_sink),
+                ):
+                    yield
+        else:
+            yield
+
+    def _display_optimization_result(self, program: "dspy.Predict"):
+        score = getattr(program, "score", None)
+        if score is None:
+            return
+        # In DSPy < 2.6.17, trial_logs contains initial score in key=-1.
+        trial_logs = getattr(program, "trial_logs", {})
+        if 1 in trial_logs:
+            initial_score = trial_logs[1].get("full_eval_score")
+        elif -1 in trial_logs:
+            initial_score = trial_logs[-1].get("full_eval_score")
+        else:
+            initial_score = None
+        if initial_score is not None:
+            if initial_score == score:
+                _logger.info(f"Optimization complete! Score remained stable at: {score}.")
+                return
+            else:
+                _logger.info(
+                    f"🎉 Optimization complete! "
+                    f"Initial score: {initial_score}. "
+                    f"Final score: {score}."
+                )
+
+    def _log_optimization_result(self, final_score: Optional[float], optimized_prompt: Prompt):
+        if not active_run():
+            return
+
+        if final_score:
+            log_metric(
+                "final_eval_score",
+                final_score,
+            )
+        log_param("optimized_prompt_uri", optimized_prompt.uri)

mlflow/genai/optimize/optimizers/dspy_optimizer.py

@@ -0,0 +1,91 @@
+import importlib.metadata
+import importlib.util
+import inspect
+from typing import TYPE_CHECKING, Callable, Optional
+
+from packaging.version import Version
+
+from mlflow.exceptions import MlflowException
+from mlflow.genai.optimize.optimizers import _BaseOptimizer
+from mlflow.genai.optimize.types import OBJECTIVE_FN, OptimizerConfig
+from mlflow.genai.optimize.util import infer_type_from_value
+from mlflow.genai.scorers import Scorer
+
+if TYPE_CHECKING:
+    import dspy
+    import pandas as pd
+
+
+class _DSPyOptimizer(_BaseOptimizer):
+    def __init__(self, optimizer_config: OptimizerConfig):
+        super().__init__(optimizer_config)
+
+        if not importlib.util.find_spec("dspy"):
+            raise ImportError("dspy is not installed. Please install it with `pip install dspy`.")
+
+        dspy_version = importlib.metadata.version("dspy")
+        if Version(dspy_version) < Version("2.6.0"):
+            raise MlflowException(
+                f"Current dspy version {dspy_version} is unsupported. "
+                "Please upgrade to version >= 2.6.0"
+            )
+
+    def _get_input_fields(self, train_data: "pd.DataFrame") -> dict[str, type]:
+        if "inputs" in train_data.columns:
+            sample_input = train_data["inputs"].values[0]
+            return {k: infer_type_from_value(v) for k, v in sample_input.items()}
+        return {}
+
+    def _get_output_fields(self, train_data: "pd.DataFrame") -> dict[str, type]:
+        if "expectations" in train_data.columns:
+            sample_output = train_data["expectations"].values[0]
+            return {k: infer_type_from_value(v) for k, v in sample_output.items()}
+        return {}
+
+    def _convert_to_dspy_dataset(self, data: "pd.DataFrame") -> list["dspy.Example"]:
+        import dspy
+
+        examples = []
+        for _, row in data.iterrows():
+            expectations = row["expectations"] if "expectations" in row else {}
+            examples.append(
+                dspy.Example(**row["inputs"], **expectations).with_inputs(*row["inputs"].keys())
+            )
+        return examples
+
+    def _convert_to_dspy_metric(
+        self,
+        input_fields: dict[str, type],
+        output_fields: dict[str, type],
+        scorers: list[Scorer],
+        objective: Optional[OBJECTIVE_FN] = None,
+    ) -> Callable[["dspy.Example"], float]:
+        def metric(example: "dspy.Example", pred: "dspy.Example", trace=None) -> float:
+            scores = {}
+            inputs = {key: example.get(key) for key in input_fields.keys()}
+            expectations = {key: example.get(key) for key in output_fields.keys()}
+            outputs = {key: pred.get(key) for key in output_fields.keys()}
+
+            for scorer in scorers:
+                kwargs = {"inputs": inputs, "outputs": outputs, "expectations": expectations}
+                signature = inspect.signature(scorer)
+                kwargs = {
+                    key: value for key, value in kwargs.items() if key in signature.parameters
+                }
+                scores[scorer.name] = scorer(**kwargs)
+            if objective is not None:
+                return objective(scores)
+            elif all(isinstance(score, (int, float, bool)) for score in scores.values()):
+                # Use total score by default if no objective is provided
+                return sum(scores.values())
+            else:
+                non_numerical_scorers = [
+                    k for k, v in scores.items() if not isinstance(v, (int, float, bool))
+                ]
+                raise MlflowException(
+                    f"Scorer [{','.join(non_numerical_scorers)}] return a string, Assessment or a "
+                    "list of Assessment. Please provide `objective` function to aggregate "
+                    "non-numerical values into a single value for optimization."
+                )
+
+        return metric

mlflow/genai/optimize/optimizers/utils/dspy_mipro_callback.py

@@ -0,0 +1,76 @@
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Optional
+
+from dspy import Prediction
+from dspy.utils.callback import BaseCallback
+
+import mlflow
+from mlflow.genai.optimize.optimizers.utils.dspy_mipro_utils import format_optimized_prompt
+
+_FULL_EVAL_NAME = "eval_full"
+
+if TYPE_CHECKING:
+    import dspy
+
+
+class _DSPyMIPROv2Callback(BaseCallback):
+    def __init__(self, prompt_name: str, input_fields: dict[str, type]):
+        self.prompt_name = prompt_name
+        self.input_fields = input_fields
+        # call_id: (key, step, program)
+        self._call_id_to_values: dict[str, tuple[str, int, "dspy.Predict"]] = {}
+        self._evaluation_counter = defaultdict(int)
+        self._best_score = None
+
+    def on_evaluate_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        key = inputs.get("callback_metadata", {}).get("metric_key", "eval")
+        step = self._evaluation_counter[key]
+        self._evaluation_counter[key] += 1
+        self._call_id_to_values[call_id] = (key, step, inputs.get("program"))
+
+    def on_evaluate_end(
+        self,
+        call_id: str,
+        outputs: Any,
+        exception: Optional[Exception] = None,
+    ):
+        if exception:
+            return
+
+        if call_id not in self._call_id_to_values:
+            return
+        key, step, program = self._call_id_to_values.pop(call_id)
+
+        # Only log the full evaluation result
+        if key != _FULL_EVAL_NAME:
+            return
+
+        if isinstance(outputs, float):
+            score = outputs
+        elif isinstance(outputs, tuple):
+            score = outputs[0]
+        elif isinstance(outputs, Prediction):
+            score = float(outputs)
+        else:
+            return
+
+        if self._best_score is None:
+            # This is the first evaluation with initial prompt,
+            # we don't register this prompt
+            self._best_score = score
+        elif score > self._best_score:
+            # When best score is updated, register the new prompt
+            self._best_score = score
+            template = format_optimized_prompt(program, self.input_fields)
+            mlflow.genai.register_prompt(
+                name=self.prompt_name,
+                template=template,
+                tags={"overall_eval_score": score, key: step},
+            )
+
+        if mlflow.active_run() is not None:
+            mlflow.log_metric(
+                key,
+                score,
+                step=step,
+            )

mlflow/genai/optimize/optimizers/utils/dspy_mipro_utils.py

@@ -0,0 +1,18 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import dspy
+
+
+def format_optimized_prompt(program: "dspy.Predict", input_fields: dict[str, type]) -> str:
+    import dspy
+
+    messages = dspy.settings.adapter.format(
+        signature=program.signature,
+        demos=program.demos,
+        inputs={key: "{{" + key + "}}" for key in input_fields.keys()},
+    )
+
+    return "\n\n".join(
+        [f"<{message['role']}>\n{message['content']}\n</{message['role']}>" for message in messages]
+    )

mlflow/genai/optimize/types.py

@@ -0,0 +1,75 @@
+import multiprocessing
+from dataclasses import dataclass, field
+from typing import Callable, Optional, Union
+
+from mlflow.entities import Feedback
+from mlflow.entities.model_registry import Prompt
+from mlflow.utils.annotations import experimental
+
+OBJECTIVE_FN = Callable[[dict[str, Union[bool, float, str, Feedback, list[Feedback]]]], float]
+
+
+@experimental(version="3.0.0")
+@dataclass
+class PromptOptimizationResult:
+    """
+    Result of the :py:func:`mlflow.genai.optimize_prompt()` API.
+
+    Args:
+        prompt: A prompt entity containing the optimized template.
+    """
+
+    prompt: Prompt
+
+
+@experimental(version="3.0.0")
+@dataclass
+class LLMParams:
+    """
+    Parameters for configuring a LLM Model.
+
+    Args:
+        model_name: Name of the model in the format `<provider>/<model name>`.
+            For example, "openai/gpt-4" or "anthropic/claude-4".
+        base_uri: Optional base URI for the API endpoint. If not provided,
+            the default endpoint for the provider will be used.
+        temperature: Optional sampling temperature for the model's outputs.
+            Higher values (e.g., 0.8) make the output more random,
+            while lower values (e.g., 0.2) make it more deterministic.
+    """
+
+    model_name: str
+    base_uri: Optional[str] = None
+    temperature: Optional[float] = None
+
+
+@experimental(version="3.0.0")
+@dataclass
+class OptimizerConfig:
+    """
+    Configuration for prompt optimization.
+
+    Args:
+        num_instruction_candidates: Number of candidate instructions to generate
+            during each optimization iteration. Higher values may lead to better
+            results but increase optimization time. Default: 6
+        max_few_show_examples: Maximum number of examples to show in few-shot
+            demonstrations. Default: 6
+        num_threads: Number of threads to use for parallel optimization.
+            Default: (number of CPU cores * 2 + 1)
+        optimizer_llm: Optional LLM parameters for the teacher model. If not provided,
+            the target LLM will be used as the teacher.
+        algorithm: The optimization algorithm to use. Default: "DSPy/MIPROv2"
+        verbose: Whether to show optimizer logs during optimization. Default: False
+        autolog: Whether to log the optimization parameters, datasets and metrics.
+            If set to True, a MLflow run is automatically created to store them.
+            Default: False
+    """
+
+    num_instruction_candidates: int = 6
+    max_few_show_examples: int = 6
+    num_threads: int = field(default_factory=lambda: (multiprocessing.cpu_count() or 1) * 2 + 1)
+    optimizer_llm: Optional[LLMParams] = None
+    algorithm: str = "DSPy/MIPROv2"
+    verbose: bool = False
+    autolog: bool = False

mlflow/genai/optimize/util.py

@@ -0,0 +1,30 @@
+from typing import Any, Union
+
+from pydantic import BaseModel, create_model
+
+
+def infer_type_from_value(value: Any, model_name: str = "Output") -> type:
+    """
+    Infer the type from the value.
+    Only supports primitive types, lists, and dict and Pydantic models.
+    """
+    if value is None:
+        return type(None)
+    elif isinstance(value, (bool, int, float, str)):
+        return type(value)
+    elif isinstance(value, list):
+        if not value:
+            return list[Any]
+        element_types = set()
+        for item in value:
+            element_types.add(infer_type_from_value(item))
+        if len(element_types) == 1:
+            return list[element_types.pop()]
+        else:
+            return list[Union[tuple(element_types)]]
+    elif isinstance(value, dict):
+        fields = {k: (infer_type_from_value(v, model_name=k), ...) for k, v in value.items()}
+        return create_model(model_name, **fields)
+    elif isinstance(value, BaseModel):
+        return type(value)
+    return Any