diff --git a/sdk/ai/azure-ai-generative/assets.json b/sdk/ai/azure-ai-generative/assets.json new file mode 100644 index 000000000000..3e89e3ae403c --- /dev/null +++ b/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": "" +} diff --git a/sdk/ai/azure-ai-generative/tests/__openai_patcher.py b/sdk/ai/azure-ai-generative/tests/__openai_patcher.py new file mode 100644 index 000000000000..4746f96b1054 --- /dev/null +++ b/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 diff --git a/sdk/ai/azure-ai-generative/tests/conftest.py b/sdk/ai/azure-ai-generative/tests/conftest.py index 5f0a00280513..670692fbf117 100644 --- a/sdk/ai/azure-ai-generative/tests/conftest.py +++ b/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, diff --git a/sdk/ai/azure-ai-generative/tests/synthetic_qa/unittests/test_qa_data_generator.py b/sdk/ai/azure-ai-generative/tests/synthetic_qa/unittests/test_qa_data_generator.py index 21266c4acbc0..04672a43f162 100644 --- a/sdk/ai/azure-ai-generative/tests/synthetic_qa/unittests/test_qa_data_generator.py +++ b/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: diff --git a/sdk/ai/azure-ai-resources/dev_requirements.txt b/sdk/ai/azure-ai-resources/dev_requirements.txt index d2aec9d69669..7d5946ecf57b 100644 --- a/sdk/ai/azure-ai-resources/dev_requirements.txt +++ b/sdk/ai/azure-ai-resources/dev_requirements.txt @@ -5,3 +5,4 @@ -e ../../ml/azure-ai-ml pytest pytest-xdist +openai diff --git a/sdk/ai/azure-ai-resources/tests/__openai_patcher.py b/sdk/ai/azure-ai-resources/tests/__openai_patcher.py new file mode 100644 index 000000000000..4746f96b1054 --- /dev/null +++ b/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 diff --git a/sdk/ai/azure-ai-resources/tests/conftest.py b/sdk/ai/azure-ai-resources/tests/conftest.py index bb339a4792f2..a87a35de87f0 100644 --- a/sdk/ai/azure-ai-resources/tests/conftest.py +++ b/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(