Repository: mlflow/mlflow
Issue category: Bug Report

[BUG] MLflow Tracing Fails to Parse Response Objects When Using OpenAI via DSPy
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Databricks

### MLflow version

- Client: 2.17.2
- Tracking server: 2.17.2


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
Linux 1106-084218-ernmjrfp-10-199-84-236 5.15.0-1071-aws #77~20.04.1-Ubuntu SMP Thu Oct 3 19:39:59 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux
No LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 22.04.5 LTS
Release:	22.04
Codename:	jammy

- **Python version**:
Python 3.11.0rc1

- **yarn version, if running the dev UI**:
apt 2.4.13 (amd64)

### Describe the problem

When using the DSPy package to interact with OpenAI, along with MLflow's OpenAI flavor, MLflow tracing fails to properly parse the content of the response objects. This issue makes it difficult to fully track or log the API responses, which can affect model output tracking and result analysis. I am seeking a solution to ensure that MLflow tracing can correctly parse and record the response object content, allowing the OpenAI flavor to function as expected.

During my experiments, the returned result appears to be a Python Response object like the image below.
<img width="641" alt="截圖 2024-11-15 下午3 26 55" src="https://github.com/user-attachments/assets/43e7a29a-72e9-475a-bc46-c4685a2c7615">

---

Here is the expected result for the Response field that I found in the official documentation:
![image](https://github.com/user-attachments/assets/42e5f6e2-6f37-4b4a-90e1-7a78d6293985)



### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
System information: Linux #77~20.04.1-Ubuntu SMP Thu Oct 3 19:39:59 UTC 2024
Python version: 3.11.0rc1
MLflow version: 2.17.2
MLflow module location: /local_disk0/.ephemeral_nfs/envs/pythonEnv-afd7f4c1-1092-44e0-9961-2b8d23b8f152/lib/python3.11/site-packages/mlflow/__init__.py
Tracking URI: databricks
Registry URI: databricks-uc
Databricks runtime version: 15.4
MLflow environment variables: 
  MLFLOW_CONDA_HOME: /databricks/conda
  MLFLOW_DEPLOYMENTS_TARGET: databricks
  MLFLOW_GATEWAY_URI: databricks
  MLFLOW_PYTHON_EXECUTABLE: /databricks/spark/scripts/mlflow_python.sh
  MLFLOW_TRACKING_URI: databricks
MLflow dependencies: 
  Flask: 2.2.5
  Jinja2: 3.1.2
  aiohttp: 3.8.5
  alembic: 1.14.0
  azure-storage-file-datalake: 12.14.0
  boto3: 1.34.39
  botocore: 1.34.39
  docker: 7.1.0
  google-cloud-storage: 2.10.0
  graphene: 3.4.3
  gunicorn: 20.1.0
  langchain: 0.1.20
  markdown: 3.4.1
  matplotlib: 3.7.2
  mlflow-skinny: 2.17.2
  numpy: 1.23.5
  pandas: 1.5.3
  pyarrow: 14.0.1
  pydantic: 2.9.2
  scikit-learn: 1.3.0
  scipy: 1.11.1
  sqlalchemy: 1.4.39
  tiktoken: 0.8.0
  virtualenv: 20.24.2
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
# Load data
```Python
import random
from typing import Literal
from dspy.datasets import DataLoader
from datasets import load_dataset

# Load the Banking77 dataset and save to ./data.
CLASSES = load_dataset("PolyAI/banking77", split="train", trust_remote_code=True, cache_dir="/Volumes/ai_default/stockllm/sample").features['label'].names
kwargs = dict(fields=("text", "label"), input_keys=("text",), split="train", trust_remote_code=True, cache_dir="/Volumes/ai_default/stockllm/sample")
# Load the first 2000 examples from the dataset, and assign a hint to each *training* example.
trainset = [
    dspy.Example(x, hint=CLASSES[x.label], label=CLASSES[x.label]).with_inputs("text", "hint")
    for x in DataLoader().from_huggingface(dataset_name="PolyAI/banking77", **kwargs)[:2000]
]
random.Random(0).shuffle(trainset)
```
# Inference
```Python
import mlflow

mlflow.start_run()
mlflow.openai.autolog()
import dspy

dspy.settings.experimental = True
dspy.configure(lm=dspy.LM('gpt-4o-mini-2024-07-18', api_key=openai_api_key))

# Define the DSPy module for classification. It will use the hint at training time, if available.
signature = dspy.Signature("text -> label").with_updated_fields('label', type_=Literal[tuple(CLASSES)])
classify = dspy.ChainOfThoughtWithHint(signature)
classify(text="<Some questions ...>")
```
<img width="534" alt="截圖 2024-11-15 下午3 35 25" src="https://github.com/user-attachments/assets/d5758bc7-9a1d-4a15-a0f1-0a1772d0cfd5">


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
There's no error message for my issues.
```
REPLACE_ME
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
There is not logs for my issues
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [X] `integrations/databricks`: Databricks integrations


Functions to edit:
mlflow/openai/_openai_autolog.py:patched_call
Functions added: mlflow/openai/_openai_autolog.py:_try_parse_raw_response

Patch:
diff --git a/mlflow/openai/_openai_autolog.py b/mlflow/openai/_openai_autolog.py
index 294153c1bf38a..757bf5360765b 100644
--- a/mlflow/openai/_openai_autolog.py
+++ b/mlflow/openai/_openai_autolog.py
@@ -4,7 +4,7 @@
 import os
 from contextlib import contextmanager
 from copy import deepcopy
-from typing import Iterator
+from typing import Any, Iterator
 
 from packaging.version import Version
 
@@ -92,6 +92,29 @@ def _get_span_type(task) -> str:
     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):
     from openai import Stream
     from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
@@ -150,7 +173,7 @@ def patched_call(original, self, *args, **kwargs):
 
     # Execute the original function
     try:
-        result = original(self, *args, **kwargs)
+        raw_result = original(self, *args, **kwargs)
     except Exception as e:
         # We have to end the trace even the exception is raised
         if config.log_traces and request_id:
@@ -161,6 +184,8 @@ def patched_call(original, self, *args, **kwargs):
                 _logger.warning(f"Encountered unexpected error when ending trace: {inner_e}")
         raise e
 
+    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
@@ -241,7 +266,7 @@ def _stream_output_logging_hook(stream: Iterator) -> Iterator:
     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 result
+    return raw_result
 
 
 def patched_agent_get_chat_completion(original, self, *args, **kwargs):

Test patch: diff --git a/tests/openai/test_openai_autolog.py b/tests/openai/test_openai_autolog.py
index 0ad67fbafd1ce..b354d43ca002d 100644
--- a/tests/openai/test_openai_autolog.py
+++ b/tests/openai/test_openai_autolog.py
@@ -362,3 +362,41 @@ def test_autolog_use_active_run_id(client, log_models):
     assert traces[1].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_2.info.run_id
     assert traces[2].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_2.info.run_id
     assert traces[3].info.request_metadata[TraceMetadataKey.SOURCE_RUN] == run_3.info.run_id
+
+
+def test_autolog_raw_response(client):
+    mlflow.openai.autolog()
+
+    with mlflow.start_run():
+        resp = client.chat.completions.with_raw_response.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "test"}],
+        )
+        resp = resp.parse()  # ensure the raw response is returned
+
+    assert resp.choices[0].message.content == '[{"role": "user", "content": "test"}]'
+    trace = mlflow.get_last_active_trace()
+    assert len(trace.data.spans) == 1
+    span = trace.data.spans[0]
+    assert isinstance(span.outputs, dict)
+    assert (
+        span.outputs["choices"][0]["message"]["content"] == '[{"role": "user", "content": "test"}]'
+    )
+
+
+def test_autolog_raw_response_stream(client):
+    mlflow.openai.autolog()
+
+    with mlflow.start_run():
+        resp = client.chat.completions.with_raw_response.create(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "test"}],
+            stream=True,
+        )
+        resp = resp.parse()  # ensure the raw response is returned
+
+    assert [c.choices[0].delta.content for c in resp] == ["Hello", " world"]
+    trace = mlflow.get_last_active_trace()
+    assert len(trace.data.spans) == 1
+    span = trace.data.spans[0]
+    assert span.outputs == "Hello world"