[AI] test: Add support for recording requests from the openai package

sdk/ai/azure-ai-generative/assets.json

@@ -0,0 +1,6 @@
+{
+  "AssetsRepo": "Azure/azure-sdk-assets",
+  "AssetsRepoPrefixPath": "python",
+  "TagPrefix": "python/ai/azure-ai-generative",
+  "Tag": ""
+}

sdk/ai/azure-ai-generative/tests/__openai_patcher.py

@@ -0,0 +1,95 @@
+"""Implementation of an httpx.Client that forwards traffic to the Azure SDK test-proxy.
+
+.. note::
+
+    This module has side-effects!
+
+    Importing this module will replace the default httpx.Client used
+    by the openai package with one that can redirect it's traffic
+    to the Azure SDK test-proxy on demand.
+
+"""
+from contextlib import contextmanager
+from typing import Iterable, Literal, Optional
+
+import httpx
+import openai._base_client
+from typing_extensions import override
+from dataclasses import dataclass
+
+
+@dataclass
+class TestProxyConfig:
+    recording_id: str
+    """The ID for the ongoing test recording."""
+
+    recording_mode: Literal["playback", "record"]
+    """The current recording mode."""
+
+    proxy_url: str
+    """The url for the Azure SDK test proxy."""
+
+
+class TestProxyHttpxClient(openai._base_client.SyncHttpxClientWrapper):
+    recording_config: Optional[TestProxyConfig] = None
+
+    @classmethod
+    def is_recording(cls) -> bool:
+        """Whether we are forwarding requests to the test proxy
+
+        :return: True if forwarding, False otherwise
+        :rtype: bool
+        """
+        return cls.recording_config is not None
+
+    @classmethod
+    @contextmanager
+    def record_with_proxy(cls, config: TestProxyConfig) -> Iterable[None]:
+        """Forward all requests made within the scope of context manager to test-proxy.
+
+        :param TestProxyConfig config: The test proxy configuration
+        """
+        cls.recording_config = config
+
+        yield
+
+        cls.recording_config = None
+
+    @override
+    def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        if self.is_recording():
+            return self._send_to_proxy(request, **kwargs)
+        else:
+            return super().send(request, **kwargs)
+
+    def _send_to_proxy(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        """Forwards a network request to the test proxy
+
+        :param httpx.Request request: The request to send
+        :keyword **kwargs: The kwargs accepted by httpx.Client.send
+        :return: The request's response
+        :rtype: httpx.Response
+        """
+        assert self.is_recording(), f"{self._send_to_proxy.__qualname__} should only be called while recording"
+        config = self.recording_config
+        original_url = request.url
+
+        request_path = original_url.copy_with(scheme="", netloc=b"")
+        request.url = httpx.URL(config.proxy_url).join(request_path)
+
+        headers = request.headers
+        if headers.get("x-recording-upstream-base-uri", None) is None:
+            headers["x-recording-upstream-base-uri"] = str(
+                httpx.URL(scheme=original_url.scheme, netloc=original_url.netloc)
+            )
+        headers["x-recording-id"] = config.recording_id
+        headers["x-recording-mode"] = config.recording_mode
+
+        response = super().send(request, **kwargs)
+
+        response.request.url = original_url
+        return response
+
+
+# openai._base_client.SyncHttpxClientWrapper is default httpx.Client instantiated by openai
+openai._base_client.SyncHttpxClientWrapper = TestProxyHttpxClient

sdk/ai/azure-ai-generative/tests/conftest.py

@@ -1,3 +1,4 @@
+from __openai_patcher import TestProxyConfig, TestProxyHttpxClient # isort: split
 import asyncio
 import base64
 import os
@@ -6,7 +7,6 @@
 import pytest
 from azure.ai.generative.synthetic.qa import QADataGenerator
 
-import pytest
 from packaging import version
 from devtools_testutils import (
     FakeTokenCredential,
@@ -17,6 +17,8 @@
     is_live,
     set_custom_default_matcher,
 )
+from devtools_testutils.config import PROXY_URL
+from devtools_testutils.helpers import get_recording_id
 from devtools_testutils.proxy_fixtures import EnvironmentVariableSanitizer
 
 from azure.ai.resources.client import AIClient
@@ -25,6 +27,18 @@
 from azure.identity import AzureCliCredential, ClientSecretCredential
 
 
+@pytest.fixture()
+def recorded_test(recorded_test):
+    """Route requests from the openai package to the test proxy."""
+
+    config = TestProxyConfig(
+        recording_id=get_recording_id(), recording_mode="record" if is_live() else "playback", proxy_url=PROXY_URL
+    )
+
+
+    with TestProxyHttpxClient.record_with_proxy(config):
+        yield recorded_test
+
 @pytest.fixture()
 def ai_client(
     e2e_subscription_id: str,

sdk/ai/azure-ai-generative/tests/synthetic_qa/unittests/test_qa_data_generator.py

@@ -91,7 +91,7 @@ def test_export_format(self, qa_type, structure):
         qa_generator = QADataGenerator(model_config)
         qas = list(zip(questions, answers))
         filepath = os.path.join(pathlib.Path(__file__).parent.parent.resolve(), "data")
-        output_file = os.path.join(filepath, f"test_{qa_type}_{structure}.jsonl")
+        output_file = os.path.join(filepath, f"test_{qa_type.value}_{structure.value}.jsonl")
         qa_generator.export_to_file(output_file, qa_type, qas, structure)
 
         if qa_type == QAType.CONVERSATION and structure == OutputStructure.CHAT_PROTOCOL:

sdk/ai/azure-ai-resources/dev_requirements.txt

@@ -5,3 +5,4 @@
 -e ../../ml/azure-ai-ml
 pytest
 pytest-xdist
+openai

sdk/ai/azure-ai-resources/tests/__openai_patcher.py

@@ -0,0 +1,95 @@
+"""Implementation of an httpx.Client that forwards traffic to the Azure SDK test-proxy.
+
+.. note::
+
+    This module has side-effects!
+
+    Importing this module will replace the default httpx.Client used
+    by the openai package with one that can redirect it's traffic
+    to the Azure SDK test-proxy on demand.
+
+"""
+from contextlib import contextmanager
+from typing import Iterable, Literal, Optional
+
+import httpx
+import openai._base_client
+from typing_extensions import override
+from dataclasses import dataclass
+
+
+@dataclass
+class TestProxyConfig:
+    recording_id: str
+    """The ID for the ongoing test recording."""
+
+    recording_mode: Literal["playback", "record"]
+    """The current recording mode."""
+
+    proxy_url: str
+    """The url for the Azure SDK test proxy."""
+
+
+class TestProxyHttpxClient(openai._base_client.SyncHttpxClientWrapper):
+    recording_config: Optional[TestProxyConfig] = None
+
+    @classmethod
+    def is_recording(cls) -> bool:
+        """Whether we are forwarding requests to the test proxy
+
+        :return: True if forwarding, False otherwise
+        :rtype: bool
+        """
+        return cls.recording_config is not None
+
+    @classmethod
+    @contextmanager
+    def record_with_proxy(cls, config: TestProxyConfig) -> Iterable[None]:
+        """Forward all requests made within the scope of context manager to test-proxy.
+
+        :param TestProxyConfig config: The test proxy configuration
+        """
+        cls.recording_config = config
+
+        yield
+
+        cls.recording_config = None
+
+    @override
+    def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        if self.is_recording():
+            return self._send_to_proxy(request, **kwargs)
+        else:
+            return super().send(request, **kwargs)
+
+    def _send_to_proxy(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        """Forwards a network request to the test proxy
+
+        :param httpx.Request request: The request to send
+        :keyword **kwargs: The kwargs accepted by httpx.Client.send
+        :return: The request's response
+        :rtype: httpx.Response
+        """
+        assert self.is_recording(), f"{self._send_to_proxy.__qualname__} should only be called while recording"
+        config = self.recording_config
+        original_url = request.url
+
+        request_path = original_url.copy_with(scheme="", netloc=b"")
+        request.url = httpx.URL(config.proxy_url).join(request_path)
+
+        headers = request.headers
+        if headers.get("x-recording-upstream-base-uri", None) is None:
+            headers["x-recording-upstream-base-uri"] = str(
+                httpx.URL(scheme=original_url.scheme, netloc=original_url.netloc)
+            )
+        headers["x-recording-id"] = config.recording_id
+        headers["x-recording-mode"] = config.recording_mode
+
+        response = super().send(request, **kwargs)
+
+        response.request.url = original_url
+        return response
+
+
+# openai._base_client.SyncHttpxClientWrapper is default httpx.Client instantiated by openai
+openai._base_client.SyncHttpxClientWrapper = TestProxyHttpxClient

sdk/ai/azure-ai-resources/tests/conftest.py

@@ -1,3 +1,4 @@
+from __openai_patcher import TestProxyConfig, TestProxyHttpxClient # isort: split
 import asyncio
 import base64
 import os
@@ -17,6 +18,8 @@
     is_live,
     set_custom_default_matcher,
 )
+from devtools_testutils.config import PROXY_URL
+from devtools_testutils.helpers import get_recording_id
 from devtools_testutils.proxy_fixtures import (
     EnvironmentVariableSanitizer,
     VariableRecorder
@@ -38,6 +41,17 @@ def generate_random_string():
     return generate_random_string
 
 
+@pytest.fixture()
+def recorded_test(recorded_test):
+    """Route requests from the openai package to the test proxy."""
+
+    config = TestProxyConfig(
+        recording_id=get_recording_id(), recording_mode="record" if is_live() else "playback", proxy_url=PROXY_URL
+    )
+
+
+    with TestProxyHttpxClient.record_with_proxy(config):
+        yield recorded_test
 
 @pytest.fixture()
 def ai_client(