Support for OpenAI's .with_raw_response. (#1138)

* Add initial fixes and tests

* Update v1 mock server responses

* Remove commented out code

* Use attribute checks instead

* Change try/except to hasattr checks

* Remove psuedo code in comments

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

newrelic/hooks/mlmodel_openai.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import sys
 import traceback
@@ -251,11 +252,17 @@ def _record_embedding_success(transaction, embedding_id, linking_metadata, kwarg
         response_headers = getattr(response, "_nr_response_headers", {})
         input = kwargs.get("input")
 
+        attribute_response = response
         # In v1, response objects are pydantic models so this function call converts the
         # object back to a dictionary for backwards compatibility.
-        attribute_response = response
         if OPENAI_V1:
-            attribute_response = response.model_dump()
+            if hasattr(response, "model_dump"):
+                attribute_response = response.model_dump()
+            elif hasattr(response, "http_response") and hasattr(response.http_response, "text"):
+                # This is for the .with_raw_response. wrapper.  This is expected
+                # to change, but the return type for now is the following:
+                # openai._legacy_response.LegacyAPIResponse
+                attribute_response = json.loads(response.http_response.text.strip())
 
         request_id = response_headers.get("x-request-id")
         response_model = attribute_response.get("model")
@@ -441,11 +448,18 @@ def _handle_completion_success(transaction, linking_metadata, completion_id, kwa
         # If response is not a stream generator, record the event data.
         # At this point, we have a response so we can grab attributes only available on the response object
         response_headers = getattr(return_val, "_nr_response_headers", {})
+        response = return_val
+
         # In v1, response objects are pydantic models so this function call converts the
         # object back to a dictionary for backwards compatibility.
-        response = return_val
         if OPENAI_V1:
-            response = response.model_dump()
+            if hasattr(response, "model_dump"):
+                response = response.model_dump()
+            elif hasattr(response, "http_response") and hasattr(response.http_response, "text"):
+                # This is for the .with_raw_response. wrapper.  This is expected
+                # to change, but the return type for now is the following:
+                # openai._legacy_response.LegacyAPIResponse
+                response = json.loads(response.http_response.text.strip())
 
         _record_completion_success(transaction, linking_metadata, completion_id, kwargs, ft, response_headers, response)
     except Exception:

tests/mlmodel_openai/conftest.py

@@ -198,7 +198,7 @@ def _wrap_httpx_client_send(wrapped, instance, args, kwargs):
         # Send request
         response = wrapped(*args, **kwargs)
 
-        if response.status_code >= 400 or response.status_code < 200:
+        if response.status_code >= 500 or response.status_code < 200:
             prompt = "error"
 
         rheaders = getattr(response, "headers")