star

venv/lib/python3.9/site-packages/mlflow/openai/_agent_tracer.py

@@ -0,0 +1,377 @@
+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 import MlflowClient
+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.utils import end_client_span_or_trace, start_client_span_or_trace
+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/main/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
+        self._mlflow_client = MlflowClient()
+
+        # 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_client_span_or_trace(
+                client=self._mlflow_client,
+                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)
+            end_client_span_or_trace(
+                client=self._mlflow_client,
+                span=mlflow_span,
+                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_client_span_or_trace(
+                client=self._mlflow_client,
+                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
+
+            end_client_span_or_trace(
+                client=self._mlflow_client,
+                span=mlflow_span,
+                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}")

venv/lib/python3.9/site-packages/mlflow/openai/_openai_autolog.py

@@ -0,0 +1,457 @@
+import functools
+import json
+import logging
+import os
+from contextlib import contextmanager
+from copy import deepcopy
+from typing import Any, AsyncIterator, Iterator, Optional
+
+from packaging.version import Version
+
+import mlflow
+from mlflow import MlflowException
+from mlflow.entities import RunTag, SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
+from mlflow.openai.utils.chat_schema import set_span_chat_attributes
+from mlflow.tracing.assessment import MlflowClient
+from mlflow.tracing.constant import (
+    STREAM_CHUNK_EVENT_NAME_FORMAT,
+    STREAM_CHUNK_EVENT_VALUE_KEY,
+    TraceMetadataKey,
+)
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import (
+    TraceJSONEncoder,
+    end_client_span_or_trace,
+    start_client_span_or_trace,
+)
+from mlflow.tracking.context import registry as context_registry
+from mlflow.tracking.fluent import _get_experiment_id
+from mlflow.utils.autologging_utils import disable_autologging, get_autologging_config
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+from mlflow.utils.autologging_utils.safety import _resolve_extra_tags
+
+MIN_REQ_VERSION = Version(_ML_PACKAGE_VERSIONS["openai"]["autologging"]["minimum"])
+MAX_REQ_VERSION = Version(_ML_PACKAGE_VERSIONS["openai"]["autologging"]["maximum"])
+
+_logger = logging.getLogger(__name__)
+
+
+def _get_input_from_model(model, kwargs):
+    from openai.resources.chat.completions import Completions as ChatCompletions
+    from openai.resources.completions import Completions
+    from openai.resources.embeddings import Embeddings
+
+    model_class_param_name_mapping = {
+        ChatCompletions: "messages",
+        Completions: "prompt",
+        Embeddings: "input",
+    }
+    if param_name := model_class_param_name_mapping.get(model.__class__):
+        # openai tasks accept only keyword arguments
+        if param := kwargs.get(param_name):
+            return param
+        input_example_exc = MlflowException(
+            "Inference function signature changes, please contact MLflow team to "
+            "fix OpenAI autologging.",
+        )
+    else:
+        input_example_exc = MlflowException(
+            "Unsupported OpenAI task. Only support chat completions, completions and embeddings."
+        )
+    _logger.warning(
+        f"Failed to gather input example of model {model.__class__.__name__} "
+        f"due to error: {input_example_exc}"
+    )
+
+
+@contextmanager
+def _set_api_key_env_var(client):
+    """
+    Gets the API key from the client and temporarily set it as an environment variable
+    """
+    api_key = client.api_key
+    original = os.environ.get("OPENAI_API_KEY", None)
+    os.environ["OPENAI_API_KEY"] = api_key
+    yield
+    if original is not None:
+        os.environ["OPENAI_API_KEY"] = original
+    else:
+        os.environ.pop("OPENAI_API_KEY")
+
+
+def _get_span_type(task: type) -> str:
+    from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
+    from openai.resources.chat.completions import Completions as ChatCompletions
+    from openai.resources.completions import AsyncCompletions, Completions
+    from openai.resources.embeddings import AsyncEmbeddings, Embeddings
+
+    span_type_mapping = {
+        ChatCompletions: SpanType.CHAT_MODEL,
+        AsyncChatCompletions: SpanType.CHAT_MODEL,
+        Completions: SpanType.LLM,
+        AsyncCompletions: SpanType.LLM,
+        Embeddings: SpanType.EMBEDDING,
+        AsyncEmbeddings: SpanType.EMBEDDING,
+    }
+
+    try:
+        # Only available in openai>=1.40.0
+        from openai.resources.beta.chat.completions import (
+            AsyncCompletions as BetaAsyncChatCompletions,
+        )
+        from openai.resources.beta.chat.completions import Completions as BetaChatCompletions
+
+        span_type_mapping[BetaChatCompletions] = SpanType.CHAT_MODEL
+        span_type_mapping[BetaAsyncChatCompletions] = SpanType.CHAT_MODEL
+    except ImportError:
+        pass
+
+    try:
+        # Responses API only available in openai>=1.66.0
+        from openai.resources.responses import AsyncResponses, Responses
+
+        span_type_mapping[Responses] = SpanType.CHAT_MODEL
+        span_type_mapping[AsyncResponses] = SpanType.CHAT_MODEL
+    except ImportError:
+        pass
+
+    return span_type_mapping.get(task, SpanType.UNKNOWN)
+
+
+def _try_parse_raw_response(response: Any) -> Any:
+    """
+    As documented at https://github.com/openai/openai-python/tree/52357cff50bee57ef442e94d78a0de38b4173fc2?tab=readme-ov-file#accessing-raw-response-data-eg-headers,
+    a `LegacyAPIResponse` (https://github.com/openai/openai-python/blob/52357cff50bee57ef442e94d78a0de38b4173fc2/src/openai/_legacy_response.py#L45)
+    object is returned when the `create` method is invoked with `with_raw_response`.
+    """
+    try:
+        from openai._legacy_response import LegacyAPIResponse
+    except ImportError:
+        _logger.debug("Failed to import `LegacyAPIResponse` from `openai._legacy_response`")
+        return response
+    if isinstance(response, LegacyAPIResponse):
+        try:
+            # `parse` returns either a `pydantic.BaseModel` or a `openai.Stream` object
+            # depending on whether the request has a `stream` parameter set to `True`.
+            return response.parse()
+        except Exception as e:
+            _logger.debug(f"Failed to parse {response} (type: {response.__class__}): {e}")
+
+    return response
+
+
+def patched_call(original, self, *args, **kwargs):
+    config = AutoLoggingConfig.init(flavor_name=mlflow.openai.FLAVOR_NAME)
+    active_run = mlflow.active_run()
+    run_id = _get_autolog_run_id(self, active_run)
+    mlflow_client = mlflow.MlflowClient()
+
+    # If optional artifacts logging are enabled e.g. log_models, we need to create a run
+    if config.should_log_optional_artifacts():
+        run_id = _start_run_or_log_tag(mlflow_client, config, run_id)
+
+    if config.log_traces:
+        span = _start_span(mlflow_client, self, kwargs, run_id)
+
+    # Execute the original function
+    try:
+        raw_result = original(self, *args, **kwargs)
+    except Exception as e:
+        if config.log_traces:
+            _end_span_on_exception(mlflow_client, span, e)
+        raise
+
+    if config.log_traces:
+        _end_span_on_success(mlflow_client, span, kwargs, raw_result)
+
+    if config.should_log_optional_artifacts():
+        _log_optional_artifacts(config, run_id, self, kwargs)
+
+    # Even if the model is not logged, we keep a single run per model
+    self._mlflow_run_id = run_id
+
+    # Terminate the run if it is not managed by the user
+    if run_id is not None and (active_run is None or active_run.info.run_id != run_id):
+        mlflow_client.set_terminated(run_id)
+
+    return raw_result
+
+
+async def async_patched_call(original, self, *args, **kwargs):
+    config = AutoLoggingConfig.init(flavor_name=mlflow.openai.FLAVOR_NAME)
+    active_run = mlflow.active_run()
+    run_id = _get_autolog_run_id(self, active_run)
+    mlflow_client = mlflow.MlflowClient()
+
+    # If optional artifacts logging are enabled e.g. log_models, we need to create a run
+    if config.should_log_optional_artifacts():
+        run_id = _start_run_or_log_tag(mlflow_client, config, run_id)
+
+    if config.log_traces:
+        span = _start_span(mlflow_client, self, kwargs, run_id)
+
+    # Execute the original function
+    try:
+        raw_result = await original(self, *args, **kwargs)
+    except Exception as e:
+        if config.log_traces:
+            _end_span_on_exception(mlflow_client, span, e)
+        raise
+
+    if config.log_traces:
+        _end_span_on_success(mlflow_client, span, kwargs, raw_result)
+
+    if config.should_log_optional_artifacts():
+        _log_optional_artifacts(config, run_id, self, kwargs)
+
+    # Even if the model is not logged, we keep a single run per model
+    self._mlflow_run_id = run_id
+
+    # Terminate the run if it is not managed by the user
+    if run_id is not None and (active_run is None or active_run.info.run_id != run_id):
+        mlflow_client.set_terminated(run_id)
+
+    return raw_result
+
+
+def _get_autolog_run_id(instance, active_run):
+    """
+    Get the run ID to use for logging artifacts and associate with the trace.
+
+    The run ID is determined as follows:
+    - If there is an active run (created by a user), use its run ID.
+    - If the model has a `_mlflow_run_id` attribute, use it. This is the run ID created
+        by autologging in a previous call to the same model.
+    """
+    return active_run.info.run_id if active_run else getattr(instance, "_mlflow_run_id", None)
+
+
+def _start_run_or_log_tag(
+    mlflow_client: MlflowClient, config: AutoLoggingConfig, run_id: Optional[str]
+) -> str:
+    """Start a new run or log models, or log extra tags if a run is already active."""
+    # include run context tags
+    resolved_tags = context_registry.resolve_tags(config.extra_tags)
+    tags = _resolve_extra_tags(mlflow.openai.FLAVOR_NAME, resolved_tags)
+    if run_id is not None:
+        mlflow_client.log_batch(
+            run_id=run_id,
+            tags=[RunTag(key, str(value)) for key, value in tags.items()],
+        )
+    else:
+        run = mlflow_client.create_run(
+            experiment_id=_get_experiment_id(),
+            tags=tags,
+        )
+        run_id = run.info.run_id
+    return run_id
+
+
+def _log_optional_artifacts(
+    config: AutoLoggingConfig, run_id: str, instance: Any, kwargs: dict[str, Any]
+):
+    if hasattr(instance, "_mlflow_model_logged"):
+        # Model is already logged for this instance, no need to log again
+        return
+
+    input_example = None
+    if config.log_input_examples:
+        input_example = deepcopy(_get_input_from_model(instance, kwargs))
+        if not config.log_model_signatures:
+            _logger.info(
+                "Signature is automatically generated for logged model if "
+                "input_example is provided. To disable log_model_signatures, "
+                "please also disable log_input_examples."
+            )
+
+    registered_model_name = get_autologging_config(
+        mlflow.openai.FLAVOR_NAME, "registered_model_name", None
+    )
+    try:
+        task = mlflow.openai._get_task_name(instance.__class__)
+        with disable_autologging():
+            # If the user is using `openai.OpenAI()` client,
+            # they do not need to set the "OPENAI_API_KEY" environment variable.
+            # This temporarily sets the API key as an environment variable
+            # so that the model can be logged.
+            with _set_api_key_env_var(instance._client):
+                mlflow.openai.log_model(
+                    kwargs.get("model"),
+                    task,
+                    "model",
+                    input_example=input_example,
+                    registered_model_name=registered_model_name,
+                    run_id=run_id,
+                )
+    except Exception as e:
+        _logger.warning(f"Failed to log model due to error: {e}")
+
+    # Even if the model is not logged, we keep a single run per model
+    instance._mlflow_model_logged = True
+
+
+def _start_span(mlflow_client: MlflowClient, instance: Any, inputs: dict[str, Any], run_id: str):
+    # Record input parameters to attributes
+    attributes = {k: v for k, v in inputs.items() if k not in ("messages", "input")}
+
+    # If there is an active span, create a child span under it, otherwise create a new trace
+    span = start_client_span_or_trace(
+        mlflow_client,
+        name=instance.__class__.__name__,
+        span_type=_get_span_type(instance.__class__),
+        inputs=inputs,
+        attributes=attributes,
+    )
+
+    # Associate run ID to the trace manually, because if a new run is created by
+    # autologging, it is not set as the active run thus not automatically
+    # associated with the trace.
+    if run_id is not None:
+        tm = InMemoryTraceManager().get_instance()
+        tm.set_request_metadata(span.request_id, TraceMetadataKey.SOURCE_RUN, run_id)
+
+    return span
+
+
+def _end_span_on_success(
+    mlflow_client: MlflowClient, span: LiveSpan, inputs: dict[str, Any], raw_result: Any
+):
+    from openai import AsyncStream, Stream
+
+    result = _try_parse_raw_response(raw_result)
+
+    if isinstance(result, Stream):
+        # If the output is a stream, we add a hook to store the intermediate chunks
+        # and then log the outputs as a single artifact when the stream ends
+        def _stream_output_logging_hook(stream: Iterator) -> Iterator:
+            output = []
+            for i, chunk in enumerate(stream):
+                output.append(_process_chunk(span, i, chunk))
+                yield chunk
+            output = chunk.response if _is_responses_final_event(chunk) else "".join(output)
+            _end_span_on_success(mlflow_client, span, inputs, output)
+
+        result._iterator = _stream_output_logging_hook(result._iterator)
+    elif isinstance(result, AsyncStream):
+
+        async def _stream_output_logging_hook(stream: AsyncIterator) -> AsyncIterator:
+            output = []
+            async for chunk in stream:
+                output.append(_process_chunk(span, len(output), chunk))
+                yield chunk
+            output = chunk.response if _is_responses_final_event(chunk) else "".join(output)
+            _end_span_on_success(mlflow_client, span, inputs, output)
+
+        result._iterator = _stream_output_logging_hook(result._iterator)
+    else:
+        try:
+            set_span_chat_attributes(span, inputs, result)
+            end_client_span_or_trace(mlflow_client, span, outputs=result)
+        except Exception as e:
+            _logger.warning(f"Encountered unexpected error when ending trace: {e}", exc_info=True)
+
+
+def _is_responses_final_event(chunk: Any) -> bool:
+    try:
+        from openai.types.responses import ResponseCompletedEvent
+
+        return isinstance(chunk, ResponseCompletedEvent)
+    except ImportError:
+        return False
+
+
+def _end_span_on_exception(mlflow_client: MlflowClient, span: LiveSpan, e: Exception):
+    try:
+        span.add_event(SpanEvent.from_exception(e))
+        mlflow_client.end_span(span.request_id, span.span_id, status=SpanStatusCode.ERROR)
+    except Exception as inner_e:
+        _logger.warning(f"Encountered unexpected error when ending trace: {inner_e}")
+
+
+def _process_chunk(span: LiveSpan, index: int, chunk: Any) -> str:
+    """Parse the chunk and log it as a span event in the trace."""
+    from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+    from openai.types.completion import Completion
+
+    # `chunk.choices` can be empty: https://github.com/mlflow/mlflow/issues/13361
+    if isinstance(chunk, Completion) and chunk.choices:
+        parsed = chunk.choices[0].text or ""
+    elif isinstance(chunk, ChatCompletionChunk) and chunk.choices:
+        parsed = chunk.choices[0].delta.content or ""
+    else:
+        parsed = ""
+
+    span.add_event(
+        SpanEvent(
+            name=STREAM_CHUNK_EVENT_NAME_FORMAT.format(index=index),
+            # OpenTelemetry SpanEvent only support str-str key-value pairs for attributes
+            attributes={STREAM_CHUNK_EVENT_VALUE_KEY: json.dumps(chunk, cls=TraceJSONEncoder)},
+        )
+    )
+    return parsed
+
+
+def patched_agent_get_chat_completion(original, self, *args, **kwargs):
+    """
+    Patch the `get_chat_completion` method of the ChatCompletion object.
+    OpenAI autolog already handles the raw completion request, but tracing
+    the swarm's method is useful to track other parameters like agent name.
+    """
+    agent = kwargs.get("agent") or args[0]
+
+    # Patch agent's functions to generate traces. Function calls only happen
+    # after the first completion is generated because of the design of
+    # function calling. Therefore, we can safely patch the tool functions here
+    # within get_chat_completion() hook.
+    # We cannot patch functions during the agent's initialization because the
+    # agent's functions can be modified after the agent is created.
+    def function_wrapper(fn):
+        if "context_variables" in fn.__code__.co_varnames:
+
+            def wrapper(*args, **kwargs):
+                # NB: Swarm uses `func.__code__.co_varnames` to inspect if the provided
+                # tool function includes 'context_variables' parameter in the signature
+                # and ingest the global context variables if so. Wrapping the function
+                # with mlflow.trace() will break this.
+                # The co_varnames is determined based on the local variables of the
+                # function, so we workaround this by declaring it here as a local variable.
+                context_variables = kwargs.get("context_variables", {})  # noqa: F841
+                return mlflow.trace(
+                    fn,
+                    name=f"{agent.name}.{fn.__name__}",
+                    span_type=SpanType.TOOL,
+                )(*args, **kwargs)
+        else:
+
+            def wrapper(*args, **kwargs):
+                return mlflow.trace(
+                    fn,
+                    name=f"{agent.name}.{fn.__name__}",
+                    span_type=SpanType.TOOL,
+                )(*args, **kwargs)
+
+        wrapped = functools.wraps(fn)(wrapper)
+        wrapped._is_mlflow_traced = True  # Marker to avoid double tracing
+        return wrapped
+
+    agent.functions = [
+        function_wrapper(fn) if not hasattr(fn, "_is_mlflow_traced") else fn
+        for fn in agent.functions
+    ]
+
+    traced_fn = mlflow.trace(
+        original, name=f"{agent.name}.get_chat_completion", span_type=SpanType.CHAIN
+    )
+    return traced_fn(self, *args, **kwargs)
+
+
+def patched_swarm_run(original, self, *args, **kwargs):
+    """
+    Patched version of `run` method of the Swarm object.
+    """
+    traced_fn = mlflow.trace(original, span_type=SpanType.AGENT)
+    return traced_fn(self, *args, **kwargs)

venv/lib/python3.9/site-packages/mlflow/openai/api_request_parallel_processor.py

@@ -0,0 +1,126 @@
+# Based ons: https://github.com/openai/openai-cookbook/blob/6df6ceff470eeba26a56de131254e775292eac22/examples/api_request_parallel_processor.py
+# Several changes were made to make it work with MLflow.
+
+"""
+API REQUEST PARALLEL PROCESSOR
+
+Using the OpenAI API to process lots of text quickly takes some care.
+If you trickle in a million API requests one by one, they'll take days to complete.
+If you flood a million API requests in parallel, they'll exceed the rate limits and fail with
+errors. To maximize throughput, parallel requests need to be throttled to stay under rate limits.
+
+This script parallelizes requests to the OpenAI API
+
+Features:
+- Makes requests concurrently, to maximize throughput
+- Retries failed requests up to {max_attempts} times, to avoid missing data
+- Logs errors, to diagnose problems with requests
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import mlflow
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StatusTracker:
+    """Stores metadata about the script's progress. Only one instance is created."""
+
+    num_tasks_started: int = 0
+    num_tasks_in_progress: int = 0  # script ends when this reaches 0
+    num_tasks_succeeded: int = 0
+    num_tasks_failed: int = 0
+    num_rate_limit_errors: int = 0
+    lock: threading.Lock = threading.Lock()
+    error = None
+
+    def start_task(self):
+        with self.lock:
+            self.num_tasks_started += 1
+            self.num_tasks_in_progress += 1
+
+    def complete_task(self, *, success: bool):
+        with self.lock:
+            self.num_tasks_in_progress -= 1
+            if success:
+                self.num_tasks_succeeded += 1
+            else:
+                self.num_tasks_failed += 1
+
+    def increment_num_rate_limit_errors(self):
+        with self.lock:
+            self.num_rate_limit_errors += 1
+
+
+def call_api(
+    index: int, results: list[tuple[int, Any]], task: Callable, status_tracker: StatusTracker
+):
+    import openai
+
+    status_tracker.start_task()
+    try:
+        result = task()
+        _logger.debug(f"Request #{index} succeeded")
+        status_tracker.complete_task(success=True)
+        results.append((index, result))
+    except openai.RateLimitError as e:
+        status_tracker.complete_task(success=False)
+        _logger.debug(f"Request #{index} failed with: {e}")
+        status_tracker.increment_num_rate_limit_errors()
+        status_tracker.error = mlflow.MlflowException(
+            f"Request #{index} failed with rate limit: {e}."
+        )
+    except Exception as e:
+        status_tracker.complete_task(success=False)
+        _logger.debug(f"Request #{index} failed with: {e}")
+        status_tracker.error = mlflow.MlflowException(
+            f"Request #{index} failed with: {e.__cause__}"
+        )
+
+
+def process_api_requests(
+    request_tasks: list[Callable[[], Any]],
+    max_workers: int = 10,
+):
+    """Processes API requests in parallel"""
+    # initialize trackers
+    status_tracker = StatusTracker()  # single instance to track a collection of variables
+
+    results: list[tuple[int, Any]] = []
+    request_tasks_iter = enumerate(request_tasks)
+    _logger.debug(f"Request pool executor will run {len(request_tasks)} requests")
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(
+                call_api,
+                index=index,
+                task=task,
+                results=results,
+                status_tracker=status_tracker,
+            )
+            for index, task in request_tasks_iter
+        ]
+        wait(futures, return_when=FIRST_EXCEPTION)
+
+    # after finishing, log final status
+    if status_tracker.num_tasks_failed > 0:
+        if status_tracker.num_tasks_failed == 1:
+            raise status_tracker.error
+        raise mlflow.MlflowException(
+            f"{status_tracker.num_tasks_failed} tasks failed. See logs for details."
+        )
+    if status_tracker.num_rate_limit_errors > 0:
+        _logger.debug(
+            f"{status_tracker.num_rate_limit_errors} rate limit errors received. "
+            "Consider running at a lower rate."
+        )
+
+    return [res for _, res in sorted(results)]

venv/lib/python3.9/site-packages/mlflow/openai/utils/chat_schema.py

@@ -0,0 +1,292 @@
+import json
+import logging
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel
+
+from mlflow.entities import SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.exceptions import MlflowException
+from mlflow.tracing import set_span_chat_messages, set_span_chat_tools
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    ContentType,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_RESPONSE_API_BUILT_IN_TOOLS = {
+    "file_search",
+    "computer_use_preview",
+    "web_search_preview",
+}
+
+
+def set_span_chat_attributes(span: LiveSpan, inputs: dict[str, Any], output: Any):
+    if span.span_type not in (SpanType.CHAT_MODEL, SpanType.LLM):
+        return
+
+    messages = _parse_inputs_output(inputs, output)
+    try:
+        set_span_chat_messages(span, messages)
+    except MlflowException:
+        _logger.debug(
+            "Failed to set chat messages on span",
+            exc_info=True,
+        )
+
+    if tools := _parse_tools(inputs):
+        try:
+            set_span_chat_tools(span, tools)
+        except MlflowException:
+            _logger.debug("Failed to set chat tools on span", exc_info=True)
+
+
+def _parse_inputs_output(inputs: dict[str, Any], output: Any) -> list[ChatMessage]:
+    from openai.types.chat import ChatCompletion
+
+    try:
+        from openai.types.responses import Response
+
+        if isinstance(output, Response):
+            messages = []
+            if _input := inputs.get("input"):
+                if isinstance(_input, str):
+                    messages.append(ChatMessage(role="user", content=_input))
+                elif isinstance(_input, list):
+                    for item in _input:
+                        messages.extend(_parse_response_item(item, messages))
+
+            for output in output.output:
+                output_dict = output.model_dump(exclude_unset=True)
+                messages.extend(_parse_response_item(output_dict, messages))
+
+            return messages
+    except ImportError:
+        pass
+
+    messages = []
+    if "messages" in inputs:
+        messages.extend(inputs["messages"])
+
+    if isinstance(output, ChatCompletion):
+        messages.extend([output.choices[0].message.to_dict(exclude_unset=True)])
+    elif isinstance(output, str):
+        messages.extend([{"role": "assistant", "content": output}])
+
+    return messages
+
+
+def _parse_response_item(
+    item: Union[dict[str, Any], BaseModel],
+    past_messages: list[ChatMessage],
+) -> list[ChatMessage]:
+    """Parse Response API output into MLflow standard chat messages"""
+    if isinstance(item, BaseModel):
+        item = item.model_dump()
+
+    item_type = item.get("type", "message")
+    if item_type == "message":
+        content, refusal = _parse_message_content(item["content"], past_messages)
+        message = ChatMessage(role=item["role"])
+        if content:
+            message.content = content
+        if refusal:
+            message.refusal = refusal
+        return [message]
+
+    elif item_type == "function_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["call_id"], tool_name=item["name"], arguments=item["arguments"]
+            )
+        ]
+
+    elif item_type == "function_call_output":
+        return [
+            ChatMessage(
+                role="tool",
+                content=item["output"],
+                tool_call_id=item["call_id"],
+            )
+        ]
+
+    elif item_type == "file_search_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["id"],
+                tool_name=item["type"],
+                arguments=json.dumps({"queries": item["queries"]}),
+            ),
+            ChatMessage(role="tool", tool_call_id=item.get("id"), content=item_type),
+        ]
+
+    elif item_type == "web_search_call":
+        return [
+            _get_tool_call_message(tool_id=item["id"], tool_name=item["type"], arguments=""),
+            ChatMessage(role="tool", tool_call_id=item.get("id"), content=item_type),
+        ]
+
+    elif item_type == "computer_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["call_id"],
+                tool_name=item["type"],
+                arguments=json.dumps({"action": item["action"]}),
+            )
+        ]
+    elif item_type == "computer_call_output":
+        output = item["output"]
+        return [
+            ChatMessage(
+                role="tool",
+                content=[
+                    # Screenshot of the computer after taking the action
+                    ImageContentPart(
+                        image_url=ImageUrl(url=output.get("image_url")),
+                        type="image_url",
+                    ),
+                ],
+                tool_call_id=item["call_id"],
+            )
+        ]
+
+    elif item_type == "reasoning":
+        summary = item["summary"][0]["text"] if item["summary"] else None
+        return [ChatMessage(role="assistant", content=summary)]
+
+    raise MlflowException(f"Unknown output type: {type(item)}")
+
+
+def _parse_message_content(
+    content: Union[str, list[dict[str, Any]]], past_messages: Optional[list[ChatMessage]] = None
+) -> tuple[ContentType, Optional[str]]:
+    if isinstance(content, str):
+        return content, None
+
+    if not isinstance(content, list):
+        raise MlflowException(f"Invalid content type: {type(content)}")
+
+    parsed_contents = []
+    refusal = None
+    for item in content:
+        content_type = item.get("type")
+        if content_type == "input_text":
+            parsed_contents.append(TextContentPart(text=item.get("text"), type="text"))
+
+        elif content_type == "input_image":
+            parsed_contents.append(
+                ImageContentPart(
+                    image_url=ImageUrl(
+                        url=item["image_url"],
+                        detail=item.get("detail"),
+                    ),
+                    type="image_url",
+                )
+            )
+
+        elif content_type == "input_file":
+            # TODO: MLflow chat schema doesn't support file content yet. Even if it does,
+            #   including the full file data in an attribute is not a good idea.
+            parsed_contents.append(
+                TextContentPart(
+                    text=f"{item.get('file_id')}:{item.get('file_name')}",
+                    type="text",
+                )
+            )
+
+        elif content_type == "output_text":
+            if annotations := item.get("annotations"):
+                _populate_tool_result_message(annotations, past_messages)
+            parsed_contents.append(TextContentPart(text=item.get("text"), type="text"))
+
+        elif content_type == "refusal":
+            refusal = item.get("refusal")
+
+        else:
+            raise MlflowException(f"Unknown content type: {content_type}")
+
+    return parsed_contents, refusal
+
+
+def _get_tool_call_message(tool_id: str, tool_name: str, arguments: str) -> ChatMessage:
+    return ChatMessage(
+        role="assistant",
+        tool_calls=[
+            ToolCall(
+                id=tool_id,
+                type="function",
+                function=Function(name=tool_name, arguments=arguments),
+            )
+        ],
+    )
+
+
+def _populate_tool_result_message(
+    annotations: dict[str, Any], messages: list[ChatMessage]
+) -> ChatMessage:
+    """
+    Parses annotations from the Response API output into MLflow standard chat spec.
+
+    In OpenAI spec, annotations are used for populating information from file/web
+    search tools in the Responses API response. When converting to MLflow standard
+    chat spec, it should be presented as a tool result message.
+    """
+    file_citations = [a for a in annotations if a.get("type") == "file_citation"]
+    web_citations = [a for a in annotations if a.get("type") == "url_citation"]
+
+    if file_search_result_msg := _find_tool_message(messages, "file_search_call"):
+        file_search_result_msg.content = json.dumps(file_citations)
+
+    if web_search_result_msg := _find_tool_message(messages, "web_search_call"):
+        web_search_result_msg.content = json.dumps(web_citations)
+
+
+def _find_tool_message(messages, tool_type):
+    return next((msg for msg in messages if msg.role == "tool" and msg.content == tool_type), None)
+
+
+def _parse_tools(inputs: dict[str, Any]) -> list[ChatTool]:
+    tools = inputs.get("tools", [])
+
+    if tools is None or not isinstance(tools, Iterable):
+        return []
+
+    parsed_tools = []
+    for tool in tools:
+        tool_type = tool.get("type", "function")
+        if tool_type == "function":
+            if "function" in tool:
+                # ChatCompletion API style
+                parsed_tools.append(ChatTool(**tool))
+            else:
+                # Responses API style
+                definition = {k: v for k, v in tool.items() if k != "type"}
+                parsed_tools.append(
+                    ChatTool(
+                        type="function",
+                        function=FunctionToolDefinition(**definition),
+                    )
+                )
+        elif tool_type in _RESPONSE_API_BUILT_IN_TOOLS:
+            parsed_tools.append(
+                ChatTool(
+                    type="function",
+                    function=FunctionToolDefinition(
+                        name=tool_type,
+                    ),
+                )
+            )
+        else:
+            raise MlflowException(f"Unknown tool type: {tool_type}")
+
+    return parsed_tools