genesis-flow 1.0.0__py3-none-any.whl

mlflow/models/wheeled_model.py

@@ -0,0 +1,280 @@
+import os
+import platform
+import shutil
+import subprocess
+import sys
+
+import yaml
+
+import mlflow
+from mlflow import MlflowClient
+from mlflow.environment_variables import MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import BAD_REQUEST
+from mlflow.pyfunc.model import MLMODEL_FILE_NAME, Model
+from mlflow.store.artifact.utils.models import _parse_model_uri, get_model_name_and_version
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.environment import (
+    _REQUIREMENTS_FILE_NAME,
+    _get_pip_deps,
+    _mlflow_additional_pip_env,
+    _overwrite_pip_deps,
+)
+from mlflow.utils.model_utils import _validate_and_prepare_target_save_path
+from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri
+
+_WHEELS_FOLDER_NAME = "wheels"
+_ORIGINAL_REQ_FILE_NAME = "original_requirements.txt"
+_PLATFORM = "platform"
+
+
+class WheeledModel:
+    """
+    Helper class to create a model with added dependency wheels from an existing registered model.
+    The `wheeled` model contains all the model dependencies as wheels stored as model artifacts.
+    .. note::
+        This utility only operates on a model that has been registered to the Model Registry.
+    """
+
+    def __init__(self, model_uri):
+        self._model_uri = model_uri
+        databricks_profile_uri = (
+            get_databricks_profile_uri_from_artifact_uri(model_uri) or mlflow.get_registry_uri()
+        )
+        client = MlflowClient(registry_uri=databricks_profile_uri)
+        self._model_name, _ = get_model_name_and_version(client, model_uri)
+
+    @classmethod
+    def log_model(cls, model_uri, registered_model_name=None):
+        """
+        Logs a registered model as an MLflow artifact for the current run. This only operates on
+        a model which has been registered to the Model Registry. Given a registered model_uri (
+        e.g. models:/<model_name>/<model_version>), this utility re-logs the model along with all
+        the required model libraries back to the Model Registry. The required model libraries are
+        stored along with the model as model artifacts. In addition, supporting files to the
+        model (e.g. conda.yaml, requirements.txt) are modified to use the added libraries.
+
+        By default, this utility creates a new model version under the same registered model
+        specified by ``model_uri``. This behavior can be overridden by specifying the
+        ``registered_model_name`` argument.
+
+        Args:
+            model_uri: A registered model uri in the Model Registry of the form
+                       models:/<model_name>/<model_version/stage/latest>
+            registered_model_name: The new model version (model with its libraries) is
+                                   registered under the inputted registered_model_name. If None,
+                                   a new version is logged to the existing model in the Model
+                                   Registry.
+
+        .. code-block:: python
+            :caption: Example
+
+            # Given a model uri, log the wheeled model
+            with mlflow.start_run():
+                WheeledModel.log_model(model_uri)
+        """
+        parsed_uri = _parse_model_uri(model_uri)
+        return Model.log(
+            artifact_path=None,
+            flavor=WheeledModel(model_uri),
+            registered_model_name=registered_model_name or parsed_uri.name,
+        )
+
+    def save_model(self, path, mlflow_model=None):
+        """
+        Given an existing registered model, saves the model along with it's dependencies stored as
+        wheels to a path on the local file system.
+
+        This does not modify existing model behavior or existing model flavors. It simply downloads
+        the model dependencies as wheels and modifies the requirements.txt and conda.yaml file to
+        point to the downloaded wheels.
+
+        The download_command defaults to downloading only binary packages using the
+        `--only-binary=:all:` option. This behavior can be overridden using an environment
+        variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows setting
+        different options such as `--prefer-binary`, `--no-binary`, etc.
+
+        Args:
+            path: Local path where the model is to be saved.
+            mlflow_model: The new :py:mod:`mlflow.models.Model` metadata file to store the
+                updated model metadata.
+        """
+        from mlflow.pyfunc import ENV, FLAVOR_NAME, _extract_conda_env
+
+        path = os.path.abspath(path)
+        _validate_and_prepare_target_save_path(path)
+
+        local_model_path = _download_artifact_from_uri(self._model_uri, output_path=path)
+
+        wheels_dir = os.path.join(local_model_path, _WHEELS_FOLDER_NAME)
+        pip_requirements_path = os.path.join(local_model_path, _REQUIREMENTS_FILE_NAME)
+        model_metadata_path = os.path.join(local_model_path, MLMODEL_FILE_NAME)
+
+        model_metadata = Model.load(model_metadata_path)
+
+        # Check if the model file has `wheels` set to True
+        if model_metadata.__dict__.get(_WHEELS_FOLDER_NAME, None) is not None:
+            raise MlflowException("Model libraries are already added", BAD_REQUEST)
+
+        conda_env = _extract_conda_env(model_metadata.flavors.get(FLAVOR_NAME, {}).get(ENV, None))
+        conda_env_path = os.path.join(local_model_path, conda_env)
+        if conda_env is None and not os.path.isfile(pip_requirements_path):
+            raise MlflowException(
+                "Cannot add libraries for model with no logged dependencies.", BAD_REQUEST
+            )
+
+        if not os.path.isfile(pip_requirements_path):
+            self._create_pip_requirement(conda_env_path, pip_requirements_path)
+
+        WheeledModel._download_wheels(
+            pip_requirements_path=pip_requirements_path, dst_path=wheels_dir
+        )
+
+        # Keep a copy of the original requirement.txt
+        shutil.copy2(pip_requirements_path, os.path.join(local_model_path, _ORIGINAL_REQ_FILE_NAME))
+
+        # Update requirements.txt with wheels
+        pip_deps = self._overwrite_pip_requirements_with_wheels(
+            pip_requirements_path=pip_requirements_path, wheels_dir=wheels_dir
+        )
+
+        # Update conda.yaml with wheels
+        self._update_conda_env(pip_deps, conda_env_path)
+
+        # Update MLModel File
+        mlflow_model = self._update_mlflow_model(
+            original_model_metadata=model_metadata, mlflow_model=mlflow_model
+        )
+        mlflow_model.save(model_metadata_path)
+        return mlflow_model
+
+    def _update_conda_env(self, new_pip_deps, conda_env_path):
+        """
+        Updates the list pip packages in the conda.yaml file to the list of wheels in the wheels
+        directory.
+        {
+            "name": "env",
+            "channels": [...],
+            "dependencies": [
+                ...,
+                "pip",
+                {"pip": [...]},  <- Overwrite this with list of wheels
+            ],
+        }
+
+        Args:
+            new_pip_deps: List of pip dependencies as wheels
+            conda_env_path: Path to conda.yaml file in the model directory
+        """
+        with open(conda_env_path) as f:
+            conda_env = yaml.safe_load(f)
+
+        new_conda_env = _overwrite_pip_deps(conda_env, new_pip_deps)
+
+        with open(conda_env_path, "w") as out:
+            yaml.safe_dump(new_conda_env, stream=out, default_flow_style=False)
+
+    def _update_mlflow_model(self, original_model_metadata, mlflow_model):
+        """
+        Modifies the MLModel file to reflect updated information such as the run_id,
+        utc_time_created. Additionally, this also adds `wheels` to the MLModel file to indicate that
+        this is a `wheeled` model.
+
+        Args:
+            original_model_metadata: The model metadata stored in the original MLmodel file.
+            mlflow_model: :py:mod:`mlflow.models.Model` configuration of the newly created
+                          wheeled model
+        """
+
+        run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
+        if mlflow_model is None:
+            mlflow_model = Model(run_id=run_id)
+
+        original_model_metadata.__dict__.update(
+            {k: v for k, v in mlflow_model.__dict__.items() if v}
+        )
+        mlflow_model.__dict__.update(original_model_metadata.__dict__)
+        mlflow_model.artifact_path = WheeledModel.get_wheel_artifact_path(
+            mlflow_model.artifact_path
+        )
+
+        mlflow_model.wheels = {_PLATFORM: platform.platform()}
+        return mlflow_model
+
+    @classmethod
+    def _download_wheels(cls, pip_requirements_path, dst_path):
+        """
+        Downloads all the wheels of the dependencies specified in the requirements.txt file.
+        The pip wheel download_command defaults to downloading only binary packages using
+        the `--only-binary=:all:` option. This behavior can be overridden using an
+        environment variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows
+        setting different options such as `--prefer-binary`, `--no-binary`, etc.
+
+        Args:
+            pip_requirements_path: Path to requirements.txt in the model directory
+            dst_path: Path to the directory where the wheels are to be downloaded
+        """
+        if not os.path.exists(dst_path):
+            os.makedirs(dst_path)
+
+        pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()
+
+        try:
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "wheel",
+                    pip_wheel_options,
+                    "--wheel-dir",
+                    dst_path,
+                    "-r",
+                    pip_requirements_path,
+                    "--no-cache-dir",
+                    "--progress-bar=off",
+                ],
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+            )
+        except subprocess.CalledProcessError as e:
+            raise MlflowException(
+                f"An error occurred while downloading the dependency wheels: {e.stdout}"
+            )
+
+    def _overwrite_pip_requirements_with_wheels(self, pip_requirements_path, wheels_dir):
+        """
+        Overwrites the requirements.txt with the wheels of the required dependencies.
+
+        Args:
+            pip_requirements_path: Path to requirements.txt in the model directory.
+            wheels_dir: Path to directory where wheels are stored.
+        """
+        wheels = []
+        with open(pip_requirements_path, "w") as wheels_requirements:
+            for wheel_file in os.listdir(wheels_dir):
+                if wheel_file.endswith(".whl"):
+                    complete_wheel_file = os.path.join(_WHEELS_FOLDER_NAME, wheel_file)
+                    wheels.append(complete_wheel_file)
+                    wheels_requirements.write(complete_wheel_file + "\n")
+        return wheels
+
+    def _create_pip_requirement(self, conda_env_path, pip_requirements_path):
+        """
+        This method creates a requirements.txt file for the model dependencies if the file does not
+        already exist. It uses the pip dependencies found in the conda.yaml env file.
+
+        Args:
+            conda_env_path: Path to conda.yaml env file which contains the required pip
+                dependencies
+            pip_requirements_path: Path where the new requirements.txt will be created.
+        """
+        with open(conda_env_path) as f:
+            conda_env = yaml.safe_load(f)
+        pip_deps = _get_pip_deps(conda_env)
+        _mlflow_additional_pip_env(pip_deps, pip_requirements_path)
+
+    @classmethod
+    def get_wheel_artifact_path(cls, original_artifact_path):
+        return original_artifact_path + "_" + _WHEELS_FOLDER_NAME

mlflow/openai/__init__.py

@@ -0,0 +1,57 @@
+"""
+The ``mlflow.openai`` module provides an API for logging and loading OpenAI models.
+
+Credential management for OpenAI on Databricks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. warning::
+
+    Specifying secrets for model serving with ``MLFLOW_OPENAI_SECRET_SCOPE`` is deprecated.
+    Use `secrets-based environment variables <https://docs.databricks.com/en/machine-learning/model-serving/store-env-variable-model-serving.html>`_
+    instead.
+
+When this flavor logs a model on Databricks, it saves a YAML file with the following contents as
+``openai.yaml`` if the ``MLFLOW_OPENAI_SECRET_SCOPE`` environment variable is set.
+
+.. code-block:: yaml
+
+    OPENAI_API_BASE: {scope}:openai_api_base
+    OPENAI_API_KEY: {scope}:openai_api_key
+    OPENAI_API_KEY_PATH: {scope}:openai_api_key_path
+    OPENAI_API_TYPE: {scope}:openai_api_type
+    OPENAI_ORGANIZATION: {scope}:openai_organization
+
+- ``{scope}`` is the value of the ``MLFLOW_OPENAI_SECRET_SCOPE`` environment variable.
+- The keys are the environment variables that the ``openai-python`` package uses to
+  configure the API client.
+- The values are the references to the secrets that store the values of the environment
+  variables.
+
+When the logged model is served on Databricks, each secret will be resolved and set as the
+corresponding environment variable. See https://docs.databricks.com/security/secrets/index.html
+for how to set up secrets on Databricks.
+"""
+
+from mlflow.openai.autolog import autolog
+from mlflow.openai.constant import FLAVOR_NAME
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+__all__ = ["autolog", "FLAVOR_NAME"]
+
+
+# Import model logging APIs only if mlflow skinny or full package is installed,
+# i.e., skip if only mlflow-tracing package is installed.
+if not IS_TRACING_SDK_ONLY:
+    from mlflow.openai.model import (
+        _load_pyfunc,
+        load_model,
+        log_model,
+        save_model,
+    )
+
+    __all__ += [
+        "load_model",
+        "log_model",
+        "save_model",
+        "_load_pyfunc",
+    ]

mlflow/openai/_agent_tracer.py

@@ -0,0 +1,364 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, Optional
+
+import agents.tracing as oai
+from agents import add_trace_processor
+from agents._run_impl import TraceCtxManager
+from agents.tracing.setup import GLOBAL_TRACE_PROVIDER
+from pydantic import BaseModel
+
+from mlflow.entities.span import LiveSpan, SpanType
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatus, SpanStatusCode
+from mlflow.openai import FLAVOR_NAME
+from mlflow.tracing.constant import SpanAttributeKey
+from mlflow.tracing.fluent import start_span_no_context
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionToolDefinition,
+    TextContentPart,
+    ToolCall,
+)
+from mlflow.utils.autologging_utils.safety import safe_patch
+
+_logger = logging.getLogger(__name__)
+
+
+class OpenAISpanType:
+    """
+    https://github.com/openai/openai-agents-python/blob/ca8e8bed5d0f33e8a0bc3eabd5f1b0a183e73765/src/agents/tracing/span_data.py#L11
+    """
+
+    AGENT = "agent"
+    FUNCTION = "function"
+    GENERATION = "generation"
+    RESPONSE = "response"
+    HANDOFF = "handoff"
+    CUSTOM = "custom"
+    GUARDRAIL = "guardrail"
+
+
+_SPAN_TYPE_MAP = {
+    OpenAISpanType.AGENT: SpanType.AGENT,
+    OpenAISpanType.FUNCTION: SpanType.TOOL,
+    OpenAISpanType.GENERATION: SpanType.CHAT_MODEL,
+    OpenAISpanType.RESPONSE: SpanType.CHAT_MODEL,
+    OpenAISpanType.GUARDRAIL: SpanType.TOOL,
+    # Default to chain type
+}
+
+
+def add_mlflow_trace_processor():
+    processors = GLOBAL_TRACE_PROVIDER._multi_processor._processors
+
+    if any(isinstance(p, MlflowOpenAgentTracingProcessor) for p in processors):
+        return
+
+    add_trace_processor(MlflowOpenAgentTracingProcessor())
+
+
+def remove_mlflow_trace_processor():
+    processors = GLOBAL_TRACE_PROVIDER._multi_processor._processors
+    non_mlflow_processors = [
+        p for p in processors if not isinstance(p, MlflowOpenAgentTracingProcessor)
+    ]
+    GLOBAL_TRACE_PROVIDER._multi_processor._processors = non_mlflow_processors
+
+
+class MlflowOpenAgentTracingProcessor(oai.TracingProcessor):
+    def __init__(
+        self,
+        project_name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._span_id_to_mlflow_span: dict[str, LiveSpan] = {}
+        self._project_name = project_name
+
+        # Patch TraceCtxManager to handle exceptions from the agent properly
+        # The original implementation does not propagate exception to the root span,
+        # resulting in the trace to have status OK even if there is an exception.
+        def _patched_exit(original, instance, exc_type, exc_val, exc_tb):
+            try:
+                if exc_val and instance.trace:
+                    span = self._span_id_to_mlflow_span.get(instance.trace.trace_id)
+                    span.add_event(SpanEvent.from_exception(exc_val))
+                    span.set_status(SpanStatusCode.ERROR)
+            except Exception:
+                _logger.debug("Failed to handle exception in MLflow trace", exc_info=True)
+
+            return original(instance, exc_type, exc_val, exc_tb)
+
+        safe_patch(
+            FLAVOR_NAME,
+            TraceCtxManager,
+            "__exit__",
+            _patched_exit,
+        )
+
+    def on_trace_start(self, trace: oai.Trace) -> None:
+        try:
+            mlflow_span = start_span_no_context(
+                name=trace.name,
+                span_type=SpanType.AGENT,
+                # TODO: Trace object doesn't contain input/output. Can we get it somehow?
+                inputs="",
+                attributes=trace.metadata,
+            )
+            # NB: Trace ID has different prefix as span ID so will not conflict
+            self._span_id_to_mlflow_span[trace.trace_id] = mlflow_span
+
+            if trace.group_id:
+                # Group ID is used for grouping multiple agent executions together
+                mlflow_span.set_tag("group_id", trace.group_id)
+
+            original_exit = trace.__exit__
+
+            # Patch __exit__ method to handle exception properly
+            def _patched_exit(self, exc_type, exc_val, exc_tb):
+                if exc_val:
+                    mlflow_span.add_event(SpanEvent.from_exception(exc_val))
+                    mlflow_span.set_status(SpanStatusCode.ERROR)
+
+                original_exit(exc_type, exc_val, exc_tb)
+
+            safe_patch(
+                FLAVOR_NAME,
+                trace.__class__,
+                "__exit__",
+                _patched_exit,
+            )
+
+        except Exception:
+            _logger.debug("Failed to start MLflow trace", exc_info=True)
+
+    def on_trace_end(self, trace: oai.Trace) -> None:
+        try:
+            mlflow_span = self._span_id_to_mlflow_span.pop(trace.trace_id, None)
+            mlflow_span.end(status=mlflow_span.status, outputs="")
+        except Exception:
+            _logger.debug("Failed to end MLflow trace", exc_info=True)
+
+    def on_span_start(self, span: oai.Span[Any]) -> None:
+        try:
+            parent_mlflow_span = self._span_id_to_mlflow_span.get(span.parent_id)
+
+            # Parent might be a trace
+            if not parent_mlflow_span:
+                parent_mlflow_span = self._span_id_to_mlflow_span.get(span.trace_id)
+
+            inputs, _, attributes = _parse_span_data(span.span_data)
+
+            mlflow_span = start_span_no_context(
+                name=_get_span_name(span.span_data),
+                span_type=_SPAN_TYPE_MAP.get(span.span_data.type, SpanType.CHAIN),
+                parent_span=parent_mlflow_span,
+                inputs=inputs,
+                attributes=attributes,
+            )
+            self._span_id_to_mlflow_span[span.span_id] = mlflow_span
+        except Exception:
+            _logger.debug("Failed to start MLflow span", exc_info=True)
+
+    def on_span_end(self, span: oai.Span[Any]) -> None:
+        try:
+            # parsed_span_data = parse_spandata(span.span_data)
+            mlflow_span = self._span_id_to_mlflow_span.pop(span.span_id, None)
+
+            inputs, outputs, attributes = _parse_span_data(span.span_data)
+
+            mlflow_span.set_inputs(inputs)
+            mlflow_span.set_outputs(outputs)
+            mlflow_span.set_attributes(attributes)
+
+            if span.error:
+                status = SpanStatus(
+                    status_code=SpanStatusCode.ERROR,
+                    description=span.error["message"],
+                )
+                mlflow_span.add_event(
+                    SpanEvent(
+                        name="exception",
+                        attributes={
+                            "exception.message": span.error["message"],
+                            "exception.type": "",
+                            "exception.stacktrace": json.dumps(span.error["data"]),
+                        },
+                    )
+                )
+            else:
+                status = SpanStatusCode.OK
+
+            mlflow_span.end(status=status)
+        except Exception:
+            _logger.debug("Failed to end MLflow span", exc_info=True)
+
+    def force_flush(self) -> None:
+        # MLflow doesn't need flush but this method is required by the interface
+        pass
+
+    def shutdown(self) -> None:
+        self.force_flush()
+
+
+def _get_span_name(span_data: oai.SpanData) -> str:
+    if hasattr(span_data, "name"):
+        return span_data.name
+    elif isinstance(span_data, oai.GenerationSpanData):
+        return "Generation"
+    elif isinstance(span_data, oai.ResponseSpanData):
+        return "Response"
+    elif isinstance(span_data, oai.HandoffSpanData):
+        return "Handoff"
+    else:
+        return "Unknown"
+
+
+def _parse_span_data(span_data: oai.SpanData) -> tuple[Any, Any, dict[str, Any]]:
+    inputs = None
+    outputs = None
+    attributes = {}
+
+    if span_data.type == OpenAISpanType.AGENT:
+        attributes = {
+            "handoffs": span_data.handoffs,
+            "tools": span_data.tools,
+            "output_type": span_data.output_type,
+        }
+        outputs = {"output_type": span_data.output_type}
+
+    elif span_data.type == OpenAISpanType.FUNCTION:
+        try:
+            inputs = json.loads(span_data.input)
+        except Exception:
+            inputs = span_data.input
+        outputs = span_data.output
+
+    elif span_data.type == OpenAISpanType.GENERATION:
+        inputs = span_data.input
+        outputs = span_data.output
+        attributes = {
+            "model": span_data.model,
+            "model_config": span_data.model_config,
+            "usage": span_data.usage,
+        }
+
+    elif span_data.type == OpenAISpanType.RESPONSE:
+        inputs, outputs, attributes = _parse_response_span_data(span_data)
+
+    elif span_data.type == OpenAISpanType.HANDOFF:
+        inputs = {"from_agent": span_data.from_agent}
+        outputs = {"to_agent": span_data.to_agent}
+
+    elif span_data.type == OpenAISpanType.CUSTOM:
+        outputs = span_data.data
+
+    elif span_data.type == OpenAISpanType.GUARDRAIL:
+        outputs = {"triggered": span_data.triggered}
+
+    return inputs, outputs, attributes
+
+
+def _parse_response_span_data(span_data: oai.ResponseSpanData) -> tuple[Any, Any, dict[str, Any]]:
+    inputs = span_data.input
+    response = span_data.response
+    response_dict = response.model_dump() if response else {}
+    outputs = response_dict.get("output")
+    attributes = {k: v for k, v in response_dict.items() if k != "output"}
+
+    # Extract chat messages
+    messages = []
+    if response and response.instructions:
+        messages.append(ChatMessage(role="system", content=span_data.response.instructions))
+    if span_data.input:
+        parsed = [_parse_message_like(m) for m in span_data.input]
+        messages.extend([m for m in parsed if m is not None])
+    if response and response.output:
+        parsed = [_parse_message_like(m) for m in span_data.response.output]
+        messages.extend(parsed)
+    attributes[SpanAttributeKey.CHAT_MESSAGES] = [m.model_dump_compat() for m in messages]
+
+    # Extract chat tools
+    chat_tools = []
+    for tool in response_dict.get("tools", []):
+        try:
+            tool = ChatTool(
+                type="function",
+                function=FunctionToolDefinition(
+                    name=tool["name"],
+                    description=tool.get("description"),
+                    parameters=tool.get("parameters"),
+                    strict=tool.get("strict"),
+                ),
+            )
+            chat_tools.append(tool)
+        except Exception as e:
+            _logger.debug(f"Failed to parse chat tool: {tool}. Error: {e}")
+
+    if chat_tools:
+        attributes[SpanAttributeKey.CHAT_TOOLS] = chat_tools
+
+    return inputs, outputs, attributes
+
+
+def _parse_message_like(message_like: Any) -> Optional[ChatMessage]:
+    try:
+        return ChatMessage.validate_compat(message_like)
+    except Exception:
+        pass
+
+    if isinstance(message_like, BaseModel):
+        message_like = message_like.model_dump()
+
+    msg_type = message_like["type"]
+    if msg_type == "message":
+        content = []
+        refusal = None
+        for content_block in message_like["content"]:
+            # Content is a list of either text or refusal https://github.com/openai/openai-python/blob/9dea82fb8cdd06683f9e8033b54cff219789af7f/src/openai/types/responses/response_output_message.py#L13C38-L13C56
+            if "text" in content_block:
+                content.append(TextContentPart(type="text", text=content_block["text"]))
+            elif "refusal" in content_block:
+                refusal = content_block["refusal"]
+            else:
+                _logger.debug(f"Unknown content type in message: {content_block}")
+        return ChatMessage(
+            role=message_like["role"],
+            content=content,
+            refusal=refusal,
+        )
+    elif msg_type == "function_call":
+        return ChatMessage(
+            role="assistant",
+            content="",
+            tool_calls=[
+                ToolCall(
+                    id=message_like["call_id"],
+                    function=Function(
+                        name=message_like["name"],
+                        arguments=message_like["arguments"],
+                    ),
+                )
+            ],
+        )
+    elif msg_type == "function_call_output":
+        return ChatMessage(
+            role="tool",
+            content=message_like["output"],
+            tool_call_id=message_like["call_id"],
+        )
+
+    # Ignore unknown message types.
+    # Response API supports the following additional message types, which is not
+    # supported by our chat standard schema yet:
+    # https://github.com/openai/openai-python/blob/9dea82fb8cdd06683f9e8033b54cff219789af7f/src/openai/types/responses/response_output_item.py#L16
+    # - File search tool call
+    # - Web search tool call
+    # - Computer tool call
+    # - Reasoning
+    _logger.debug(f"Unknown message type: {msg_type}")