star

venv/lib/python3.9/site-packages/mlflow/dspy/__init__.py

@@ -0,0 +1,18 @@
+from mlflow.dspy.autolog import autolog
+from mlflow.dspy.load import _load_pyfunc, load_model
+from mlflow.dspy.save import (
+    get_default_conda_env,
+    get_default_pip_requirements,
+    log_model,
+    save_model,
+)
+
+__all__ = [
+    "autolog",
+    "get_default_conda_env",
+    "get_default_pip_requirements",
+    "save_model",
+    "log_model",
+    "load_model",
+    "_load_pyfunc",
+]

venv/lib/python3.9/site-packages/mlflow/dspy/autolog.py

@@ -0,0 +1,198 @@
+import importlib
+
+from packaging.version import Version
+
+import mlflow
+from mlflow.dspy.save import FLAVOR_NAME
+from mlflow.tracing.provider import trace_disabled
+from mlflow.tracing.utils import construct_full_inputs
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import (
+    autologging_integration,
+    get_autologging_config,
+    safe_patch,
+)
+
+
+@experimental
+def autolog(
+    log_traces: bool = True,
+    log_traces_from_compile: bool = False,
+    log_traces_from_eval: bool = True,
+    log_compiles: bool = False,
+    log_evals: bool = False,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from DSPy to MLflow. Currently, the
+    MLflow DSPy flavor only supports autologging for tracing.
+
+    Args:
+        log_traces: If ``True``, traces are logged for DSPy models by using. If ``False``,
+            no traces are collected during inference. Default to ``True``.
+        log_traces_from_compile: If ``True``, traces are logged when compiling (optimizing)
+            DSPy programs. If ``False``, traces are only logged from normal model inference and
+            disabled when compiling. Default to ``False``.
+        log_traces_from_eval: If ``True``, traces are logged for DSPy models when running DSPy's
+            `built-in evaluator <https://dspy.ai/learn/evaluation/metrics/#evaluation>`_.
+            If ``False``, traces are only logged from normal model inference and disabled when
+            running the evaluator. Default to ``True``.
+        log_compiles: If ``True``, information about the optimization process is logged when
+            `Teleprompter.compile()` is called.
+        log_evals: If ``True``, information about the evaluation call is logged when
+            `Evaluate.__call__()` is called.
+        disable: If ``True``, disables the DSPy autologging integration. If ``False``,
+            enables the DSPy autologging integration.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during DSPy
+            autologging. If ``False``, show all events and warnings.
+    """
+    # NB: The @autologging_integration annotation is used for adding shared logic. However, one
+    # caveat is that the wrapped function is NOT executed when disable=True is passed. This prevents
+    # us from running cleaning up logging when autologging is turned off. To workaround this, we
+    # annotate _autolog() instead of this entrypoint, and define the cleanup logic outside it.
+    # This needs to be called before doing any safe-patching (otherwise safe-patch will be no-op).
+    # TODO: since this implementation is inconsistent, explore a universal way to solve the issue.
+    _autolog(
+        log_traces=log_traces,
+        log_traces_from_compile=log_traces_from_compile,
+        log_traces_from_eval=log_traces_from_eval,
+        log_compiles=log_compiles,
+        log_evals=log_evals,
+        disable=disable,
+        silent=silent,
+    )
+
+    import dspy
+
+    from mlflow.dspy.callback import MlflowCallback
+    from mlflow.dspy.util import log_dspy_dataset, save_dspy_module_state
+
+    # Enable tracing by setting the MlflowCallback
+    if not disable:
+        if not any(isinstance(c, MlflowCallback) for c in dspy.settings.callbacks):
+            dspy.settings.configure(callbacks=[*dspy.settings.callbacks, MlflowCallback()])
+    else:
+        dspy.settings.configure(
+            callbacks=[c for c in dspy.settings.callbacks if not isinstance(c, MlflowCallback)]
+        )
+
+    def patch_fn(original, self, *args, **kwargs):
+        # NB: Since calling mlflow.dspy.autolog() again does not unpatch a function, we need to
+        # check this flag at runtime to determine if we should generate traces.
+        # method to disable tracing for compile and evaluate by default
+        @trace_disabled
+        def _trace_disabled_fn(self, *args, **kwargs):
+            return original(self, *args, **kwargs)
+
+        def _compile_fn(self, *args, **kwargs):
+            if callback := _active_callback():
+                callback.optimizer_stack_level += 1
+            try:
+                if get_autologging_config(FLAVOR_NAME, "log_traces_from_compile"):
+                    result = original(self, *args, **kwargs)
+                else:
+                    result = _trace_disabled_fn(self, *args, **kwargs)
+                return result
+            finally:
+                if callback:
+                    callback.optimizer_stack_level -= 1
+                    if callback.optimizer_stack_level == 0:
+                        # Reset the callback state after the completion of root compile
+                        callback.reset()
+
+        if isinstance(self, Teleprompter):
+            if not get_autologging_config(FLAVOR_NAME, "log_compiles"):
+                return _compile_fn(self, *args, **kwargs)
+
+            program = _compile_fn(self, *args, **kwargs)
+            # Save the state of the best model in json format
+            # so that users can see the demonstrations and instructions.
+            save_dspy_module_state(program, "best_model.json")
+
+            # Teleprompter.get_params is introduced in dspy 2.6.15
+            params = (
+                self.get_params()
+                if Version(importlib.metadata.version("dspy")) >= Version("2.6.15")
+                else {}
+            )
+            # Construct the dict of arguments passed to the compile call
+            inputs = construct_full_inputs(original, self, *args, **kwargs)
+            # Update params with the arguments passed to the compile call
+            params.update(inputs)
+            mlflow.log_params(
+                {k: v for k, v in inputs.items() if isinstance(v, (int, float, str, bool))}
+            )
+
+            if trainset := inputs.get("trainset"):
+                log_dspy_dataset(trainset, "trainset.json")
+            if valset := inputs.get("valset"):
+                log_dspy_dataset(valset, "valset.json")
+            return program
+
+        if isinstance(self, Teleprompter) and get_autologging_config(
+            FLAVOR_NAME, "log_traces_from_compile"
+        ):
+            return original(self, *args, **kwargs)
+
+        if isinstance(self, Evaluate) and get_autologging_config(
+            FLAVOR_NAME, "log_traces_from_eval"
+        ):
+            return original(self, *args, **kwargs)
+
+        return _trace_disabled_fn(self, *args, **kwargs)
+
+    from dspy.evaluate import Evaluate
+    from dspy.teleprompt import Teleprompter
+
+    compile_patch = "compile"
+    for cls in Teleprompter.__subclasses__():
+        # NB: This is to avoid the abstraction inheritance of superclasses that are defined
+        # only for the purposes of abstraction. The recursion behavior of the
+        # __subclasses__ dunder method will target the appropriate subclasses we need to patch.
+        if hasattr(cls, compile_patch):
+            safe_patch(
+                FLAVOR_NAME,
+                cls,
+                compile_patch,
+                patch_fn,
+                manage_run=get_autologging_config(FLAVOR_NAME, "log_compiles"),
+            )
+
+    call_patch = "__call__"
+    if hasattr(Evaluate, call_patch):
+        safe_patch(
+            FLAVOR_NAME,
+            Evaluate,
+            call_patch,
+            patch_fn,
+        )
+
+
+# This is required by mlflow.autolog()
+autolog.integration_name = FLAVOR_NAME
+
+
+@autologging_integration(FLAVOR_NAME)
+def _autolog(
+    log_traces: bool,
+    log_traces_from_compile: bool,
+    log_traces_from_eval: bool,
+    log_compiles: bool,
+    log_evals: bool,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    TODO: Implement patching logic for autologging artifacts.
+    """
+
+
+def _active_callback():
+    import dspy
+
+    from mlflow.dspy.callback import MlflowCallback
+
+    for callback in dspy.settings.callbacks:
+        if isinstance(callback, MlflowCallback):
+            return callback

venv/lib/python3.9/site-packages/mlflow/dspy/callback.py

@@ -0,0 +1,403 @@
+import logging
+import threading
+from collections import defaultdict
+from functools import wraps
+from typing import Any, Optional, Union
+
+import dspy
+from dspy.utils.callback import BaseCallback
+
+import mlflow
+from mlflow.dspy.save import FLAVOR_NAME
+from mlflow.dspy.util import log_dspy_module_params, save_dspy_module_state
+from mlflow.entities import SpanStatusCode, SpanType
+from mlflow.entities.run_status import RunStatus
+from mlflow.entities.span_event import SpanEvent
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc.context import get_prediction_context, maybe_set_prediction_context
+from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
+from mlflow.tracing.utils import (
+    end_client_span_or_trace,
+    set_span_chat_messages,
+    start_client_span_or_trace,
+)
+from mlflow.tracing.utils.token import SpanWithToken
+from mlflow.utils.autologging_utils import (
+    get_autologging_config,
+)
+
+_logger = logging.getLogger(__name__)
+_lock = threading.Lock()
+
+
+def skip_if_trace_disabled(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if get_autologging_config(FLAVOR_NAME, "log_traces"):
+            func(*args, **kwargs)
+
+    return wrapper
+
+
+class MlflowCallback(BaseCallback):
+    """Callback for generating MLflow traces for DSPy components"""
+
+    def __init__(self, dependencies_schema: Optional[dict[str, Any]] = None):
+        self._client = mlflow.MlflowClient()
+        self._dependencies_schema = dependencies_schema
+        # call_id: (LiveSpan, OTel token)
+        self._call_id_to_span: dict[str, SpanWithToken] = {}
+
+        ###### state management for optimization process ######
+        # The current callback logic assumes there is no optimization running in parallel.
+        # The state management may not work when multiple optimizations are running in parallel.
+        # optimizer_stack_level is used to determine if the callback is called within compile
+        # we cannot use boolean flag because the callback can be nested
+        self.optimizer_stack_level = 0
+        # call_id: (key, step)
+        self._call_id_to_metric_key: dict[str, tuple[str, int]] = {}
+        self._evaluation_counter = defaultdict(int)
+
+    def set_dependencies_schema(self, dependencies_schema: dict[str, Any]):
+        if self._dependencies_schema:
+            raise MlflowException(
+                "Dependencies schema should be set only once to the callback.",
+                error_code=MlflowException.INVALID_PARAMETER_VALUE,
+            )
+        self._dependencies_schema = dependencies_schema
+
+    @skip_if_trace_disabled
+    def on_module_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        span_type = self._get_span_type_for_module(instance)
+        attributes = self._get_span_attribute_for_module(instance)
+
+        # The __call__ method of dspy.Module has a signature of (self, *args, **kwargs),
+        # while all built-in modules only accepts keyword arguments. To avoid recording
+        # empty "args" key in the inputs, we remove it if it's empty.
+        if "args" in inputs and not inputs["args"]:
+            inputs.pop("args")
+
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.forward",
+            span_type=span_type,
+            inputs=self._unpack_kwargs(inputs),
+            attributes=attributes,
+        )
+
+    @skip_if_trace_disabled
+    def on_module_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        # NB: DSPy's Prediction object is a customized dictionary-like object, but its repr
+        # is not easy to read on UI. Therefore, we unpack it to a dictionary.
+        # https://github.com/stanfordnlp/dspy/blob/6fe693528323c9c10c82d90cb26711a985e18b29/dspy/primitives/prediction.py#L21-L28
+        if isinstance(outputs, dspy.Prediction):
+            outputs = outputs.toDict()
+
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_lm_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        span_type = (
+            SpanType.CHAT_MODEL if getattr(instance, "model_type", None) == "chat" else SpanType.LLM
+        )
+
+        attributes = {
+            **instance.kwargs,
+            "model": instance.model,
+            "model_type": instance.model_type,
+            "cache": instance.cache,
+        }
+
+        inputs = self._unpack_kwargs(inputs)
+
+        span = self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.__call__",
+            span_type=span_type,
+            inputs=inputs,
+            attributes=attributes,
+        )
+
+        if messages := self._extract_messages_from_lm_inputs(inputs):
+            try:
+                set_span_chat_messages(span, messages)
+            except Exception as e:
+                _logger.debug(f"Failed to set input messages for {span}. Error: {e}")
+
+    @skip_if_trace_disabled
+    def on_lm_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        st = self._call_id_to_span.get(call_id)
+        try:
+            output_msg = self._extract_messages_from_lm_outputs(outputs)
+            set_span_chat_messages(st.span, output_msg, append=True)
+        except Exception as e:
+            _logger.debug(f"Failed to set output messages for {call_id}. Error: {e}")
+
+        self._end_span(call_id, outputs, exception)
+
+    def _extract_messages_from_lm_inputs(self, inputs: dict[str, Any]) -> list[dict[str, str]]:
+        # LM input is either a list of messages or a prompt string
+        # https://github.com/stanfordnlp/dspy/blob/ac5bf56bb1ed7261d9637168563328c1dfeb27af/dspy/clients/lm.py#L92
+        # TODO: Extract tool definition once https://github.com/stanfordnlp/dspy/pull/2023 is merged
+        return inputs.get("messages") or [{"role": "user", "content": inputs.get("prompt")}]
+
+    def _extract_messages_from_lm_outputs(
+        self, outputs: list[Union[str, dict[str, Any]]]
+    ) -> list[dict[str, str]]:
+        # LM output is either a string or a dictionary of text and logprobs
+        # https://github.com/stanfordnlp/dspy/blob/ac5bf56bb1ed7261d9637168563328c1dfeb27af/dspy/clients/lm.py#L105-L114
+        # TODO: Extract tool calls once https://github.com/stanfordnlp/dspy/pull/2023 is merged
+        return [
+            {"role": "assistant", "content": o.get("text") if isinstance(o, dict) else o}
+            for o in outputs
+        ]
+
+    @skip_if_trace_disabled
+    def on_adapter_format_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.format",
+            span_type=SpanType.PARSER,
+            inputs=self._unpack_kwargs(inputs),
+            attributes={},
+        )
+
+    @skip_if_trace_disabled
+    def on_adapter_format_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_adapter_parse_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.parse",
+            span_type=SpanType.PARSER,
+            inputs=self._unpack_kwargs(inputs),
+            attributes={},
+        )
+
+    @skip_if_trace_disabled
+    def on_adapter_parse_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_tool_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        # DSPy uses the special "finish" tool to signal the end of the agent.
+        if instance.name == "finish":
+            return
+
+        inputs = self._unpack_kwargs(inputs)
+        # Tools are always called with keyword arguments only.
+        inputs.pop("args", None)
+
+        self._start_span(
+            call_id,
+            name=f"Tool.{instance.name}",
+            span_type=SpanType.TOOL,
+            inputs=inputs,
+            attributes={
+                "name": instance.name,
+                "description": instance.desc,
+                "args": instance.args,
+            },
+        )
+
+    @skip_if_trace_disabled
+    def on_tool_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        if call_id in self._call_id_to_span:
+            self._end_span(call_id, outputs, exception)
+
+    def on_evaluate_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        """
+        Callback handler at the beginning of evaluation call. Available with DSPy>=2.6.9.
+        This callback starts a nested run for each evaluation call inside optimization.
+        If called outside optimization and no active run exists, it creates a new run.
+        """
+        if not get_autologging_config(FLAVOR_NAME, "log_evals"):
+            return
+
+        key = "eval"
+        if callback_metadata := inputs.get("callback_metadata"):
+            if "metric_key" in callback_metadata:
+                key = callback_metadata["metric_key"]
+        if self.optimizer_stack_level > 0:
+            with _lock:
+                # we may want to include optimizer_stack_level in the key
+                # to handle nested optimization
+                step = self._evaluation_counter[key]
+                self._evaluation_counter[key] += 1
+            self._call_id_to_metric_key[call_id] = (key, step)
+            mlflow.start_run(run_name=f"{key}_{step}", nested=True)
+        else:
+            mlflow.start_run(run_name=key, nested=True)
+        if program := inputs.get("program"):
+            save_dspy_module_state(program, "model.json")
+            log_dspy_module_params(program)
+
+    def on_evaluate_end(
+        self,
+        call_id: str,
+        outputs: Any,
+        exception: Optional[Exception] = None,
+    ):
+        """
+        Callback handler at the end of evaluation call. Available with DSPy>=2.6.9.
+        This callback logs the evaluation score to the individual run
+        and add eval metric to the parent run if called inside optimization.
+        """
+        if not get_autologging_config(FLAVOR_NAME, "log_evals"):
+            return
+        if exception:
+            mlflow.end_run(status=RunStatus.to_string(RunStatus.FAILED))
+            return
+        score = None
+        if isinstance(outputs, float):
+            score = outputs
+        elif isinstance(outputs, tuple):
+            score = outputs[0]
+        elif isinstance(outputs, dspy.Prediction):
+            score = float(outputs)
+            try:
+                mlflow.log_table(self._generate_result_table(outputs.results), "result_table.json")
+            except Exception:
+                _logger.debug("Failed to log result table.", exc_info=True)
+        if score is not None:
+            mlflow.log_metric("eval", score)
+
+        mlflow.end_run()
+        # Log the evaluation score to the parent run if called inside optimization
+        if self.optimizer_stack_level > 0 and mlflow.active_run() is not None:
+            if call_id not in self._call_id_to_metric_key:
+                return
+            key, step = self._call_id_to_metric_key.pop(call_id)
+            if score is not None:
+                mlflow.log_metric(
+                    key,
+                    score,
+                    step=step,
+                )
+
+    def reset(self):
+        self._call_id_to_metric_key: dict[str, tuple[str, int]] = {}
+        self._evaluation_counter = defaultdict(int)
+
+    def _start_span(
+        self,
+        call_id: str,
+        name: str,
+        span_type: SpanType,
+        inputs: dict[str, Any],
+        attributes: dict[str, Any],
+    ):
+        prediction_context = get_prediction_context()
+        if prediction_context and self._dependencies_schema:
+            prediction_context.update(**self._dependencies_schema)
+
+        with maybe_set_prediction_context(prediction_context):
+            span = start_client_span_or_trace(
+                self._client,
+                name=name,
+                span_type=span_type,
+                parent_span=mlflow.get_current_active_span(),
+                inputs=inputs,
+                attributes=attributes,
+            )
+
+        token = set_span_in_context(span)
+        self._call_id_to_span[call_id] = SpanWithToken(span, token)
+
+        return span
+
+    def _end_span(
+        self,
+        call_id: str,
+        outputs: Optional[Any],
+        exception: Optional[Exception] = None,
+    ):
+        st = self._call_id_to_span.pop(call_id, None)
+
+        if not st.span:
+            _logger.warning(f"Failed to end a span. Span not found for call_id: {call_id}")
+            return
+
+        status = SpanStatusCode.OK if exception is None else SpanStatusCode.ERROR
+
+        if exception:
+            st.span.add_event(SpanEvent.from_exception(exception))
+
+        try:
+            end_client_span_or_trace(
+                client=self._client,
+                span=st.span,
+                outputs=outputs,
+                status=status,
+            )
+        finally:
+            detach_span_from_context(st.token)
+
+    def _get_span_type_for_module(self, instance):
+        if isinstance(instance, dspy.Retrieve):
+            return SpanType.RETRIEVER
+        elif isinstance(instance, dspy.ReAct):
+            return SpanType.AGENT
+        elif isinstance(instance, dspy.Predict):
+            return SpanType.LLM
+        elif isinstance(instance, dspy.Adapter):
+            return SpanType.PARSER
+        else:
+            return SpanType.CHAIN
+
+    def _get_span_attribute_for_module(self, instance):
+        if isinstance(instance, dspy.Predict):
+            return {"signature": instance.signature.signature}
+        elif isinstance(instance, dspy.ChainOfThought):
+            if hasattr(instance, "signature"):
+                signature = instance.signature.signature
+            else:
+                signature = instance.predict.signature.signature
+
+            attributes = {"signature": signature}
+            if hasattr(instance, "extended_signature"):
+                attributes["extended_signature"] = instance.extended_signature.signature
+            return attributes
+        return {}
+
+    def _unpack_kwargs(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        """Unpacks the kwargs from the inputs dictionary"""
+        # NB: Not using pop() to avoid modifying the original inputs dictionary
+        kwargs = inputs.get("kwargs", {})
+        inputs_wo_kwargs = {k: v for k, v in inputs.items() if k != "kwargs"}
+        return {**inputs_wo_kwargs, **kwargs}
+
+    def _generate_result_table(
+        self, outputs: list[tuple[dspy.Example, dspy.Prediction, Any]]
+    ) -> dict[str, list[Any]]:
+        result = {"score": []}
+        for i, (example, prediction, score) in enumerate(outputs):
+            for k, v in example.items():
+                if f"example_{k}" not in result:
+                    result[f"example_{k}"] = [None] * i
+                result[f"example_{k}"].append(v)
+
+            for k, v in prediction.items():
+                if f"pred_{k}" not in result:
+                    result[f"pred_{k}"] = [None] * i
+                result[f"pred_{k}"].append(v)
+
+            result["score"].append(score)
+
+            for k, v in result.items():
+                if len(v) != i + 1:
+                    result[k].append(None)
+
+        return result

venv/lib/python3.9/site-packages/mlflow/dspy/load.py

@@ -0,0 +1,89 @@
+import os
+
+import cloudpickle
+
+from mlflow.models import Model
+from mlflow.models.dependencies_schemas import _get_dependencies_schema_from_model
+from mlflow.tracing.provider import trace_disabled
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.annotations import experimental
+from mlflow.utils.model_utils import (
+    _add_code_from_conf_to_system_path,
+    _get_flavor_configuration,
+)
+
+_DEFAULT_MODEL_PATH = "data/model.pkl"
+
+
+def _set_dependency_schema_to_tracer(model_path, callbacks):
+    """
+    Set dependency schemas from the saved model metadata to the tracer
+    to propagate it to inference traces.
+    """
+    from mlflow.dspy.callback import MlflowCallback
+
+    tracer = next((cb for cb in callbacks if isinstance(cb, MlflowCallback)), None)
+    if tracer is None:
+        return
+
+    model = Model.load(model_path)
+    tracer.set_dependencies_schema(_get_dependencies_schema_from_model(model))
+
+
+def _load_model(model_uri, dst_path=None):
+    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
+    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name="dspy")
+
+    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
+    model_path = flavor_conf.get("model_path", _DEFAULT_MODEL_PATH)
+    with open(os.path.join(local_model_path, model_path), "rb") as f:
+        loaded_wrapper = cloudpickle.load(f)
+
+    _set_dependency_schema_to_tracer(local_model_path, loaded_wrapper.dspy_settings["callbacks"])
+
+    return loaded_wrapper
+
+
+@experimental
+@trace_disabled  # Suppress traces for internal calls while loading model
+def load_model(model_uri, dst_path=None):
+    """
+    Load a Dspy model from a run.
+
+    This function will also set the global dspy settings `dspy.settings` by the saved settings.
+
+    Args:
+        model_uri: The location, in URI format, of the MLflow model. For example:
+
+            - ``/Users/me/path/to/local/model``
+            - ``relative/path/to/local/model``
+            - ``s3://my_bucket/path/to/model``
+            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
+            - ``mlflow-artifacts:/path/to/model``
+
+            For more information about supported URI schemes, see
+            `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
+            artifact-locations>`_.
+        dst_path: The local filesystem path to utilize for downloading the model artifact.
+            This directory must already exist if provided. If unspecified, a local output
+            path will be created.
+
+    Returns:
+        An `dspy.module` instance, representing the dspy model.
+    """
+    import dspy
+
+    wrapper = _load_model(model_uri, dst_path)
+
+    # Set the global dspy settings for reproducing the model's behavior when the model is
+    # loaded via `mlflow.dspy.load_model`. Note that for the model to be loaded as pyfunc,
+    # settings will be set in the wrapper's `predict` method via local context to avoid the
+    # "dspy.settings can only be changed by the thread that initially configured it" error
+    # in Databricks model serving.
+    dspy.settings.configure(**wrapper.dspy_settings)
+
+    return wrapper.model
+
+
+def _load_pyfunc(path):
+    return _load_model(path)

venv/lib/python3.9/site-packages/mlflow/dspy/save.py

@@ -0,0 +1,365 @@
+"""Functions for saving DSPY models to MLflow."""
+
+import os
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import cloudpickle
+import yaml
+
+import mlflow
+from mlflow import pyfunc
+from mlflow.dspy.wrapper import DspyChatModelWrapper, DspyModelWrapper
+from mlflow.entities.model_registry.prompt import Prompt
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.models import (
+    Model,
+    ModelInputExample,
+    ModelSignature,
+    infer_pip_requirements,
+)
+from mlflow.models.dependencies_schemas import _get_dependencies_schemas
+from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.models.rag_signatures import SIGNATURE_FOR_LLM_INFERENCE_TASK
+from mlflow.models.resources import Resource, _ResourceBuilder
+from mlflow.models.signature import _infer_signature_from_input_example
+from mlflow.models.utils import _save_example
+from mlflow.tracing.provider import trace_disabled
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+from mlflow.utils.annotations import experimental
+from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
+from mlflow.utils.environment import (
+    _CONDA_ENV_FILE_NAME,
+    _CONSTRAINTS_FILE_NAME,
+    _PYTHON_ENV_FILE_NAME,
+    _REQUIREMENTS_FILE_NAME,
+    _mlflow_conda_env,
+    _process_conda_env,
+    _process_pip_requirements,
+    _PythonEnv,
+)
+from mlflow.utils.file_utils import get_total_file_size, write_to
+from mlflow.utils.model_utils import (
+    _validate_and_copy_code_paths,
+    _validate_and_prepare_target_save_path,
+)
+from mlflow.utils.requirements_utils import _get_pinned_requirement
+
+FLAVOR_NAME = "dspy"
+
+_MODEL_SAVE_PATH = "model"
+_MODEL_DATA_PATH = "data"
+
+
+def get_default_pip_requirements():
+    """
+    Returns:
+        A list of default pip requirements for MLflow Models produced by Dspy flavor. Calls to
+        `save_model()` and `log_model()` produce a pip environment that, at minimum, contains these
+        requirements.
+    """
+    return [_get_pinned_requirement("dspy")]
+
+
+def get_default_conda_env():
+    """
+    Returns:
+        The default Conda environment for MLflow Models produced by calls to `save_model()` and
+        `log_model()`.
+    """
+    return _mlflow_conda_env(additional_pip_deps=get_default_pip_requirements())
+
+
+@experimental
+@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
+@trace_disabled  # Suppress traces for internal predict calls while logging model
+def save_model(
+    model,
+    path: str,
+    task: Optional[str] = None,
+    model_config: Optional[dict[str, Any]] = None,
+    code_paths: Optional[list[str]] = None,
+    mlflow_model: Optional[Model] = None,
+    conda_env: Optional[Union[list[str], str]] = None,
+    signature: Optional[ModelSignature] = None,
+    input_example: Optional[ModelInputExample] = None,
+    pip_requirements: Optional[Union[list[str], str]] = None,
+    extra_pip_requirements: Optional[Union[list[str], str]] = None,
+    metadata: Optional[dict[str, Any]] = None,
+    resources: Optional[Union[str, Path, list[Resource]]] = None,
+):
+    """
+    Save a Dspy model.
+
+    This method saves a Dspy model along with metadata such as model signature and conda
+    environments to local file system. This method is called inside `mlflow.dspy.log_model()`.
+
+    Args:
+        model: an instance of `dspy.Module`. The Dspy model/module to be saved.
+        path: local path where the MLflow model is to be saved.
+        task: defaults to None. The task type of the model. Can only be `llm/v1/chat` or None for
+            now.
+        model_config: keyword arguments to be passed to the Dspy Module at instantiation.
+        code_paths: {{ code_paths }}
+        mlflow_model: an instance of `mlflow.models.Model`, defaults to None. MLflow model
+            configuration to which to add the Dspy model metadata. If None, a blank instance will
+            be created.
+        conda_env: {{ conda_env }}
+        signature: {{ signature }}
+        input_example: {{ input_example }}
+        pip_requirements: {{ pip_requirements }}
+        extra_pip_requirements: {{ extra_pip_requirements }}
+        metadata: {{ metadata }}
+        resources: A list of model resources or a resources.yaml file containing a list of
+            resources required to serve the model.
+    """
+
+    import dspy
+
+    from mlflow.transformers.llm_inference_utils import (
+        _LLM_INFERENCE_TASK_KEY,
+        _METADATA_LLM_INFERENCE_TASK_KEY,
+    )
+
+    if signature:
+        num_inputs = len(signature.inputs.inputs)
+        if num_inputs == 0:
+            raise MlflowException(
+                "The model signature's input schema must contain at least one field.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+    if task and task not in SIGNATURE_FOR_LLM_INFERENCE_TASK:
+        raise MlflowException(
+            "Invalid task: {task} at `mlflow.dspy.save_model()` call. The task must be None or one "
+            f"of: {list(SIGNATURE_FOR_LLM_INFERENCE_TASK.keys())}",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    if mlflow_model is None:
+        mlflow_model = Model()
+    if signature is not None:
+        mlflow_model.signature = signature
+    saved_example = None
+    if input_example is not None:
+        path = os.path.abspath(path)
+        _validate_and_prepare_target_save_path(path)
+        saved_example = _save_example(mlflow_model, input_example, path)
+    if metadata is not None:
+        mlflow_model.metadata = metadata
+
+    with _get_dependencies_schemas() as dependencies_schemas:
+        schema = dependencies_schemas.to_dict()
+        if schema is not None:
+            if mlflow_model.metadata is None:
+                mlflow_model.metadata = {}
+            mlflow_model.metadata.update(schema)
+
+    model_data_subpath = _MODEL_DATA_PATH
+    # Construct new data folder in existing path.
+    data_path = os.path.join(path, model_data_subpath)
+    os.makedirs(data_path, exist_ok=True)
+    # Set the model path to end with ".pkl" as we use cloudpickle for serialization.
+    model_subpath = os.path.join(model_data_subpath, _MODEL_SAVE_PATH) + ".pkl"
+    model_path = os.path.join(path, model_subpath)
+    # Dspy has a global context `dspy.settings`, and we need to save it along with the model.
+    dspy_settings = dict(dspy.settings.config)
+
+    # Don't save the trace in the model, which is only useful during the training phase.
+    dspy_settings.pop("trace", None)
+
+    # Store both dspy model and settings in `DspyChatModelWrapper` or `DspyModelWrapper` for
+    # serialization.
+    if task == "llm/v1/chat":
+        wrapped_dspy_model = DspyChatModelWrapper(model, dspy_settings, model_config)
+    else:
+        wrapped_dspy_model = DspyModelWrapper(model, dspy_settings, model_config)
+
+    with open(model_path, "wb") as f:
+        cloudpickle.dump(wrapped_dspy_model, f)
+
+    flavor_options = {
+        "model_path": model_subpath,
+    }
+
+    if task:
+        if mlflow_model.signature is None:
+            mlflow_model.signature = SIGNATURE_FOR_LLM_INFERENCE_TASK[task]
+        flavor_options.update({_LLM_INFERENCE_TASK_KEY: task})
+        if mlflow_model.metadata:
+            mlflow_model.metadata[_METADATA_LLM_INFERENCE_TASK_KEY] = task
+        else:
+            mlflow_model.metadata = {_METADATA_LLM_INFERENCE_TASK_KEY: task}
+
+    if saved_example and mlflow_model.signature is None:
+        signature = _infer_signature_from_input_example(saved_example, wrapped_dspy_model)
+        mlflow_model.signature = signature
+    code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
+
+    # Add flavor info to `mlflow_model`.
+    mlflow_model.add_flavor(FLAVOR_NAME, code=code_dir_subpath, **flavor_options)
+    # Add loader_module, data and env data to `mlflow_model`.
+    pyfunc.add_to_model(
+        mlflow_model,
+        loader_module="mlflow.dspy",
+        code=code_dir_subpath,
+        conda_env=_CONDA_ENV_FILE_NAME,
+        python_env=_PYTHON_ENV_FILE_NAME,
+    )
+
+    # Add model file size to `mlflow_model`.
+    if size := get_total_file_size(path):
+        mlflow_model.model_size_bytes = size
+
+    # Add resources if specified.
+    if resources is not None:
+        if isinstance(resources, (Path, str)):
+            serialized_resource = _ResourceBuilder.from_yaml_file(resources)
+        else:
+            serialized_resource = _ResourceBuilder.from_resources(resources)
+
+        mlflow_model.resources = serialized_resource
+
+    # Save mlflow_model to path/MLmodel.
+    mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
+
+    if conda_env is None:
+        if pip_requirements is None:
+            default_reqs = get_default_pip_requirements()
+            # To ensure `_load_pyfunc` can successfully load the model during the dependency
+            # inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
+            inferred_reqs = infer_pip_requirements(path, FLAVOR_NAME, fallback=default_reqs)
+            default_reqs = sorted(set(inferred_reqs).union(default_reqs))
+        else:
+            default_reqs = None
+        conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
+            default_reqs,
+            pip_requirements,
+            extra_pip_requirements,
+        )
+    else:
+        conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
+
+    with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
+        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
+
+    # Save `constraints.txt` if necessary.
+    if pip_constraints:
+        write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
+
+    # Save `requirements.txt`.
+    write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
+
+    _PythonEnv.current().to_yaml(os.path.join(path, _PYTHON_ENV_FILE_NAME))
+
+
+@experimental
+@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
+@trace_disabled  # Suppress traces for internal predict calls while logging model
+def log_model(
+    dspy_model,
+    artifact_path: str,
+    task: Optional[str] = None,
+    model_config: Optional[dict[str, Any]] = None,
+    code_paths: Optional[list[str]] = None,
+    conda_env: Optional[Union[list[str], str]] = None,
+    signature: Optional[ModelSignature] = None,
+    input_example: Optional[ModelInputExample] = None,
+    registered_model_name: Optional[str] = None,
+    await_registration_for: int = DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+    pip_requirements: Optional[Union[list[str], str]] = None,
+    extra_pip_requirements: Optional[Union[list[str], str]] = None,
+    metadata: Optional[dict[str, Any]] = None,
+    resources: Optional[Union[str, Path, list[Resource]]] = None,
+    prompts: Optional[list[Union[str, Prompt]]] = None,
+):
+    """
+    Log a Dspy model along with metadata to MLflow.
+
+    This method saves a Dspy model along with metadata such as model signature and conda
+    environments to MLflow.
+
+    Args:
+        dspy_model: an instance of `dspy.Module`. The Dspy model to be saved.
+        artifact_path: the run-relative path to which to log model artifacts.
+        task: defaults to None. The task type of the model. Can only be `llm/v1/chat` or None for
+            now.
+        model_config: keyword arguments to be passed to the Dspy Module at instantiation.
+        code_paths: {{ code_paths }}
+        conda_env: {{ conda_env }}
+        signature: {{ signature }}
+        input_example: {{ input_example }}
+        registered_model_name: defaults to None. If set, create a model version under
+            `registered_model_name`, also create a registered model if one with the given name does
+            not exist.
+        await_registration_for: defaults to
+            `mlflow.tracking._model_registry.DEFAULT_AWAIT_MAX_SLEEP_SECONDS`. Number of
+            seconds to wait for the model version to finish being created and is in ``READY``
+            status. By default, the function waits for five minutes. Specify 0 or None to skip
+            waiting.
+        pip_requirements: {{ pip_requirements }}
+        extra_pip_requirements: {{ extra_pip_requirements }}
+        metadata: Custom metadata dictionary passed to the model and stored in the MLmodel
+            file.
+        resources: A list of model resources or a resources.yaml file containing a list of
+            resources required to serve the model.
+        prompts: {{ prompts }}
+
+    .. code-block:: python
+        :caption: Example
+
+        import dspy
+        import mlflow
+        from mlflow.models import ModelSignature
+        from mlflow.types.schema import ColSpec, Schema
+
+        # Set up the LM.
+        lm = dspy.LM(model="openai/gpt-4o-mini", max_tokens=250)
+        dspy.settings.configure(lm=lm)
+
+
+        class CoT(dspy.Module):
+            def __init__(self):
+                super().__init__()
+                self.prog = dspy.ChainOfThought("question -> answer")
+
+            def forward(self, question):
+                return self.prog(question=question)
+
+
+        dspy_model = CoT()
+
+        mlflow.set_tracking_uri("http://127.0.0.1:5000")
+        mlflow.set_experiment("test-dspy-logging")
+
+        from mlflow.dspy import log_model
+
+        input_schema = Schema([ColSpec("string")])
+        output_schema = Schema([ColSpec("string")])
+        signature = ModelSignature(inputs=input_schema, outputs=output_schema)
+
+        with mlflow.start_run():
+            log_model(
+                dspy_model,
+                "model",
+                input_example="what is 2 + 2?",
+                signature=signature,
+            )
+    """
+    return Model.log(
+        artifact_path=artifact_path,
+        flavor=mlflow.dspy,
+        model=dspy_model,
+        task=task,
+        model_config=model_config,
+        code_paths=code_paths,
+        conda_env=conda_env,
+        registered_model_name=registered_model_name,
+        signature=signature,
+        input_example=input_example,
+        await_registration_for=await_registration_for,
+        pip_requirements=pip_requirements,
+        extra_pip_requirements=extra_pip_requirements,
+        metadata=metadata,
+        resources=resources,
+        prompts=prompts,
+    )

venv/lib/python3.9/site-packages/mlflow/dspy/util.py

@@ -0,0 +1,109 @@
+import logging
+import tempfile
+from collections import defaultdict
+from pathlib import Path
+from typing import Optional
+
+from dspy import Example
+
+import mlflow
+
+_logger = logging.getLogger(__name__)
+
+
+def save_dspy_module_state(program, file_name: str = "model.json"):
+    """
+    Save states of dspy `Module` to a temporary directory and log it as an artifact.
+
+    Args:
+        program: The dspy `Module` to be saved.
+        file_name: The name of the file to save the dspy module state. Default is `model.json`.
+    """
+    try:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir, file_name)
+            program.save(path)
+            mlflow.log_artifact(path)
+    except Exception as e:
+        _logger.warning(f"Failed to save dspy module state: {e}")
+
+
+def log_dspy_module_params(program):
+    """
+    Log the parameters of the dspy `Module` as run parameters.
+
+    Args:
+        program: The dspy `Module` to be logged.
+    """
+    try:
+        states = program.dump_state()
+        flat_state_dict = _flatten_dspy_module_state(
+            states, exclude_keys=("metadata", "lm", "traces", "train")
+        )
+        mlflow.log_params(
+            {f"{program.__class__.__name__}.{k}": v for k, v in flat_state_dict.items()}
+        )
+    except Exception as e:
+        _logger.warning(f"Failed to log dspy module params: {e}")
+
+
+def log_dspy_dataset(dataset: list["Example"], file_name: str):
+    """
+    Log the DSPy dataset as a table.
+
+    Args:
+        dataset: The dataset to be logged.
+        file_name: The name of the file to save the dataset.
+    """
+    result = defaultdict(list)
+    try:
+        for example in dataset:
+            for k, v in example.items():
+                result[k].append(v)
+        mlflow.log_table(result, file_name)
+    except Exception as e:
+        _logger.warning(f"Failed to log dataset: {e}")
+
+
+def _flatten_dspy_module_state(
+    d, parent_key="", sep=".", exclude_keys: Optional[set] = None
+) -> dict:
+    """
+    Flattens a nested dictionary and accumulates the key names.
+
+    Args:
+        d: The dictionary or list to flatten.
+        parent_key: The base key used in recursion. Defaults to "".
+        sep: Separator for nested keys. Defaults to '.'.
+        exclude_keys: Keys to exclude from the flattened dictionary. Defaults to ().
+
+    Returns:
+        dict: A flattened dictionary with accumulated keys.
+
+    Example:
+        >>> _flatten_dspy_module_state({"a": {"b": [5, 6]}})
+        {'a.b.0': 5, 'a.b.1': 6}
+    """
+    items = {}
+
+    if isinstance(d, dict):
+        for k, v in d.items():
+            if exclude_keys and k in exclude_keys:
+                continue
+            new_key = f"{parent_key}{sep}{k}" if parent_key else k
+            if isinstance(v, Example):
+                # Don't flatten Example objects further even if it has dict or list values
+                v = {key: str(value) for key, value in v.items()}
+            items.update(_flatten_dspy_module_state(v, new_key, sep))
+    elif isinstance(d, list):
+        for i, v in enumerate(d):
+            new_key = f"{parent_key}{sep}{i}" if parent_key else str(i)
+            if isinstance(v, Example):
+                # Don't flatten Example objects further even if it has dict or list values
+                v = {key: str(value) for key, value in v.items()}
+            items.update(_flatten_dspy_module_state(v, new_key, sep))
+    else:
+        if d is not None:
+            items[parent_key] = d
+
+    return items

venv/lib/python3.9/site-packages/mlflow/dspy/wrapper.py

@@ -0,0 +1,157 @@
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    import dspy
+
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+from mlflow.protos.databricks_pb2 import (
+    INVALID_PARAMETER_VALUE,
+)
+from mlflow.pyfunc import PythonModel
+
+_logger = logging.getLogger(__name__)
+
+
+class DspyModelWrapper(PythonModel):
+    """MLflow PyFunc wrapper class for Dspy models.
+
+    This wrapper serves two purposes:
+        - It stores the Dspy model along with dspy global settings, which are required for seamless
+            saving and loading.
+        - It provides a `predict` method so that it can be loaded as an MLflow pyfunc, which is
+            used at serving time.
+    """
+
+    def __init__(
+        self,
+        model: "dspy.Module",
+        dspy_settings: dict[str, Any],
+        model_config: Optional[dict[str, Any]] = None,
+    ):
+        self.model = model
+        self.dspy_settings = dspy_settings
+        self.model_config = model_config or {}
+
+    def predict(self, inputs: Any, params: Optional[dict[str, Any]] = None):
+        import dspy
+        import numpy as np
+        import pandas as pd
+
+        supported_input_types = (np.ndarray, pd.DataFrame, str, dict)
+        if not isinstance(inputs, supported_input_types):
+            raise MlflowException(
+                f"`inputs` must be one of: {[x.__name__ for x in supported_input_types]}, but "
+                f"received type: {type(inputs)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+        if isinstance(inputs, pd.DataFrame):
+            inputs = inputs.values
+        if isinstance(inputs, np.ndarray):
+            flatten = inputs.reshape(-1)
+            if len(flatten) > 1:
+                raise MlflowException(
+                    "Dspy model doesn't support multiple inputs or batch inference. Please "
+                    "provide a single input.",
+                    INVALID_PARAMETER_VALUE,
+                )
+            inputs = str(flatten[0])
+
+        with dspy.context(**self.dspy_settings):
+            if isinstance(inputs, dict):
+                return self.model(**inputs).toDict()
+            if isinstance(inputs, str):
+                return self.model(inputs).toDict()
+
+
+class DspyChatModelWrapper(DspyModelWrapper):
+    """MLflow PyFunc wrapper class for Dspy chat models."""
+
+    def predict(self, inputs: Any, params: Optional[dict[str, Any]] = None):
+        import dspy
+        import pandas as pd
+
+        if isinstance(inputs, dict):
+            converted_inputs = inputs["messages"]
+        elif isinstance(inputs, pd.DataFrame):
+            converted_inputs = inputs.messages[0]
+        else:
+            raise MlflowException(
+                f"Unsupported input type: {type(inputs)}. To log a DSPy model with task "
+                "'llm/v1/chat', the input must be a dict or a pandas DataFrame.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        # `dspy.settings` cannot be shared across threads, so we are setting the context at every
+        # predict call.
+        with dspy.context(**self.dspy_settings):
+            outputs = self.model(converted_inputs)
+
+        choices = []
+        if isinstance(outputs, str):
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": outputs},
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, dict):
+            role = outputs.get("role", "assistant")
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {"role": role, "content": json.dumps(outputs)},
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, dspy.Prediction):
+            choices.append(
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": json.dumps(outputs.toDict()),
+                    },
+                    "finish_reason": "stop",
+                }
+            )
+        elif isinstance(outputs, list):
+            for output in outputs:
+                if isinstance(output, dict):
+                    role = output.get("role", "assistant")
+                    choices.append(
+                        {
+                            "index": 0,
+                            "message": {"role": role, "content": json.dumps(outputs)},
+                            "finish_reason": "stop",
+                        }
+                    )
+                elif isinstance(output, dspy.Prediction):
+                    choices.append(
+                        {
+                            "index": 0,
+                            "message": {
+                                "role": role,
+                                "content": json.dumps(outputs.toDict()),
+                            },
+                            "finish_reason": "stop",
+                        }
+                    )
+                else:
+                    raise MlflowException(
+                        f"Unsupported output type: {type(output)}. To log a DSPy model with task "
+                        "'llm/v1/chat', the DSPy model must return a dict, a dspy.Prediction, or a "
+                        "list of dicts or dspy.Prediction.",
+                        INVALID_PARAMETER_VALUE,
+                    )
+        else:
+            raise MlflowException(
+                f"Unsupported output type: {type(outputs)}. To log a DSPy model with task "
+                "'llm/v1/chat', the DSPy model must return a dict, a dspy.Prediction, or a list of "
+                "dicts or dspy.Prediction.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        return {"choices": choices}