Repository: deepset-ai/haystack

Issue:
Update OpenAI Encoder to optionally use the Azure OpenAI service
**Is your feature request related to a problem? Please describe.**
I think it could be very useful to update the OpenAI nodes in Haystack to be able to ping the [Azure OpenAI](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/overview) service that was recently announced. This is mostly from experience that the OpenAI API is often overloaded and unresponsive. I believe the Azure OpenAI service hosted on Azure cloud will be more reliable and usable for production workloads.

**Describe the solution you'd like**
Update the existing OpenAI nodes to optionally make requests to the Azure OpenAI service that can be toggled using a boolean parameter. 

**Describe alternatives you've considered**
We could also create new nodes for the Azure OpenAI API.



Discussion:
Cool, how big is this change? Do we need to point the REST client to some new URL, or is the change more substantial? 
These values should be changed, if using official library and not using Azure AD Authentication.

```
openai.api_version = '2022-12-01'
openai.api_base = 'ENDPOINT'
```
```
openai.api_type = 'azure'
openai.api_key = 'KEY'
```

Basically, changes the URL and headers.
[https://github.com/danielbichuetti/openai-python/blob/main/openai/api_requestor.py](https://github.com/danielbichuetti/openai-python/blob/main/openai/api_requestor.py)
There are 2 ways of using Azure OpenAI endpoints (it can be checked on official client, [here](https://github.com/danielbichuetti/openai-python/blob/48b69293a3dff649165c1d0e0f9d77d1b52436c5/openai/util.py#L28)):
- Azure API Key
- Azure AD

Main changes are related to API version and endpoint:

```
openai.api_version = '2022-12-01'
openai.api_base = '' # Please add your endpoint here
```

API key header difference can be seen [here](https://github.com/danielbichuetti/openai-python/blob/48b69293a3dff649165c1d0e0f9d77d1b52436c5/openai/util.py#L21):
```
api_key_to_header = (
    lambda api, key: {"Authorization": f"Bearer {key}"}
    if api in (ApiType.OPEN_AI, ApiType.AZURE_AD)
    else {"api-key": f"{key}"}
)
```

I'll propose to add the ability to set up organization ID too, as this is recommended for approved projects on OpenAI service.
```
headers["OpenAI-Organization"] = self. Organization
```


**UPDATE**
There is a REST API reference for the Azure service [here](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference)
Code is starting to be repeated in a lot of places.

Isn't it time to start using OpenAI official client? It will be a slight increase in the total project size. Furthermore, it will help with maintenance.
Also when using the Azure OpenAI services we will be able to share it with more clients because of compliance issues with OpenAI. E.g. OpenAI doesnt have SOC2 or ISO 27001.
There was a Discord topic about the support for Azure OpenAI.

User asked about it, and mentioned regarding the AnswerGenerator.

[Discord topic](https://discord.com/channels/993534733298450452/993539071815200889/1072485828686794772)


Resolving patch:
diff --git a/haystack/nodes/answer_generator/openai.py b/haystack/nodes/answer_generator/openai.py
--- a/haystack/nodes/answer_generator/openai.py
+++ b/haystack/nodes/answer_generator/openai.py
@@ -53,8 +53,11 @@ class OpenAIAnswerGenerator(BaseGenerator):
     def __init__(
         self,
         api_key: str,
+        azure_base_url: Optional[str] = None,
+        azure_deployment_name: Optional[str] = None,
         model: str = "text-davinci-003",
         max_tokens: int = 50,
+        api_version: str = "2022-12-01",
         top_k: int = 5,
         temperature: float = 0.2,
         presence_penalty: float = 0.1,
@@ -68,6 +71,11 @@ def __init__(
     ):
         """
         :param api_key: Your API key from OpenAI. It is required for this node to work.
+        :param azure_base_url: The base URL for the Azure OpenAI API. If not supplied, Azure OpenAI API will not be used.
+                               This parameter is an OpenAI Azure endpoint, usually in the form `https://<your-endpoint>.openai.azure.com'
+
+        :param azure_deployment_name: The name of the Azure OpenAI API deployment. If not supplied, Azure OpenAI API
+                                     will not be used.
         :param model: ID of the engine to use for generating the answer. You can select one of `"text-ada-001"`,
                      `"text-babbage-001"`, `"text-curie-001"`, or `"text-davinci-003"`
                      (from worst to best and from cheapest to most expensive). For more information about the models,
@@ -75,6 +83,7 @@ def __init__(
         :param max_tokens: The maximum number of tokens reserved for the generated Answer.
                            A higher number allows for longer answers without exceeding the max prompt length of the OpenAI model.
                            A lower number allows longer prompts with more documents passed as context, but the generated answer might be cut after max_tokens.
+        :param api_version: The version of the Azure OpenAI API to use. The default is `2022-12-01` version.
         :param top_k: Number of generated Answers.
         :param temperature: What sampling temperature to use. Higher values mean the model will take more risks and
                             value 0 (argmax sampling) works better for scenarios with a well-defined Answer.
@@ -159,6 +168,9 @@ def __init__(
                 )
 
         self.api_key = api_key
+        self.azure_base_url = azure_base_url
+        self.azure_deployment_name = azure_deployment_name
+        self.api_version = api_version
         self.model = model
         self.max_tokens = max_tokens
         self.top_k = top_k
@@ -179,6 +191,8 @@ def __init__(
         else:
             self.MAX_TOKENS_LIMIT = 2048
 
+        self.using_azure = self.azure_deployment_name is not None and self.azure_base_url is not None
+
         if USE_TIKTOKEN:
             logger.debug("Using tiktoken %s tokenizer", tokenizer)
             self._tk_tokenizer: tiktoken.Encoding = tiktoken.get_encoding(tokenizer)
@@ -229,9 +243,6 @@ def predict(
         prompt, input_docs = self._build_prompt_within_max_length(query=query, documents=documents)
         logger.debug("Prompt being sent to OpenAI API with prompt %s.", prompt)
 
-        # get answers from OpenAI API
-        url = "https://api.openai.com/v1/completions"
-
         payload = {
             "model": self.model,
             "prompt": prompt,
@@ -242,9 +253,17 @@ def predict(
             "presence_penalty": self.presence_penalty,
             "frequency_penalty": self.frequency_penalty,
         }
+        url = "https://api.openai.com/v1/completions"
+        if self.using_azure:
+            url = f"{self.azure_base_url}/openai/deployments/{self.azure_deployment_name}/completions?api-version={self.api_version}"
+
+        headers = {"Content-Type": "application/json"}
+        if self.using_azure:
+            headers = {"api-key": self.api_key, **headers}
+        else:
+            headers = {"Authorization": f"Bearer {self.api_key}", **headers}
 
-        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
-        response = requests.request("POST", url, headers=headers, data=json.dumps(payload), timeout=timeout)
+        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
         res = json.loads(response.text)
 
         if response.status_code != 200 or "choices" not in res:
diff --git a/haystack/nodes/prompt/prompt_node.py b/haystack/nodes/prompt/prompt_node.py
--- a/haystack/nodes/prompt/prompt_node.py
+++ b/haystack/nodes/prompt/prompt_node.py
@@ -212,11 +212,13 @@ def invoke(self, *args, **kwargs):
         pass
 
     @classmethod
-    def supports(cls, model_name_or_path: str) -> bool:
+    def supports(cls, model_name_or_path: str, **kwargs) -> bool:
         """
         Checks if the given model is supported by this invocation layer.
 
         :param model_name_or_path: The name or path of the model.
+        :param kwargs: additional keyword arguments passed to the underlying model which might be used to determine
+        if the model is supported.
         :return: True if this invocation layer supports the model, False otherwise.
         """
         return False
@@ -371,7 +373,7 @@ def invoke(self, *args, **kwargs):
         return generated_texts
 
     @classmethod
-    def supports(cls, model_name_or_path: str) -> bool:
+    def supports(cls, model_name_or_path: str, **kwargs) -> bool:
         try:
             config = AutoConfig.from_pretrained(model_name_or_path)
         except OSError:
@@ -422,7 +424,6 @@ def __init__(
                 f"api_key {api_key} must be a valid OpenAI key. Visit https://openai.com/api/ to get one."
             )
         self.api_key = api_key
-        self.url = "https://api.openai.com/v1/completions"
 
         # Due to reflective construction of all invocation layers we might receive some
         # unknown kwargs, so we need to take only the relevant.
@@ -446,6 +447,14 @@ def __init__(
             if key in kwargs
         }
 
+    @property
+    def url(self) -> str:
+        return "https://api.openai.com/v1/completions"
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
     @retry_with_exponential_backoff(
         backoff_in_seconds=float(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 5)),
         max_retries=int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5)),
@@ -494,11 +503,9 @@ def invoke(self, *args, **kwargs):
             "best_of": kwargs_with_defaults.get("best_of", 1),
             "logit_bias": kwargs_with_defaults.get("logit_bias", {}),
         }
-        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
-        response = requests.request(
-            "POST",
+        response = requests.post(
             self.url,
-            headers=headers,
+            headers=self.headers,
             data=json.dumps(payload),
             timeout=float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30)),
         )
@@ -530,15 +537,62 @@ def invoke(self, *args, **kwargs):
         return responses
 
     @classmethod
-    def supports(cls, model_name_or_path: str) -> bool:
-        return any(m for m in ["ada", "babbage", "davinci", "curie"] if m in model_name_or_path)
+    def supports(cls, model_name_or_path: str, **kwargs) -> bool:
+        valid_model = any(m for m in ["ada", "babbage", "davinci", "curie"] if m in model_name_or_path)
+        return valid_model and kwargs.get("azure_base_url") is None
+
+
+class AzureOpenAIInvocationLayer(OpenAIInvocationLayer):
+    """
+    Azure OpenAI Invocation Layer
+
+    This layer is used to invoke the OpenAI API on Azure. It is essentially the same as the OpenAIInvocationLayer
+    with additional two parameters: azure_base_url and azure_deployment_name. The azure_base_url is the URL of the Azure OpenAI
+    endpoint and the azure_deployment_name is the name of the deployment.
+    """
+
+    def __init__(
+        self,
+        azure_base_url: str,
+        azure_deployment_name: str,
+        api_key: str,
+        api_version: str = "2022-12-01",
+        model_name_or_path: str = "text-davinci-003",
+        max_length: Optional[int] = 100,
+        **kwargs,
+    ):
+        super().__init__(api_key, model_name_or_path, max_length, **kwargs)
+        self.azure_base_url = azure_base_url
+        self.azure_deployment_name = azure_deployment_name
+        self.api_version = api_version
+
+    @property
+    def url(self) -> str:
+        return f"{self.azure_base_url}/openai/deployments/{self.azure_deployment_name}/completions?api-version={self.api_version}"
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        return {"api-key": self.api_key, "Content-Type": "application/json"}
+
+    @classmethod
+    def supports(cls, model_name_or_path: str, **kwargs) -> bool:
+        """
+        Ensures Azure OpenAI Invocation Layer is selected when azure_base_url and azure_deployment_name are provided in
+        addition to a list of supported models.
+        """
+        valid_model = any(m for m in ["ada", "babbage", "davinci", "curie"] if m in model_name_or_path)
+        return (
+            valid_model and kwargs.get("azure_base_url") is not None and kwargs.get("azure_deployment_name") is not None
+        )
 
 
 class PromptModel(BaseComponent):
     """
     The PromptModel class is a component that uses a pre-trained model to perform tasks based on a prompt. Out of
-    the box, it supports two model invocation layers: Hugging Face transformers and OpenAI, with the ability to
-    register additional custom invocation layers.
+    the box, it supports model invocation layers for:
+    - Hugging Face transformers (all text2text-generation models)
+    - OpenAI InstructGPT models
+    - Azure OpenAI InstructGPT models
 
     Although it is possible to use PromptModel to make prompt invocations on the underlying model, use
     PromptNode to interact with the model. PromptModel instances are a way for multiple
@@ -569,6 +623,11 @@ def __init__(
         :param use_gpu: Whether to use GPU or not.
         :param devices: The devices to use where the model is loaded.
         :param model_kwargs: Additional keyword arguments passed to the underlying model.
+
+        Note that Azure OpenAI InstructGPT models require two additional parameters: azure_base_url (The URL for the
+        Azure OpenAI API endpoint, usually in the form `https://<your-endpoint>.openai.azure.com') and
+        azure_deployment_name (The name of the Azure OpenAI API deployment). These parameters should be supplied
+        in the `model_kwargs` dictionary.
         """
         super().__init__()
         self.model_name_or_path = model_name_or_path
@@ -584,6 +643,7 @@ def __init__(
 
         self.register(HFLocalInvocationLayer)  # pylint: disable=W0108
         self.register(OpenAIInvocationLayer)  # pylint: disable=W0108
+        self.register(AzureOpenAIInvocationLayer)  # pylint: disable=W0108
 
         self.model_invocation_layer = self.create_invocation_layer()
 
@@ -597,7 +657,7 @@ def create_invocation_layer(self) -> PromptModelInvocationLayer:
         all_kwargs = {**self.model_kwargs, **kwargs}
 
         for invocation_layer in self.invocation_layers:
-            if invocation_layer.supports(self.model_name_or_path):
+            if invocation_layer.supports(self.model_name_or_path, **all_kwargs):
                 return invocation_layer(
                     model_name_or_path=self.model_name_or_path, max_length=self.max_length, **all_kwargs
                 )
@@ -647,6 +707,9 @@ def run_batch(
     ):
         raise NotImplementedError("This method should never be implemented in the derived class")
 
+    def __repr__(self):
+        return "{}({!r})".format(self.__class__.__name__, self.__dict__)
+
 
 def get_predefined_prompt_templates() -> List[PromptTemplate]:
     return [
@@ -721,7 +784,8 @@ class PromptNode(BaseComponent):
     """
     The PromptNode class is the central abstraction in Haystack's large language model (LLM) support. PromptNode
     supports multiple NLP tasks out of the box. You can use it to perform tasks, such as
-    summarization, question answering, question generation, and more, using a single, unified model within the Haystack framework.
+    summarization, question answering, question generation, and more, using a single, unified model within the Haystack
+    framework.
 
     One of the benefits of PromptNode is that you can use it to define and add additional prompt templates
      the model supports. Defining additional prompt templates makes it possible to extend the model's capabilities
@@ -732,9 +796,13 @@ class PromptNode(BaseComponent):
     Using an instance of the PromptModel class, you can create multiple PromptNodes that share the same model, saving
     the memory and time required to load the model multiple times.
 
-    PromptNode also supports multiple model invocation layers: Hugging Face transformers and OpenAI with an
-    ability to register additional custom invocation layers. However, we currently support only
-    T5 Flan and OpenAI InstructGPT models.
+    PromptNode also supports multiple model invocation layers:
+    - Hugging Face transformers (all text2text-generation models)
+    - OpenAI InstructGPT models
+    - Azure OpenAI InstructGPT models
+
+    However, users are not limited to above-mentioned models only as there is a built-in ability to register
+    additional custom model invocation layers.
 
     We recommend using LLMs fine-tuned on a collection of datasets phrased as instructions, otherwise we find that the
     LLM does not "follow" prompt instructions well. This is why we recommend using T5 flan or OpenAI InstructGPT models.
@@ -772,6 +840,12 @@ def __init__(
         :param top_k: The number of independently generated texts to return per prompt. For example, if you set top_k=3, the model's going to generate three answers to the query.
         :param stop_words: Stops text generation if any of the stop words is generated.
         :param model_kwargs: Additional keyword arguments passed when loading the model specified in `model_name_or_path`.
+
+        Note that Azure OpenAI InstructGPT models require two additional parameters: azure_base_url (The URL for the
+        Azure OpenAI API endpoint, usually in the form `https://<your-endpoint>.openai.azure.com') and
+        azure_deployment_name (The name of the Azure OpenAI API deployment).
+        These parameters should be supplied in the `model_kwargs` dictionary.
+
         """
         send_event("PromptNode initialized")
         super().__init__()


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -537,7 +537,17 @@ def rag_generator():
 
 @pytest.fixture
 def openai_generator():
-    return OpenAIAnswerGenerator(api_key=os.environ.get("OPENAI_API_KEY", ""), model="text-babbage-001", top_k=1)
+    azure_conf = haystack_azure_conf()
+    if azure_conf:
+        return OpenAIAnswerGenerator(
+            api_key=azure_conf["api_key"],
+            azure_base_url=azure_conf["azure_base_url"],
+            azure_deployment_name=azure_conf["azure_deployment_name"],
+            model="text-babbage-001",
+            top_k=1,
+        )
+    else:
+        return OpenAIAnswerGenerator(api_key=os.environ.get("OPENAI_API_KEY", ""), model="text-babbage-001", top_k=1)
 
 
 @pytest.fixture
@@ -974,6 +984,28 @@ def prompt_node():
     return PromptNode("google/flan-t5-small", devices=["cpu"])
 
 
+def haystack_azure_conf():
+    api_key = os.environ.get("AZURE_OPENAI_API_KEY", None)
+    azure_base_url = os.environ.get("AZURE_OPENAI_BASE_URL", None)
+    azure_deployment_name = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", None)
+    if api_key and azure_base_url and azure_deployment_name:
+        return {"api_key": api_key, "azure_base_url": azure_base_url, "azure_deployment_name": azure_deployment_name}
+    else:
+        return {}
+
+
+@pytest.fixture
+def haystack_openai_config(request):
+    if request.param == "openai":
+        api_key = os.environ.get("OPENAI_API_KEY", None)
+        if not api_key:
+            return {}
+        else:
+            return {"api_key": api_key}
+    elif request.param == "azure":
+        return haystack_azure_conf()
+
+
 @pytest.fixture
 def prompt_model(request):
     if request.param == "openai":
@@ -981,5 +1013,15 @@ def prompt_model(request):
         if api_key is None or api_key == "":
             api_key = "KEY_NOT_FOUND"
         return PromptModel("text-davinci-003", api_key=api_key)
+    elif request.param == "azure":
+        api_key = os.environ.get("AZURE_OPENAI_API_KEY", "KEY_NOT_FOUND")
+        if api_key is None or api_key == "":
+            api_key = "KEY_NOT_FOUND"
+        return PromptModel("text-davinci-003", api_key=api_key, model_kwargs=haystack_azure_conf())
     else:
         return PromptModel("google/flan-t5-base", devices=["cpu"])
+
+
+@pytest.fixture
+def azure_conf():
+    return haystack_azure_conf()
diff --git a/test/nodes/test_generator.py b/test/nodes/test_generator.py
--- a/test/nodes/test_generator.py
+++ b/test/nodes/test_generator.py
@@ -133,22 +133,29 @@ def __call__(self, some_invalid_para: str, another_invalid_param: str) -> None:
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="No OpenAI API key provided. Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_openai_answer_generator(openai_generator, docs):
+@pytest.mark.parametrize("haystack_openai_config", ["openai", "azure"], indirect=True)
+def test_openai_answer_generator(haystack_openai_config, docs):
+    if not haystack_openai_config:
+        pytest.skip("No API key found, skipping test")
+
+    openai_generator = OpenAIAnswerGenerator(
+        api_key=haystack_openai_config["api_key"],
+        azure_base_url=haystack_openai_config.get("azure_base_url", None),
+        azure_deployment_name=haystack_openai_config.get("azure_deployment_name", None),
+        model="text-babbage-001",
+        top_k=1,
+    )
     prediction = openai_generator.predict(query="Who lives in Berlin?", documents=docs, top_k=1)
     assert len(prediction["answers"]) == 1
     assert "Carla" in prediction["answers"][0].answer
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="No OpenAI API key provided. Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_openai_answer_generator_custom_template(docs):
+@pytest.mark.parametrize("haystack_openai_config", ["openai", "azure"], indirect=True)
+def test_openai_answer_generator_custom_template(haystack_openai_config, docs):
+    if not haystack_openai_config:
+        pytest.skip("No API key found, skipping test")
+
     lfqa_prompt = PromptTemplate(
         name="lfqa",
         prompt_text="""
@@ -157,20 +164,29 @@ def test_openai_answer_generator_custom_template(docs):
         prompt_params=["context", "query"],
     )
     node = OpenAIAnswerGenerator(
-        api_key=os.environ.get("OPENAI_API_KEY", ""), model="text-babbage-001", top_k=1, prompt_template=lfqa_prompt
+        api_key=haystack_openai_config["api_key"],
+        azure_base_url=haystack_openai_config.get("azure_base_url", None),
+        azure_deployment_name=haystack_openai_config.get("azure_deployment_name", None),
+        model="text-babbage-001",
+        top_k=1,
+        prompt_template=lfqa_prompt,
     )
     prediction = node.predict(query="Who lives in Berlin?", documents=docs, top_k=1)
     assert len(prediction["answers"]) == 1
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="No OpenAI API key provided. Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_openai_answer_generator_max_token(docs, caplog):
+@pytest.mark.parametrize("haystack_openai_config", ["openai", "azure"], indirect=True)
+def test_openai_answer_generator_max_token(haystack_openai_config, docs, caplog):
+    if not haystack_openai_config:
+        pytest.skip("No API key found, skipping test")
+
     openai_generator = OpenAIAnswerGenerator(
-        api_key=os.environ.get("OPENAI_API_KEY", ""), model="text-babbage-001", top_k=1
+        api_key=haystack_openai_config["api_key"],
+        azure_base_url=haystack_openai_config.get("azure_base_url", None),
+        azure_deployment_name=haystack_openai_config.get("azure_deployment_name", None),
+        model="text-babbage-001",
+        top_k=1,
     )
     openai_generator.MAX_TOKENS_LIMIT = 116
     with caplog.at_level(logging.INFO):
diff --git a/test/nodes/test_prompt_node.py b/test/nodes/test_prompt_node.py
--- a/test/nodes/test_prompt_node.py
+++ b/test/nodes/test_prompt_node.py
@@ -11,8 +11,17 @@
 from haystack.nodes.prompt.prompt_node import HFLocalInvocationLayer
 
 
-def is_openai_api_key_set(api_key: str):
-    return len(api_key) > 0 and api_key != "KEY_NOT_FOUND"
+def skip_test_for_invalid_key(prompt_model):
+    if prompt_model.api_key is not None and prompt_model.api_key == "KEY_NOT_FOUND":
+        pytest.skip("No API key found, skipping test")
+
+
+@pytest.fixture
+def get_api_key(request):
+    if request.param == "openai":
+        return os.environ.get("OPENAI_API_KEY", None)
+    elif request.param == "azure":
+        return os.environ.get("AZURE_OPENAI_API_KEY", None)
 
 
 @pytest.mark.unit
@@ -253,41 +262,31 @@ def test_invalid_state_ops(prompt_node):
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_open_ai_prompt_with_params():
-    pm = PromptModel("text-davinci-003", api_key=os.environ["OPENAI_API_KEY"])
-    pn = PromptNode(pm)
+@pytest.mark.parametrize("prompt_model", ["openai", "azure"], indirect=True)
+def test_open_ai_prompt_with_params(prompt_model):
+    skip_test_for_invalid_key(prompt_model)
+    pn = PromptNode(prompt_model)
     optional_davinci_params = {"temperature": 0.5, "max_tokens": 10, "top_p": 1, "frequency_penalty": 0.5}
     r = pn.prompt("question-generation", documents=["Berlin is the capital of Germany."], **optional_davinci_params)
     assert len(r) == 1 and len(r[0]) > 0
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_open_ai_prompt_with_default_params():
-    pn = PromptNode(
-        model_name_or_path="text-davinci-003",
-        api_key=os.environ["OPENAI_API_KEY"],
-        model_kwargs={"temperature": 0.5, "max_tokens": 2, "top_p": 1, "frequency_penalty": 0.5},
-    )
+def test_open_ai_prompt_with_default_params(azure_conf):
+    if not azure_conf:
+        pytest.skip("No Azure API key found, skipping test")
+    model_kwargs = {"temperature": 0.5, "max_tokens": 2, "top_p": 1, "frequency_penalty": 0.5}
+    model_kwargs.update(azure_conf)
+    pn = PromptNode(model_name_or_path="text-davinci-003", api_key=azure_conf["api_key"], model_kwargs=model_kwargs)
     result = pn.prompt("question-generation", documents=["Berlin is the capital of Germany."])
     assert len(result) == 1 and len(result[0]) > 0
 
 
 @pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_open_ai_warn_if_max_tokens_is_too_short(caplog):
-    pm = PromptModel("text-davinci-003", api_key=os.environ["OPENAI_API_KEY"])
-    pn = PromptNode(pm)
+@pytest.mark.parametrize("prompt_model", ["openai", "azure"], indirect=True)
+def test_open_ai_warn_if_max_tokens_is_too_short(prompt_model, caplog):
+    skip_test_for_invalid_key(prompt_model)
+    pn = PromptNode(prompt_model)
     optional_davinci_params = {"temperature": 0.5, "max_tokens": 2, "top_p": 1, "frequency_penalty": 0.5}
     with caplog.at_level(logging.WARNING):
         _ = pn.prompt("question-generation", documents=["Berlin is the capital of Germany."], **optional_davinci_params)
@@ -295,10 +294,9 @@ def test_open_ai_warn_if_max_tokens_is_too_short(caplog):
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+@pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_stop_words(prompt_model):
-    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-        pytest.skip("No API key found for OpenAI, skipping test")
+    skip_test_for_invalid_key(prompt_model)
 
     # test stop words for both HF and OpenAI
     # set stop words in PromptNode
@@ -339,10 +337,9 @@ def test_stop_words(prompt_model):
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+@pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_simple_pipeline(prompt_model):
-    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-        pytest.skip("No API key found for OpenAI, skipping test")
+    skip_test_for_invalid_key(prompt_model)
 
     node = PromptNode(prompt_model, default_prompt_template="sentiment-analysis", output_variable="out")
 
@@ -353,10 +350,9 @@ def test_simple_pipeline(prompt_model):
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+@pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_complex_pipeline(prompt_model):
-    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-        pytest.skip("No API key found for OpenAI, skipping test")
+    skip_test_for_invalid_key(prompt_model)
 
     node = PromptNode(prompt_model, default_prompt_template="question-generation", output_variable="questions")
     node2 = PromptNode(prompt_model, default_prompt_template="question-answering")
@@ -370,10 +366,9 @@ def test_complex_pipeline(prompt_model):
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+@pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_simple_pipeline_with_topk(prompt_model):
-    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-        pytest.skip("No API key found for OpenAI, skipping test")
+    skip_test_for_invalid_key(prompt_model)
 
     node = PromptNode(prompt_model, default_prompt_template="question-generation", top_k=2)
 
@@ -385,13 +380,12 @@ def test_simple_pipeline_with_topk(prompt_model):
 
 
 @pytest.mark.integration
-@pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+@pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
 def test_complex_pipeline_with_qa(prompt_model):
     """Test the PromptNode where the `query` is a string instead of a list what the PromptNode would expects,
     because in a question-answering pipeline the retrievers need `query` as a string, so the PromptNode
     need to be able to handle the `query` being a string instead of a list."""
-    if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-        pytest.skip("No API key found for OpenAI, skipping test")
+    skip_test_for_invalid_key(prompt_model)
 
     prompt_template = PromptTemplate(
         name="question-answering-new",
@@ -696,13 +690,19 @@ def run_batch(
     assert "questions" in result["invocation_context"] and len(result["invocation_context"]["questions"]) > 0
 
 
-@pytest.mark.integration
-@pytest.mark.skipif(
-    not os.environ.get("OPENAI_API_KEY", None),
-    reason="Please export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
-)
-def test_complex_pipeline_with_all_features(tmp_path):
-    api_key = os.environ.get("OPENAI_API_KEY", None)
+@pytest.mark.parametrize("haystack_openai_config", ["openai", "azure"], indirect=True)
+def test_complex_pipeline_with_all_features(tmp_path, haystack_openai_config):
+    if not haystack_openai_config:
+        pytest.skip("No API key found, skipping test")
+
+    if "azure_base_url" in haystack_openai_config:
+        # don't change this indentation, it's important for the yaml to be valid
+        azure_conf_yaml_snippet = f"""
+                  azure_base_url: {haystack_openai_config['azure_base_url']}
+                  azure_deployment_name: {haystack_openai_config['azure_deployment_name']}
+        """
+    else:
+        azure_conf_yaml_snippet = ""
     with open(tmp_path / "tmp_config_with_prompt_template.yml", "w") as tmp_file:
         tmp_file.write(
             f"""
@@ -721,7 +721,8 @@ def test_complex_pipeline_with_all_features(tmp_path):
                 model_kwargs:
                   temperature: 0.9
                   max_tokens: 64
-                api_key: {api_key}
+                  {azure_conf_yaml_snippet}
+                api_key: {haystack_openai_config["api_key"]}
             - name: question_generation_template
               type: PromptTemplate
               params:
@@ -799,10 +800,9 @@ def test_complex_pipeline_with_multiple_same_prompt_node_components_yaml(tmp_pat
 
 class TestRunBatch:
     @pytest.mark.integration
-    @pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+    @pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
     def test_simple_pipeline_batch_no_query_single_doc_list(self, prompt_model):
-        if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-            pytest.skip("No API key found for OpenAI, skipping test")
+        skip_test_for_invalid_key(prompt_model)
 
         node = PromptNode(prompt_model, default_prompt_template="sentiment-analysis")
 
@@ -818,10 +818,9 @@ def test_simple_pipeline_batch_no_query_single_doc_list(self, prompt_model):
         assert "negative" in result["results"][1][0].casefold()
 
     @pytest.mark.integration
-    @pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+    @pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
     def test_simple_pipeline_batch_no_query_multiple_doc_list(self, prompt_model):
-        if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-            pytest.skip("No API key found for OpenAI, skipping test")
+        skip_test_for_invalid_key(prompt_model)
 
         node = PromptNode(prompt_model, default_prompt_template="sentiment-analysis", output_variable="out")
 
@@ -841,10 +840,9 @@ def test_simple_pipeline_batch_no_query_multiple_doc_list(self, prompt_model):
         assert "negative" in result["out"][1][0].casefold()
 
     @pytest.mark.integration
-    @pytest.mark.parametrize("prompt_model", ["hf", "openai"], indirect=True)
+    @pytest.mark.parametrize("prompt_model", ["hf", "openai", "azure"], indirect=True)
     def test_simple_pipeline_batch_query_multiple_doc_list(self, prompt_model):
-        if prompt_model.api_key is not None and not is_openai_api_key_set(prompt_model.api_key):
-            pytest.skip("No API key found for OpenAI, skipping test")
+        skip_test_for_invalid_key(prompt_model)
 
         prompt_template = PromptTemplate(
             name="question-answering-new",