Repository: simonkurtz-MSFT/python-openai-loadbalancer

Issue:
Support for authentication keys
Examples in the code/tests leverage using the managed identity authorization with via `token_provider`.
If the current solution supports connecting to endpoints with API keys, ti would be good to update documentation.

If the solution does not support connecting to endpoints via API keys at the moment, this is a feature request :)

Discussion:
pamelafox: @simonkurtz-MSFT I think it's fine for packages to support multiple forms of credentials. We avoid keys in our azd template samples since we provide the full infrastructure to create the RBAC roles, but I know devs sometimes use keys in other codebases.

Patch:
diff --git a/PACKAGE_README.md b/PACKAGE_README.md
index e69f02d..a4bc300 100644
--- a/PACKAGE_README.md
+++ b/PACKAGE_README.md
@@ -79,7 +79,9 @@ from openai_priority_loadbalancer import AsyncLoadBalancer, Backend
 *Importing `httpx` lets us use `httpx.Client` and `httpx.AsyncClient`. This avoids having to update openai to at least
 [1.17.0](https://github.com/openai/openai-python/releases/tag/v1.17.0). The `openai` properties for `DefaultHttpxClient` and `DefaultAsyncHttpxClient` are mere wrappers for `httpx.Client` and `httpx.AsyncClient`.*
 
-### Configuring the Backends and Load Balancer
+### Configuring the Backends and Load Balancer with a Token Provider
+
+**We strongly recommend the use of a managed identity in Azure and the use of the `AzureDefaultCredential` locally.** This section details that approach.
 
 1. Define a list of backends according to the *Load Balancer Backend Configuration* section below.
 
@@ -88,8 +90,8 @@ from openai_priority_loadbalancer import AsyncLoadBalancer, Backend
     ```python
     backends: List[Backend] = [
         Backend("oai-eastus.openai.azure.com", 1),
-        Backend("oai-eastus.openai.azure.com", 1, "/ai"),
         Backend("oai-southcentralus.openai.azure.com", 1)
+        Backend("oai-westus.openai.azure.com", 1, "/ai"),
     ]
     ```
 
@@ -121,6 +123,53 @@ from openai_priority_loadbalancer import AsyncLoadBalancer, Backend
     )
     ```
 
+### Configuring the Backends and Load Balancer with individual Azure OpenAI API Keys
+
+It's best to avoid using the Azure OpenAI instances' keys as that could a) accidentally leave credentials in your source code, and b) the keys are different for each instance, requiring maintenance, environment-specific keys, key rotations, etc.
+However, if you do need to use keys, it is possible to set them for each Azure OpenAI backend starting with release `1.1.0`.
+
+When a backend's `api_key` property is set, the `api-key` header will be replaced with the `<api_key>` value prior to sending the request to the corresponding Azure OpenAI instance. Please see below for examples.
+
+1. Define a list of backends according to the *Load Balancer Backend Configuration* section below. This includes the API key as the last parameter (below values are mock placeholders).
+
+   *Optionally, a path can be added (e.g. `"/ai"`), which gets prepended to the request path. This is an extraordinary, not a commonly needed functionality.*
+
+    ```python
+    backends: List[Backend] = [
+        Backends("oai-eastus.openai.azure.com", 1, None, 'c3d116584360f9960b38cccc5f44caba'),
+        Backends("oai-southcentralus.openai.azure.com", 1, None, '21c14252762502e8fc78b61e21db114f'),
+        Backends("oai-westus.openai.azure.com", 1, "/ai", 'd6370785453b2b9c331a94cb1b7aaa36')
+    ]
+    ```
+
+1. Instantiate the load balancer and inject a new httpx client with the load balancer as the new transport.
+
+    **Synchronous**
+
+    ```python
+    lb = LoadBalancer(backends)
+
+    client = AzureOpenAI(
+        azure_endpoint = f"https://{backends[0].host}",         # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+        api_key = "obtain_from_load_balancer",                  # the value is not used, but it must be set
+        api_version = "2024-04-01-preview",
+        http_client = httpx.Client(transport = lb)              # Inject the synchronous load balancer as the transport in a new default httpx client.
+    )
+    ```
+
+    **Asynchronous**
+
+    ```python
+    lb = AsyncLoadBalancer(backends)
+
+    client = AsyncAzureOpenAI(
+        azure_endpoint = f"https://{backends[0].host}",         # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+        api_key = "obtain_from_load_balancer",                  # the value is not used, but it must be set
+        api_version = "2024-04-01-preview",
+        http_client = httpx.AsyncClient(transport = lb)         # Inject the asynchronous load balancer as the transport in a new default async httpx client.
+    )
+    ```
+
 ## Load Balancer Backend Configuration
 
 At its core, the Load Balancer Backend configuration requires one or more backend hosts and a numeric priority starting at 1. Please take note that you define a host, not a URL.
@@ -191,6 +240,19 @@ backends = [
 ]
 ```
 
+### Backend Authentication
+
+While we strongly recommend the use of managed identities, it is possible to use the Azure OpenAI API keys for each respective Azure OpenAI instance. Note that you are solely responsible for the safeguarding and injection of these keys.
+
+```python
+# Define the backends and their priority
+backends = [
+    Backends("oai-eastus-xxxxxxxx.openai.azure.com", 1, None, 'c3d116584360f9960b38cccc5f44caba'),
+    Backends("oai-southcentralus-xxxxxxxx.openai.azure.com", 1 None, '21c14252762502e8fc78b61e21db114f'),
+    Backends("oai-westus-xxxxxxxx.openai.azure.com", 1, None, 'd6370785453b2b9c331a94cb1b7aaa36')
+]
+```
+
 ### Using the Load Balancer
 
 As these are the only changes to the [OpenAI Python API library](https://github.com/openai/openai-python) implementation, simply execute your python code.
diff --git a/README.md b/README.md
index 90f9c1b..a387667 100644
--- a/README.md
+++ b/README.md
@@ -38,7 +38,15 @@ It's also good to have some knowledge of authentication and identities.
 
 ## Authentication
 
-Locally, you can log into Azure via the CLI and the steps below and use the `AzureDefaultCredential` (what I use in my example). When deploying this application in Azure, it's recommended to use a managed identity for authentication. It's best to avoid using the Azure OpenAI instances' keys as that could a) accidentally leave credentials in your source code, and b) the keys are different for each instance, which would probably require expanding upon the `Backends` class. Best to just avoid keys.
+**We strongly recommend the use of a managed identity in Azure and the use of the `AzureDefaultCredential` locally.**
+
+Locally, you can log into Azure via the CLI and the steps below and use the `AzureDefaultCredential` (what I use in my example). When deploying this application in Azure, it's recommended to use a managed identity for authentication.
+
+### Azure OpenAI Keys
+
+It's best to avoid using the Azure OpenAI instances' keys as that could a) accidentally leave credentials in your source code, and b) the keys are different for each instance, requiring maintenance, environment-specific keys, key rotations, etc. However, if you need to use keys, it is possible to set them for each Azure OpenAI backend.
+
+When a backend's `api_key` property is set, the `api-key` header will be replaced with the `<api_key>` value prior to sending the request to the corresponding Azure OpenAI instance.
 
 ## Getting Started
 
@@ -51,9 +59,12 @@ Locally, you can log into Azure via the CLI and the steps below and use the `Azu
 
 ### Configuration
 
+Execute the following git command to ensure that updates to `config.py` are not tracked and therefore not committed. This prevents accidental check-ins of real keys and values:
+`git update-index --assume-unchanged config.py`
+
 For the load-balanced approach, please use the same model across all instances.
 
-1. Open [aoai.py](./aoai.py).
+1. Open [config.py](./config.py).
 1. Replace `<your-aoai-model>` with the name of your Azure OpenAI model.
 1. Replace `<your-aoai-instance>` with the primary/single Azure OpenAI instance.
 1. Replace `<your-aoai-instance-1>`, `<your-aoai-instance-2>`, `<your-aoai-instance-3>` with all the Azure OpenAI instances you want to load-balance across. Delete entries you don't need. See [Load Balancer Backend Configuration](#load-balancer-backend-configuration) for details.
@@ -193,3 +204,16 @@ backends = [
     Backends("oai-westus-xxxxxxxx.openai.azure.com", 3)
 ]
 ```
+
+### Backend Authentication
+
+While we strongly recommend the use of managed identities, it is possible to use the Azure OpenAI API keys for each respective Azure OpenAI instance. Note that you are solely responsible for the safeguarding and injection of these keys.
+
+```python
+# Define the backends and their priority
+backends = [
+    Backends("oai-eastus-xxxxxxxx.openai.azure.com", 1, None, 'c3d116584360f9960b38cccc5f44caba'),
+    Backends("oai-southcentralus-xxxxxxxx.openai.azure.com", 1 None, '21c14252762502e8fc78b61e21db114f'),
+    Backends("oai-westus-xxxxxxxx.openai.azure.com", 1, None, 'd6370785453b2b9c331a94cb1b7aaa36')
+]
+```
diff --git a/aoai.py b/aoai.py
index 29621b0..f24ef28 100644
--- a/aoai.py
+++ b/aoai.py
@@ -4,7 +4,6 @@ import asyncio
 import logging
 import time
 import traceback
-from typing import List
 from datetime import datetime
 # Using httpx.Client and httpx.AsyncClient avoids having to update openai to 1.17.1 or newer.
 # The openai properties for DefaultHttpxClient and DefaultAsyncHttpxClient are mere wrappers for httpx.Client and httpx.AsyncClient.
@@ -12,7 +11,8 @@ from datetime import datetime
 import httpx
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 from openai import AzureOpenAI, AsyncAzureOpenAI, NotFoundError, APIError
-from src.openai_priority_loadbalancer.openai_priority_loadbalancer import AsyncLoadBalancer, LoadBalancer, Backend
+from src.openai_priority_loadbalancer.openai_priority_loadbalancer import AsyncLoadBalancer, LoadBalancer
+import config
 
 ##########################################################################################################################################################
 
@@ -22,21 +22,15 @@ class TestExecutions:
     """Class representing the tests that can be performed."""
 
     def __init__(self):
-        self.standard                   = True
-        self.load_balanced              = True
-        self.async_load_balanced        = True
-        self.stream_load_balanced       = True
-        self.async_stream_load_balanced = True
-
-LOG_LEVEL                               = logging.INFO
-NUM_OF_REQUESTS                         = 5
-MODEL                                   = "<your-aoai-model>"  # the model, also known as the Deployment in Azure OpenAI, is common across standard and load-balanced requests
-AZURE_ENDPOINT                          = "https://oai-eastus-xxxxxxxx.openai.azure.com"
-backends: List[Backend] = [
-    Backend("oai-eastus-xxxxxxxx.openai.azure.com", 1),
-    Backend("oai-southcentralus-xxxxxxxx.openai.azure.com", 1),
-    Backend("oai-westus-xxxxxxxx.openai.azure.com", 1)
-]
+        self.standard                           = True
+        self.load_balanced                      = True
+        self.load_balanced_with_api_keys        = True
+        self.async_load_balanced                = True
+        self.async_load_balanced_with_api_keys  = True
+        self.stream_load_balanced               = True
+        self.async_stream_load_balanced         = True
+
+LOG_LEVEL = logging.INFO     # change to DEBUG for detailed information
 
 ##########################################################################################################################################################
 
@@ -61,7 +55,7 @@ def send_request(num_of_requests: int, azure_endpoint: str):
             print(f"{datetime.now()}: Standard request {i+1}/{num_of_requests}")
 
             response = client.chat.completions.create(
-                model = MODEL,
+                model = config.MODEL,
                 messages = [
                     {"role": "system", "content": "You are a helpful assistant."},
                     {"role": "user", "content": "Does Azure OpenAI support customer managed keys?"}
@@ -87,10 +81,10 @@ def send_loadbalancer_request(num_of_requests: int):
 
     try:
         # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
-        lb = LoadBalancer(backends)
+        lb = LoadBalancer(config.backends)
 
         client = AzureOpenAI(
-            azure_endpoint = f"https://{backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            azure_endpoint = f"https://{config.backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
             azure_ad_token_provider = token_provider,
             api_version = "2024-04-01-preview",
             http_client = httpx.Client(transport = lb)      # Inject the load balancer as the transport in a new default httpx client
@@ -101,7 +95,58 @@ def send_loadbalancer_request(num_of_requests: int):
 
             try:
                 response = client.chat.completions.create(
-                    model = MODEL,
+                    model = config.MODEL,
+                    messages = [
+                        {"role": "system", "content": "You are a helpful assistant."},
+                        {"role": "user", "content": "Does Azure OpenAI support customer managed keys?"}
+                    ]
+                )
+
+                success_counter += 1
+                print(f"{datetime.now()}:\n{response}\n\n\n")
+            except APIError as e:
+                if e.code == 429:
+                    print(f"{datetime.now()}: Rate limit exceeded. Python OpenAI Library has exhausted all of its retries.")
+                else:
+                    print(f"{datetime.now()}: Python OpenAI Library request failure.")
+
+                failure_counter += 1
+            except Exception:
+                traceback.print_exc()
+                failure_counter += 1
+
+            counter += 1
+
+    except NotFoundError as e:
+        print("openai.NotFoundError:", vars(e))
+        traceback.print_exc()
+    except Exception as e:
+        print("Exception:", vars(e))
+        traceback.print_exc()
+
+# Load-balanced Azure OpenAI Implementation (Multiple Backends)
+def send_loadbalancer_request_with_api_keys(num_of_requests: int):
+    """Function to send load-balanced requests to the Azure OpenAI API using API keys."""
+
+    global counter, failure_counter, success_counter
+
+    try:
+        # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
+        lb = LoadBalancer(config.backends_with_api_keys)
+
+        client = AzureOpenAI(
+            azure_endpoint = f"https://{config.backends_with_api_keys[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            api_key = "obtain_from_load_balancer",          # the value is not used, but it must be set
+            api_version = "2024-04-01-preview",
+            http_client = httpx.Client(transport = lb)      # Inject the load balancer as the transport in a new default httpx client
+        )
+
+        for i in range(num_of_requests):
+            print(f"{datetime.now()}: LoadBalancer request {i+1}/{num_of_requests}")
+
+            try:
+                response = client.chat.completions.create(
+                    model = config.MODEL,
                     messages = [
                         {"role": "system", "content": "You are a helpful assistant."},
                         {"role": "user", "content": "Does Azure OpenAI support customer managed keys?"}
@@ -137,10 +182,10 @@ async def send_async_loadbalancer_request(num_of_requests: int):
 
     try:
         # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
-        lb = AsyncLoadBalancer(backends)
+        lb = AsyncLoadBalancer(config.backends)
 
         client = AsyncAzureOpenAI(
-            azure_endpoint = f"https://{backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            azure_endpoint = f"https://{config.backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
             azure_ad_token_provider = token_provider,
             api_version = "2024-04-01-preview",
             http_client = httpx.AsyncClient(transport = lb) # Inject the load balancer as the transport in a new default httpx client
@@ -151,7 +196,57 @@ async def send_async_loadbalancer_request(num_of_requests: int):
 
             try:
                 response = await client.chat.completions.create(
-                    model = MODEL,
+                    model = config.MODEL,
+                    messages = [
+                        {"role": "system", "content": "You are a helpful assistant."},
+                        {"role": "user", "content": "Does Azure OpenAI support customer managed keys?"}
+                    ]
+                )
+
+                success_counter += 1
+                print(f"{datetime.now()}:\n{response}\n\n\n")
+            except APIError as e:
+                if e.code == 429:
+                    print(f"{datetime.now()}: Rate limit exceeded. Python OpenAI Library has exhausted all of its retries.")
+                else:
+                    print(f"{datetime.now()}: Python OpenAI Library request failure.")
+
+                failure_counter += 1
+            except Exception:
+                traceback.print_exc()
+                failure_counter += 1
+
+            counter += 1
+
+    except NotFoundError as e:
+        print("openai.NotFoundError:", vars(e))
+        traceback.print_exc()
+    except Exception as e:
+        print("Exception:", vars(e))
+        traceback.print_exc()
+
+async def send_async_loadbalancer_request_with_api_keys(num_of_requests: int):
+    """Function to send load-balanced requests to the Azure OpenAI API using API keys."""
+
+    global counter, failure_counter, success_counter
+
+    try:
+        # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
+        lb = AsyncLoadBalancer(config.backends_with_api_keys)
+
+        client = AsyncAzureOpenAI(
+            azure_endpoint = f"https://{config.backends_with_api_keys[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            api_key = "obtain_from_load_balancer",          # the value is not used, but it must be set
+            api_version = "2024-04-01-preview",
+            http_client = httpx.AsyncClient(transport = lb) # Inject the load balancer as the transport in a new default httpx client
+        )
+
+        for i in range(num_of_requests):
+            print(f"{datetime.now()}: Async LoadBalancer request {i+1}/{num_of_requests}")
+
+            try:
+                response = await client.chat.completions.create(
+                    model = config.MODEL,
                     messages = [
                         {"role": "system", "content": "You are a helpful assistant."},
                         {"role": "user", "content": "Does Azure OpenAI support customer managed keys?"}
@@ -188,10 +283,10 @@ def send_stream_loadbalancer_request(num_of_requests: int):
 
     try:
         # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
-        lb = LoadBalancer(backends)
+        lb = LoadBalancer(config.backends)
 
         client = AzureOpenAI(
-            azure_endpoint = f"https://{backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            azure_endpoint = f"https://{config.backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
             azure_ad_token_provider = token_provider,
             api_version = "2024-04-01-preview",
             http_client = httpx.Client(transport = lb)      # Inject the load balancer as the transport in a new default httpx client
@@ -204,7 +299,7 @@ def send_stream_loadbalancer_request(num_of_requests: int):
                 stream_start_time = time.time()
 
                 response = client.chat.completions.create(
-                    model = MODEL,
+                    model = config.MODEL,
                     messages = [
                         {"role": "system", "content": "You are a helpful assistant."},
                         {'role': 'user', 'content': 'Count to 5, with a comma between each number and no newlines. E.g., 1, 2, 3, ...'}
@@ -261,10 +356,10 @@ async def send_async_stream_loadbalancer_request(num_of_requests: int):
 
     try:
         # Instantiate the LoadBalancer class and create a new https client with the load balancer as the injected transport.
-        lb = AsyncLoadBalancer(backends)
+        lb = AsyncLoadBalancer(config.backends)
 
         client = AsyncAzureOpenAI(
-            azure_endpoint = f"https://{backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
+            azure_endpoint = f"https://{config.backends[0].host}", # Must be seeded, so we use the first host. It will get overwritten by the load balancer.
             azure_ad_token_provider = token_provider,
             api_version = "2024-04-01-preview",
             http_client = httpx.AsyncClient(transport = lb) # Inject the load balancer as the transport in a new default httpx client
@@ -277,7 +372,7 @@ async def send_async_stream_loadbalancer_request(num_of_requests: int):
                 stream_start_time = time.time()
 
                 response = await client.chat.completions.create(
-                    model = MODEL,
+                    model = config.MODEL,
                     messages = [
                         {"role": "system", "content": "You are a helpful assistant."},
                         {'role': 'user', 'content': 'Count to 5, with a comma between each number and no newlines. E.g., 1, 2, 3, ...'}
@@ -332,7 +427,7 @@ async def send_async_stream_loadbalancer_request(num_of_requests: int):
 
 # >>> TEST HARNESS <<<
 
-success_counter = failure_counter = counter = 0
+success_counter = failure_counter = counter = 0 # pylint: disable=C0103
 
 # Set up the logger: https://www.machinelearningplus.com/python/python-logging-guide/
 logging.basicConfig(
@@ -342,13 +437,13 @@ logging.basicConfig(
 )
 
 # Ensure that variables are set.
-if MODEL == "<your-aoai-model>":
+if config.MODEL == "<your-aoai-model>":
     raise ValueError("MODEL must be set to a valid AOAI model.\n")
 
-if "xxxxxxxx" in AZURE_ENDPOINT:
+if "xxxxxxxx" in config.AZURE_ENDPOINT:
     raise ValueError("AZURE_ENDPOINT must be set to a valid endpoint.\n")
 
-for backend in backends:
+for backend in config.backends:
     if "xxxxxxxx" in backend.host:
         raise ValueError(f"Backend {backend.host} must be set to a valid endpoint.\n")
 
@@ -359,35 +454,49 @@ test_executions = TestExecutions()
 if test_executions.standard:
     print(f"\nStandard Requests\n{'-' * 17}\n")
     start_time = time.time()
-    send_request(NUM_OF_REQUESTS, AZURE_ENDPOINT)
+    send_request(config.NUM_OF_REQUESTS, config.AZURE_ENDPOINT)
     end_time = time.time()
 
 # 2: Load-balanced requests to one or more AOAI backends
 if test_executions.load_balanced:
     print(f"\nLoad Balanced Requests\n{'-' * 22}\n")
     lb_start_time = time.time()
-    send_loadbalancer_request(NUM_OF_REQUESTS)
+    send_loadbalancer_request(config.NUM_OF_REQUESTS)
     lb_end_time = time.time()
 
-# 3: Async Load-balanced requests to one or more AOAI backends
+# 3: Load-balanced requests to one or more AOAI backends with API keys
+if test_executions.load_balanced_with_api_keys:
+    print(f"\nLoad Balanced Requests With API Keys\n{'-' * 36}\n")
+    lb_with_api_keys_start_time = time.time()
+    send_loadbalancer_request_with_api_keys(config.NUM_OF_REQUESTS)
+    lb_with_api_keys_end_time = time.time()
+
+# 4: Async Load-balanced requests to one or more AOAI backends
 if test_executions.async_load_balanced:
     print(f"\nAsync Load Balanced Requests\n{'-' * 28}\n")
     async_lb_start_time = time.time()
-    asyncio.run(send_async_loadbalancer_request(NUM_OF_REQUESTS))
+    asyncio.run(send_async_loadbalancer_request(config.NUM_OF_REQUESTS))
     async_lb_end_time = time.time()
 
-# 4: Load-balanced streaming requests to one or more AOAI backends
+# 5: Async Load-balanced requests to one or more AOAI backends with API keys
+if test_executions.async_load_balanced_with_api_keys:
+    print(f"\nAsync Load Balanced Requests With API Keys\n{'-' * 42}\n")
+    async_lb_with_api_keys_start_time = time.time()
+    asyncio.run(send_async_loadbalancer_request_with_api_keys(config.NUM_OF_REQUESTS))
+    async_lb_with_api_keys_end_time = time.time()
+
+# : Load-balanced streaming requests to one or more AOAI backends
 if test_executions.stream_load_balanced:
     print(f"\nStream Load Balanced Requests\n{'-' * 29}\n")
     stream_lb_start_time = time.time()
-    send_stream_loadbalancer_request(NUM_OF_REQUESTS)
+    send_stream_loadbalancer_request(config.NUM_OF_REQUESTS)
     stream_lb_end_time = time.time()
 
 # 5: Async Load-balanced streaming requests to one or more AOAI backends
 if test_executions.async_stream_load_balanced:
     print(f"\nStream Async Load Balanced Requests\n{'-' * 35}\n")
     async_stream_lb_start_time = time.time()
-    asyncio.run(send_async_stream_loadbalancer_request(NUM_OF_REQUESTS))
+    asyncio.run(send_async_stream_loadbalancer_request(config.NUM_OF_REQUESTS))
     async_stream_lb_end_time = time.time()
 
 # Statistics
@@ -395,24 +504,28 @@ WIDTH = 16
 SECONDS_WIDTH = WIDTH - 8
 
 print(f"\n{'*' * 100}\n")
-print(f"Requests per approach                           : {str(NUM_OF_REQUESTS).rjust(WIDTH)}")
-print(f"Number of approaches                            : {str(sum(1 for value in vars(test_executions).values() if value is True)).rjust(WIDTH)}\n")
+print(f"Requests per approach                                   : {str(config.NUM_OF_REQUESTS).rjust(WIDTH)}")
+print(f"Number of approaches                                    : {str(sum(1 for value in vars(test_executions).values() if value is True)).rjust(WIDTH)}\n")
 
-print(f"Total requests                                  : {str(counter).rjust(WIDTH)}")
-print(f"Total successful requests                       : {str(success_counter).rjust(WIDTH)}")
-print(f"Total failed requests                           : {str(failure_counter).rjust(WIDTH)}")
-print(f"Total successful requests percentage            : {('{:.2%}'.format(success_counter / counter)).rjust(WIDTH)}")
-print(f"Total Failed requests percentage                : {('{:.2%}'.format(failure_counter / counter)).rjust(WIDTH)}\n")
+print(f"Total requests                                          : {str(counter).rjust(WIDTH)}")
+print(f"Total successful requests                               : {str(success_counter).rjust(WIDTH)}")
+print(f"Total failed requests                                   : {str(failure_counter).rjust(WIDTH)}")
+print(f"Total successful requests percentage                    : {('{:.2%}'.format(success_counter / counter)).rjust(WIDTH)}")
+print(f"Total Failed requests percentage                        : {('{:.2%}'.format(failure_counter / counter)).rjust(WIDTH)}\n")
 
 if test_executions.standard:
-    print(f"Single instance operation duration              : {end_time - start_time:>{SECONDS_WIDTH}.2f} seconds")
+    print(f"Single instance operation duration                      : {end_time - start_time:>{SECONDS_WIDTH}.2f} seconds")
 if test_executions.load_balanced:
-    print(f"Load-balancer operation duration                : {lb_end_time - lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+    print(f"Load-balancer operation duration                        : {lb_end_time - lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+if test_executions.load_balanced_with_api_keys:
+    print(f"Load-balancer with API keys operation duration          : {lb_with_api_keys_end_time - lb_with_api_keys_start_time:>{SECONDS_WIDTH}.2f} seconds")
 if test_executions.async_load_balanced:
-    print(f"Async Load-balancer operation duration          : {async_lb_end_time - async_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+    print(f"Async Load-balancer operation duration                  : {async_lb_end_time - async_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+if test_executions.async_load_balanced_with_api_keys:
+    print(f"Async Load-balancer with API keys operation duration    : {async_lb_with_api_keys_end_time - async_lb_with_api_keys_start_time:>{SECONDS_WIDTH}.2f} seconds")
 if test_executions.stream_load_balanced:
-    print(f"Stream Load-balancer operation duration         : {stream_lb_end_time - stream_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+    print(f"Stream Load-balancer operation duration                 : {stream_lb_end_time - stream_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
 if test_executions.async_stream_load_balanced:
-    print(f"Stream Async Load-balancer operation duration   : {async_stream_lb_end_time - async_stream_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
+    print(f"Stream Async Load-balancer operation duration           : {async_stream_lb_end_time - async_stream_lb_start_time:>{SECONDS_WIDTH}.2f} seconds")
 
 print("\n\n")
diff --git a/config.py b/config.py
new file mode 100644
index 0000000..95fe442
--- /dev/null
+++ b/config.py
@@ -0,0 +1,23 @@
+# pylint: disable=C0115
+
+"""Module providing configuration settings for the OpenAI Priority Load Balancer test harness."""
+
+from typing import List
+from src.openai_priority_loadbalancer.openai_priority_loadbalancer import Backend
+
+class Config:
+    NUM_OF_REQUESTS = 5
+    MODEL           = "<your-aoai-model>"  # the model, also known as the Deployment in Azure OpenAI, is common across standard and load-balanced requests
+    AZURE_ENDPOINT  = "https://oai-eastus-xxxxxxxx.openai.azure.com"
+
+    backends: List[Backend] = [
+        Backend("oai-eastus-xxxxxxxx.openai.azure.com", 1),
+        Backend("oai-southcentralus-xxxxxxxx.openai.azure.com", 1),
+        Backend("oai-westus-xxxxxxxx.openai.azure.com", 1)
+    ]
+
+    backends_with_api_keys: List[Backend] = [
+        Backend("oai-eastus-xxxxxxxx.openai.azure.com", 1, None, 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'),
+        Backend("oai-southcentralus-xxxxxxxx.openai.azure.com", 1, None, 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'),
+        Backend("oai-westus-xxxxxxxx.openai.azure.com", 1, None, 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx')
+    ]
diff --git a/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py b/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
index 493de0b..324e68d 100644
--- a/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
+++ b/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
@@ -13,14 +13,15 @@ class Backend:
     """Class representing a backend object used with Azure OpenAI, etc."""
 
     # Constructor
-    def __init__(self, host: str, priority: int, path: str = None):
+    def __init__(self, host: str, priority: int, path: str = None, api_key: str = None):
         # Public instance variables
-        self.host = host
-        self.path = '' if path is None else path
-        self.is_throttling = False
-        self.priority = priority
-        self.retry_after = datetime.min
-        self.successful_call_count = 0
+        self.api_key: str = api_key
+        self.host: str = host
+        self.is_throttling: bool = False
+        self.path: str = '' if path is None else path
+        self.priority: int = priority
+        self.retry_after: datetime = datetime.min
+        self.successful_call_count: int = 0
 
 # Reference design at https://github.com/encode/httpx/blob/master/httpx/_transports/base.py
 # BaseLoadBalancer providing functionality to both synchronous and asynchronous load balancers
@@ -158,7 +159,8 @@ class BaseLoadBalancer():
 
         backend: Backend = self.backends[backend_index]
 
-        # Modify the request. Note that only the URL and Host header are being modified on the original request object. We make the smallest incision possible to avoid side effects.
+        # Modify the request. Note that only the URL and Host header are being modified on the original request object. Additionally, if an API key is defined, set
+        # the api-key header with that value. We make the smallest incision possible to avoid side effects.
         # Update URL and host header as both must match the backend server.
         request.url = request.url.copy_with(host = backend.host)
 
@@ -181,7 +183,11 @@ class BaseLoadBalancer():
         request.headers = request.headers.copy()    # We need to create a mutable copy of the headers before we modify and assign them back to the request object.
         request.headers['host'] = backend.host
 
-        self._log.debug("URL %s and host %s", request.url, request.headers['host'])
+        if backend.api_key is not None and backend.api_key != "":
+            request.headers['api-key'] = backend.api_key
+            self._log.debug("URL = [%s]; host header = [%s]; api-key header = [%s]", request.url, request.headers['host'], request.headers['api-key'])
+        else:
+            self._log.debug("URL = [%s]; host header = [%s]", request.url, request.headers['host'])
 
     def _return_429(self) -> httpx.Response:
         """Return an HTTP 429 response with a Retry-After header value. This is returned to the caller of this load balancer when no backends are available."""


Test patch:
diff --git a/tests/lib/test_openai_priority_loadbalancer.py b/tests/lib/test_openai_priority_loadbalancer.py
index 45e8036..681b598 100644
--- a/tests/lib/test_openai_priority_loadbalancer.py
+++ b/tests/lib/test_openai_priority_loadbalancer.py
@@ -59,7 +59,7 @@ def create_final_request_options() -> FinalRequestOptions:
 # Factory fixture for backends
 @pytest.fixture
 def backends_factory():
-    def _backends_factory(priority: int, is_throttling: bool, retry_after: int, path: str = None):
+    def _backends_factory(priority: int, is_throttling: bool, retry_after: int, path: str = None, api_key: str = None):
         # Start with a basic list of backends that will be modified depending on the passed arguments.
         backends: List[Backend] = [
             Backend("oai-eastus.openai.azure.com", 1),
@@ -67,11 +67,16 @@ def backends_factory():
             Backend("oai-westus.openai.azure.com", 1),
         ]
 
-        for backend, _priority, throttling, secs, _path in zip(backends, priority, is_throttling, retry_after, path):
-            backend.priority = _priority
+        for backend, _priority, throttling, secs, _path, _api_key in zip(backends, priority, is_throttling, retry_after, path, api_key):
             backend.is_throttling = throttling
+            backend.priority = _priority
+
+            if _api_key is not None:
+                backend.api_key = _api_key
+
             if secs is not None:
                 backend.retry_after = datetime.now(timezone.utc) + timedelta(seconds = secs)
+
             if _path is not None:
                 backend.path = _path
 
@@ -81,27 +86,31 @@ def backends_factory():
 
 @pytest.fixture
 def backends_same_priority(backends_factory) -> List[Backend]:
-    return backends_factory([1, 1, 1], [False, False, False], [None, None, None], [None, None, None])
+    return backends_factory([1, 1, 1], [False, False, False], [None, None, None], [None, None, None], [None, None, None])
 
 @pytest.fixture
 def backends_same_priority_custom_paths(backends_factory) -> List[Backend]:
-    return backends_factory([1, 1, 1], [False, False, False], [None, None, None], ["/ai", "ai/", "ai"])
+    return backends_factory([1, 1, 1], [False, False, False], [None, None, None], ["/ai", "ai/", "ai"], [None, None, None])
+
+@pytest.fixture
+def backends_same_priority_api_keys(backends_factory) -> List[Backend]:
+    return backends_factory([1, 1, 1], [False, False, False], [None, None, None], [None, None, None], ["c3d116584360f9960b38cccc5f44caba", "21c14252762502e8fc78b61e21db114f", "d6370785453b2b9c331a94cb1b7aaa36"])
 
 @pytest.fixture
 def backends_tiered_priority(backends_factory) -> List[Backend]:
-    return backends_factory([1, 2, 2], [False, False, False], [None, None, None], [None, None, None])
+    return backends_factory([1, 2, 2], [False, False, False], [None, None, None], [None, None, None], [None, None, None])
 
 @pytest.fixture
 def backends_0_and_1_throttling(backends_factory) -> List[Backend]:
-    return backends_factory([1, 1, 1], [True, True, False], [3, 1, None], [None, None, None])
+    return backends_factory([1, 1, 1], [True, True, False], [3, 1, None], [None, None, None], [None, None, None])
 
 @pytest.fixture
 def all_backends_throttling(backends_factory) -> List[Backend]:
-    return backends_factory([1, 1, 1], [True, True, True], [3, 1, 5], [None, None, None])
+    return backends_factory([1, 1, 1], [True, True, True], [3, 1, 5], [None, None, None], [None, None, None])
 
 @pytest.fixture
 def priority_backend_0_throttling(backends_factory) -> List[Backend]:
-    return backends_factory([1, 2, 2], [True, False, False], [3, None, None], [None, None, None])
+    return backends_factory([1, 2, 2], [True, False, False], [3, None, None], [None, None, None], [None, None, None])
 
 @pytest.fixture(params=[backends_same_priority, backends_tiered_priority, backends_0_and_1_throttling, priority_backend_0_throttling])
 def success_backends(request):
@@ -121,6 +130,10 @@ def client_same_priority(backends_same_priority) -> AzureOpenAI:
 def client_same_priority_custom_paths(backends_same_priority_custom_paths) -> AzureOpenAI:
     return create_client(backends_same_priority_custom_paths)
 
+@pytest.fixture
+def client_same_priority_api_keys(backends_same_priority_api_keys) -> AzureOpenAI:
+    return create_client(backends_same_priority_api_keys)
+
 @pytest.fixture
 def client_successful_backends(success_backends) -> AzureOpenAI:
     return create_client(success_backends)
@@ -139,6 +152,10 @@ def async_client_same_priority(backends_same_priority) -> AsyncAzureOpenAI:
 def async_client_same_priority_custom_paths(backends_same_priority_custom_paths) -> AsyncAzureOpenAI:
     return create_async_client(backends_same_priority_custom_paths)
 
+@pytest.fixture
+def async_client_same_priority_api_keys(backends_same_priority_api_keys) -> AsyncAzureOpenAI:
+    return create_async_client(backends_same_priority_api_keys)
+
 @pytest.fixture
 def async_client_successful_backends(success_backends) -> AsyncAzureOpenAI:
     return create_async_client(success_backends)
@@ -275,13 +292,32 @@ class TestSynchronous:
 
             client._client._transport.handle_request(req)
 
-            #assert response.status_code == 200
             assert req.url in (
                 'https://oai-eastus.openai.azure.com/ai/openai/completions?api-version=2024-04-01-preview',
                 'https://oai-westus.openai.azure.com/ai/openai/completions?api-version=2024-04-01-preview',
                 'https://oai-southcentralus.openai.azure.com/ai/openai/completions?api-version=2024-04-01-preview'
             )
 
+    @pytest.mark.loadbalancer
+    def test_loadbalancer_use_api_keys(self, client_same_priority_api_keys):
+        client = client_same_priority_api_keys
+
+        # Create a mock response for the transport
+        mock_response = httpx.Response(200)
+
+        with patch('httpx.Client.send', return_value = mock_response):
+            req = client._build_request(create_final_request_options())
+
+            assert req.url == 'https://foo.openai.azure.com/openai/completions?api-version=2024-04-01-preview'
+
+            client._client._transport.handle_request(req)
+
+            assert req.headers['api-key'] in (
+                'c3d116584360f9960b38cccc5f44caba',
+                '21c14252762502e8fc78b61e21db114f',
+                'd6370785453b2b9c331a94cb1b7aaa36'
+            )
+
     @pytest.mark.loadbalancer
     def test_loadbalancer_handle_failure_requests(self, client_failure_backends):
         client = client_failure_backends
@@ -466,6 +502,27 @@ class TestAsynchronous:
                 'https://oai-southcentralus.openai.azure.com/ai/openai/completions?api-version=2024-04-01-preview'
             )
 
+    @pytest.mark.asyncio
+    @pytest.mark.loadbalancer
+    async def test_async_loadbalancer_use_api_keys(self, async_client_same_priority_api_keys):
+        client = async_client_same_priority_api_keys
+
+        # Create a mock response for the transport
+        mock_response = httpx.Response(200)
+
+        with patch('httpx.AsyncClient.send', return_value = mock_response):
+            req = client._build_request(create_final_request_options())
+
+            assert req.url == 'https://foo.openai.azure.com/openai/completions?api-version=2024-04-01-preview'
+
+            await client._client._transport.handle_async_request(req)
+
+            assert req.headers['api-key'] in (
+                'c3d116584360f9960b38cccc5f44caba',
+                '21c14252762502e8fc78b61e21db114f',
+                'd6370785453b2b9c331a94cb1b7aaa36'
+            )
+
     @pytest.mark.asyncio
     @pytest.mark.async_loadbalancer
     async def test_async_loadbalancer_handle_failure_requests(self, async_client_failure_backends):