Repository: simonkurtz-MSFT/python-openai-loadbalancer

Issue:
Bug Report: AttributeError when closing the httpx.AsyncClient
**Description:**
When trying to close the `httpx.AsyncClient` in an asynchronous load balancing setup, an `AttributeError` is raised:
```
AttributeError: 'AsyncLoadBalancer' object has no attribute 'aclose'
```

This happens because the custom `AsyncLoadBalancer` object does not implement the `aclose()` method, which is required when closing the transport in `httpx.AsyncClient`.

**Steps to Reproduce:**
1. Setup an `AsyncLoadBalancer` as the transport for `httpx.AsyncClient`.
2. Make an async request using the `client.chat.completions.create()` method.
3. Try to close the `httpx.AsyncClient` with `await async_client.aclose()`.
4. Observe the traceback showing an `AttributeError`.

**Expected Behavior:**
The `async_client.aclose()` should close without raising an exception, properly releasing resources used by the transport (in this case, the `AsyncLoadBalancer`).

**Actual Behavior:**
An exception is raised when trying to close the `httpx.AsyncClient` because the custom load balancer (`AsyncLoadBalancer`) lacks the `aclose()` method.

**Traceback:**
```
Traceback (most recent call last):
  File "/workspaces/NeuralExtractor/tests/load_balancer.py", line 50, in <module>
    asyncio.run(main())
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/asyncio/base_events.py", line 654, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/workspaces/NeuralExtractor/tests/load_balancer.py", line 46, in main
    await async_client.aclose()
  File "/workspaces/NeuralExtractor/src/quartapp/.venv/lib/python3.11/site-packages/httpx/_client.py", line 2018, in aclose
    await self._transport.aclose()
          ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'AsyncLoadBalancer' object has no attribute 'aclose'
```

**Possible Cause:**
The `AsyncLoadBalancer` class does not implement the `aclose()` method required by `httpx.AsyncClient` to close the transport gracefully.

**Suggested Fix:**
Implement the `aclose()` method in the `AsyncLoadBalancer` class to close any resources properly when shutting down the client. Here's an example of how this might be added:

```python
class AsyncLoadBalancer:
    # Your current methods and attributes

    async def aclose(self):
        # Close any active connections or resources here
        pass
```

**Additional Context:**
This issue is observed when working with OpenAI's `AsyncAzureOpenAI` client and `httpx` in combination with custom load balancing logic. Properly closing the client is necessary to prevent resource leaks or unhandled exceptions during shutdown.


Discussion:

Patch:
diff --git a/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py b/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
index 324e68d..736098c 100644
--- a/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
+++ b/src/openai_priority_loadbalancer/openai_priority_loadbalancer.py
@@ -3,7 +3,7 @@
 # Python Standard Library
 import logging
 import random
-from typing import List
+from typing import List, Union
 from datetime import datetime, MAXYEAR, MINYEAR, timedelta, timezone
 
 # Third-Party Libraries
@@ -29,7 +29,7 @@ class BaseLoadBalancer():
     """Logically abstracts the BaseLoadBalancer class which should be inherited by the synchronous and asynchronous load balancer classes."""
 
     # Constructor
-    def __init__(self, transport: httpx.BaseTransport, backends: List[Backend]):
+    def __init__(self, transport: Union[httpx.BaseTransport, httpx.AsyncBaseTransport], backends: List[Backend]):
         # Public instance variables
         self.backends = backends
 
@@ -39,6 +39,12 @@ class BaseLoadBalancer():
         self._available_backends = 1
         self._transport = transport
 
+    # Magic Methods
+
+    # If a method in the BaseTransport or AsyncBaseTransport classes is not found, it will be looked up in base _transport object.
+    def __getattr__(self, name):
+        return getattr(self._transport, name)
+
     # "Protected" Methods
     def _check_throttling(self) -> None:
         """Check if any backend is throttling and reset if necessary."""


Test patch:
diff --git a/tests/lib/test_openai_priority_loadbalancer.py b/tests/lib/test_openai_priority_loadbalancer.py
index 681b598..bae260c 100644
--- a/tests/lib/test_openai_priority_loadbalancer.py
+++ b/tests/lib/test_openai_priority_loadbalancer.py
@@ -382,6 +382,26 @@ class TestSynchronous:
             # Assert that the final response status code was 400
             assert response.status_code == 400
 
+    @pytest.mark.loadbalancer
+    def test_loadbalancer_loadbalancer_close(self, client_same_priority):
+        client = client_same_priority
+
+        # 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())
+            response = client._client._transport.handle_request(req)
+
+            assert response.status_code == 200
+
+            # Assert that the transport is not closed, then close it, then assert that it was closed.
+            assert client._client._transport.is_closed is False
+
+            client._client._transport.close()
+
+            assert client._client._transport.is_closed is True
+
 # Asynchronous Tests
 
 class TestAsynchronous:
@@ -591,3 +611,24 @@ class TestAsynchronous:
 
             # Assert that the final response status code was 400
             assert response.status_code == 400
+
+    @pytest.mark.asyncio
+    @pytest.mark.async_loadbalancer
+    async def test_async_loadbalancer_close(self, async_client_same_priority):
+        client = async_client_same_priority
+
+        # 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())
+            response = await client._client._transport.handle_async_request(req)
+
+            assert response.status_code == 200
+
+            # Assert that the transport is not closed, then close it, then assert that it was closed.
+            assert client._client._transport.is_closed is False
+
+            await client._client._transport.aclose()
+
+            assert client._client._transport.is_closed is True