star

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

@@ -0,0 +1,7 @@
+import mlflow.pyfunc.loaders.chat_agent
+import mlflow.pyfunc.loaders.chat_model
+import mlflow.pyfunc.loaders.code_model
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if IS_PYDANTIC_V2_OR_NEWER:
+    import mlflow.pyfunc.loaders.responses_agent  # noqa: F401

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/chat_agent.py

@@ -0,0 +1,117 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.agent import (
+    ChatAgentChunk,
+    ChatAgentMessage,
+    ChatAgentResponse,
+    ChatContext,
+)
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, chat_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ChatAgentPyfuncWrapper(chat_agent)
+
+
+@experimental
+class _ChatAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatAgent`.
+    """
+
+    def __init__(self, chat_agent):
+        """
+        Args:
+            chat_agent: An instance of a subclass of :class:`~ChatAgent`.
+        """
+        self.chat_agent = chat_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_agent
+
+    def _convert_input(
+        self, model_input
+    ) -> tuple[list[ChatAgentMessage], Optional[ChatContext], Optional[dict[str, Any]]]:
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatAgentMessage(**message) for message in dict_input.get("messages", [])]
+        context = ChatContext(**dict_input["context"]) if "context" in dict_input else None
+        custom_inputs = dict_input.get("custom_inputs", None)
+
+        return messages, context, custom_inputs
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the (:py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>`)
+                schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        response = self.chat_agent.predict(messages, context, custom_inputs)
+        return self._response_to_dict(response, ChatAgentResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>`) schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        for response in self.chat_agent.predict_stream(messages, context, custom_inputs):
+            yield self._response_to_dict(response, ChatAgentChunk)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/chat_model.py

@@ -0,0 +1,127 @@
+import inspect
+import logging
+from typing import Any, Generator, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.llm import ChatCompletionChunk, ChatCompletionResponse, ChatMessage, ChatParams
+from mlflow.utils.annotations import experimental
+
+_logger = logging.getLogger(__name__)
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, chat_model, signature = _load_context_model_and_signature(model_path, model_config)
+    return _ChatModelPyfuncWrapper(chat_model=chat_model, context=context, signature=signature)
+
+
+@experimental
+class _ChatModelPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatModel`.
+    """
+
+    def __init__(self, chat_model, context, signature):
+        """
+        Args:
+            chat_model: An instance of a subclass of :class:`~ChatModel`.
+            context: A :class:`~PythonModelContext` instance containing artifacts that
+                        ``chat_model`` may use when performing inference.
+            signature: :class:`~ModelSignature` instance describing model input and output.
+        """
+        self.chat_model = chat_model
+        self.context = context
+        self.signature = signature
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_model
+
+    def _convert_input(self, model_input):
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, "
+                f"but got {type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatMessage.from_dict(message) for message in dict_input.pop("messages", [])]
+        params = ChatParams.from_dict(dict_input)
+        return messages, params
+
+    def predict(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Model predictions in :py:class:`~ChatCompletionResponse` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict).parameters
+        if "context" in parameters or len(parameters) == 3:
+            response = self.chat_model.predict(self.context, messages, params)
+        else:
+            response = self.chat_model.predict(messages, params)
+        return self._response_to_dict(response)
+
+    def _response_to_dict(self, response: ChatCompletionResponse) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionResponse):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionResponse, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def _streaming_response_to_dict(self, response: ChatCompletionChunk) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionChunk):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionChunk, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Generator over model predictions in :py:class:`~ChatCompletionChunk` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict_stream).parameters
+        if "context" in parameters or len(parameters) == 3:
+            stream = self.chat_model.predict_stream(self.context, messages, params)
+        else:
+            stream = self.chat_model.predict_stream(messages, params)
+
+        for response in stream:
+            yield self._streaming_response_to_dict(response)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/code_model.py

@@ -0,0 +1,31 @@
+from typing import Any, Optional
+
+from mlflow.pyfunc.loaders.chat_agent import _ChatAgentPyfuncWrapper
+from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
+from mlflow.pyfunc.model import (
+    ChatAgent,
+    ChatModel,
+    _load_context_model_and_signature,
+    _PythonModelPyfuncWrapper,
+)
+
+try:
+    from mlflow.pyfunc.model import ResponsesAgent
+
+    IS_RESPONSES_AGENT_AVAILABLE = True
+except ImportError:
+    IS_RESPONSES_AGENT_AVAILABLE = False
+
+
+def _load_pyfunc(local_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, model, signature = _load_context_model_and_signature(local_path, model_config)
+    if isinstance(model, ChatModel):
+        return _ChatModelPyfuncWrapper(model, context, signature)
+    elif isinstance(model, ChatAgent):
+        return _ChatAgentPyfuncWrapper(model)
+    elif IS_RESPONSES_AGENT_AVAILABLE and isinstance(model, ResponsesAgent):
+        from mlflow.pyfunc.loaders.responses_agent import _ResponsesAgentPyfuncWrapper
+
+        return _ResponsesAgentPyfuncWrapper(model)
+    else:
+        return _PythonModelPyfuncWrapper(model, context, signature)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/responses_agent.py

@@ -0,0 +1,108 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import _load_context_model_and_signature
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if not IS_PYDANTIC_V2_OR_NEWER:
+    raise ImportError(
+        "ResponsesAgent and its pydantic classes are not supported in pydantic v1. "
+        "Please upgrade to pydantic v2 or newer to use ResponsesAgent.",
+    )
+
+from mlflow.types.responses import ResponsesRequest, ResponsesResponse, ResponsesStreamEvent
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, responses_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ResponsesAgentPyfuncWrapper(responses_agent)
+
+
+@experimental
+class _ResponsesAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by
+    :class:`~ResponsesAgent`.
+    """
+
+    def __init__(self, responses_agent):
+        self.responses_agent = responses_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.responses_agent
+
+    def _convert_input(self, model_input) -> ResponsesRequest:
+        import pandas
+
+        if isinstance(model_input, pandas.DataFrame):
+            model_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        elif not isinstance(model_input, dict):
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return ResponsesRequest(**model_input)
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the
+            (:py:class:`ResponsesResponse <mlflow.types.responses.ResponsesResponse>`)
+            schema.
+        """
+        request = self._convert_input(model_input)
+        response = self.responses_agent.predict(request)
+        return self._response_to_dict(response, ResponsesResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ResponsesStreamEvent <mlflow.types.responses.ResponsesStreamEvent>`)
+                schema.
+        """
+        request = self._convert_input(model_input)
+        for response in self.responses_agent.predict_stream(request):
+            yield self._response_to_dict(response, ResponsesStreamEvent)