fix: return 502 for errors caused by particular upstream configuration

aidial_adapter_openai/exception_handlers.py

@@ -22,6 +22,14 @@ def to_adapter_exception(exc: Exception) -> AdapterException:
     return _expose_error_message_to_user(_convert_to_adapter_exception(exc))
 
 
+def _get_status_code(exc: APIStatusError) -> int:
+    has_api_key_header = "api-key" in exc.request.headers
+    code = exc.status_code
+    if code == 500 or (code in (401, 403) and has_api_key_header):
+        code = 502
+    return code
+
+
 def _convert_to_adapter_exception(exc: Exception) -> AdapterException:
 
     if isinstance(exc, (DialException, ResponseWrapper)):
@@ -50,8 +58,10 @@ def _convert_to_adapter_exception(exc: Exception) -> AdapterException:
         if "Content-Encoding" in httpx_headers:
             del httpx_headers["Content-Encoding"]
 
+        status_code = _get_status_code(exc)
+
         return parse_adapter_exception(
-            status_code=r.status_code,
+            status_code=status_code,
             headers=httpx_headers,
             content=r.text,
         )

tests/unit_tests/test_errors.py

@@ -32,6 +32,20 @@
     single_choice_chunk,
 )
 
+_API_VERSION = "api-version=2023-03-15-preview"
+_UPSTREAM_ENDPOINT = (
+    "http://localhost:5001/openai/deployments/gpt-4/chat/completions"
+)
+
+
+@pytest.fixture(autouse=True)
+def mock_azure_ad_token():
+    with patch(
+        "aidial_adapter_openai.utils.auth.get_api_key",
+        return_value="test-azure-ad-token",
+    ):
+        yield
+
 
 def assert_equal(actual: Any, expected: Any):
     assert actual == expected
@@ -1044,7 +1058,7 @@ async def test_error_from_gpt_multi_modal(stream: bool):
             },
         )
 
-        assert response.status_code == 500
+        assert response.status_code == 502
         assert response.content == b"Something went wrong"
 
 
@@ -1276,3 +1290,74 @@ async def test_rate_limit_exceeded_during_streaming():
                 "type": "server_error",
             }
         }
+
+
+@respx.mock
+@pytest.mark.parametrize("upstream_error_code", [400, 401, 403, 500, 502, 503])
+@pytest.mark.parametrize(
+    "error_during_streaming",
+    [True, False],
+    ids=["error-in-stream", "error-before-stream"],
+)
+@pytest.mark.parametrize(
+    "with_api_key", [True, False], ids=["with-key", "no-key"]
+)
+@pytest.mark.parametrize("stream", [True, False], ids=["stream", "block"])
+async def test_upstream_errors_to_adapter_errors(
+    test_app: httpx.AsyncClient,
+    upstream_error_code: int,
+    error_during_streaming: bool,
+    with_api_key: bool,
+    stream: bool,
+):
+    mock_stream = OpenAIStream(
+        single_choice_chunk(delta={"role": "assistant", "content": "test"}),
+        {
+            "error": {
+                "message": "Upstream streaming error",
+                "code": int(upstream_error_code),
+            }
+        },
+    )
+
+    def chat_completion_handler(request: httpx.Request):
+        streaming = json.loads(request.content).get("stream", False)
+        if streaming and error_during_streaming:
+            return httpx.Response(
+                status_code=200,
+                headers={"Content-Type": "text/event-stream"},
+                content=mock_stream.to_content(),
+            )
+        else:
+            return httpx.Response(
+                status_code=upstream_error_code,
+                content="Upstream error",
+            )
+
+    respx.post(f"{_UPSTREAM_ENDPOINT}?{_API_VERSION}").mock(
+        side_effect=chat_completion_handler
+    )
+
+    headers = {"X-UPSTREAM-ENDPOINT": _UPSTREAM_ENDPOINT}
+    if with_api_key:
+        headers["X-UPSTREAM-KEY"] = "TEST_API_KEY"
+
+    response = await test_app.post(
+        f"/openai/deployments/gpt-4/chat/completions?{_API_VERSION}",
+        json={
+            "messages": [{"role": "user", "content": "test"}],
+            "stream": stream,
+        },
+        headers=headers,
+    )
+
+    if stream and error_during_streaming:
+        assert response.status_code == 200
+        return
+
+    if upstream_error_code == 500 or (
+        upstream_error_code in (401, 403) and with_api_key
+    ):
+        assert response.status_code == 502
+    else:
+        assert response.status_code == upstream_error_code