AAP-18027: WCAClient.get_token: Distinguish between API Key related failures and others (#697)

Co-authored-by: Michael Anstis <manstis@redhat.com>

Author: manstis

ansible_wisdom/ai/api/model_client/exceptions.py

@@ -47,6 +47,11 @@ class WcaTokenFailure(WcaException):
     """An attempt to retrieve a WCA Token failed."""
 
 
+@dataclass
+class WcaTokenFailureApiKeyError(WcaException):
+    """An attempt to retrieve a WCA Token failed due to a problem with the provided API Key."""
+
+
 @dataclass
 class WcaCloudflareRejection(WcaException):
     """Cloudflare rejected the request."""

ansible_wisdom/ai/api/model_client/wca_client.py

@@ -9,6 +9,8 @@
     ContentMatchResponseChecks,
     InferenceContext,
     InferenceResponseChecks,
+    TokenContext,
+    TokenResponseChecks,
 )
 from django.apps import apps
 from django.conf import settings
@@ -205,6 +207,8 @@ def post_request():
 
         try:
             response = post_request()
+            context = TokenContext(response)
+            TokenResponseChecks().run_checks(context)
             response.raise_for_status()
 
         except HTTPError as e:

ansible_wisdom/ai/api/model_client/wca_utils.py

@@ -1,7 +1,12 @@
 from abc import abstractmethod
 from typing import Generic, TypeVar
 
-from .exceptions import WcaCloudflareRejection, WcaEmptyResponse, WcaInvalidModelId
+from .exceptions import (
+    WcaCloudflareRejection,
+    WcaEmptyResponse,
+    WcaInvalidModelId,
+    WcaTokenFailureApiKeyError,
+)
 
 T = TypeVar('T')
 
@@ -23,6 +28,34 @@ def run_checks(self, context: T):
             check.check(context)
 
 
+class TokenContext:
+    def __init__(self, result):
+        self.result = result
+
+
+class TokenResponseChecks(Checks[TokenContext]):
+    class ResponseStatusCode400Missing(Check[TokenContext]):
+        def check(self, context: TokenContext):
+            if context.result.status_code == 400:
+                if "Property missing or empty" in context.result.json()["errorMessage"]:
+                    raise WcaTokenFailureApiKeyError()
+
+    class ResponseStatusCode400NotFound(Check[TokenContext]):
+        def check(self, context: TokenContext):
+            if context.result.status_code == 400:
+                if "Provided API key could not be found" in context.result.json()["errorMessage"]:
+                    raise WcaTokenFailureApiKeyError()
+
+    def __init__(self):
+        super().__init__(
+            [
+                # The ordering of these checks is important!
+                TokenResponseChecks.ResponseStatusCode400Missing(),
+                TokenResponseChecks.ResponseStatusCode400NotFound(),
+            ]
+        )
+
+
 class InferenceContext:
     def __init__(self, model_id, result, is_multi_task_prompt):
         self.model_id = model_id

ansible_wisdom/ai/api/wca/api_key_views.py

@@ -3,7 +3,7 @@
 
 from ai.api.aws.exceptions import WcaSecretManagerError
 from ai.api.aws.wca_secret_manager import Suffixes
-from ai.api.model_client.exceptions import WcaTokenFailure
+from ai.api.model_client.exceptions import WcaTokenFailureApiKeyError
 from ai.api.permissions import (
     AcceptedTermsPermission,
     IsOrganisationAdministrator,
@@ -140,18 +140,18 @@ def post(self, request, *args, **kwargs):
             logger.info(e, exc_info=True)
             return Response(status=HTTP_400_BAD_REQUEST)
 
-        except WcaTokenFailure as e:
+        except WcaTokenFailureApiKeyError as e:
             exception = e
             logger.info(
                 f"An error occurred trying to retrieve a WCA Token for Organisation '{org_id}'.",
                 exc_info=True,
             )
             return Response(status=HTTP_400_BAD_REQUEST)
 
-        except WcaSecretManagerError as e:
+        except Exception as e:
             exception = e
             logger.exception(e)
-            raise ServiceUnavailable
+            raise ServiceUnavailable(cause=e)
 
         finally:
             duration = round((time.time() - start_time) * 1000, 2)
@@ -199,23 +199,25 @@ def get(self, request, *args, **kwargs):
             model_mesh_client = apps.get_app_config("ai").wca_client
             secret_manager = apps.get_app_config("ai").get_wca_secret_manager()
             api_key = secret_manager.get_secret(org_id, Suffixes.API_KEY)
-            token = model_mesh_client.get_token(api_key['SecretString'])
+            if api_key is None:
+                return Response(status=HTTP_400_BAD_REQUEST)
+            token = model_mesh_client.get_token(api_key)
             if token is None:
                 return Response(status=HTTP_400_BAD_REQUEST)
 
-        except WcaSecretManagerError as e:
-            exception = e
-            logger.exception(e)
-            raise ServiceUnavailable
-
-        except WcaTokenFailure as e:
+        except WcaTokenFailureApiKeyError as e:
             exception = e
             logger.info(
                 f"An error occurred trying to retrieve a WCA Token for Organisation '{org_id}'.",
                 exc_info=True,
             )
             return Response(status=HTTP_400_BAD_REQUEST)
 
+        except Exception as e:
+            exception = e
+            logger.exception(e)
+            raise ServiceUnavailable(cause=e)
+
         finally:
             duration = round((time.time() - start_time) * 1000, 2)
             event = {

ansible_wisdom/ai/api/wca/model_id_views.py

@@ -240,11 +240,6 @@ def do_validated_operation(request, api_key_provider, model_id_provider, on_succ
         logger.info(e, exc_info=True)
         raise WcaBadRequestException(cause=e)
 
-    except WcaSecretManagerError as e:
-        exception = e
-        logger.exception(e)
-        return Response(status=HTTP_500_INTERNAL_SERVER_ERROR)
-
     except Exception as e:
         exception = e
         logger.exception(e)

ansible_wisdom/ai/api/wca/tests/test_api_key_views.py

@@ -4,7 +4,7 @@
 
 from ai.api.aws.exceptions import WcaSecretManagerError
 from ai.api.aws.wca_secret_manager import Suffixes, WcaSecretManager
-from ai.api.model_client.exceptions import WcaTokenFailure
+from ai.api.model_client.exceptions import WcaTokenFailure, WcaTokenFailureApiKeyError
 from ai.api.model_client.wca_client import WCAClient
 from ai.api.permissions import (
     AcceptedTermsPermission,
@@ -194,7 +194,9 @@ def test_set_key_with_invalid_value(self, *args):
         self.mock_secret_manager.get_secret.assert_called_with('123', Suffixes.API_KEY)
 
         # Set Key
-        self.mock_wca_client.get_token.side_effect = WcaTokenFailure('Something went wrong')
+        self.mock_wca_client.get_token.side_effect = WcaTokenFailureApiKeyError(
+            'Something went wrong'
+        )
         with self.assertLogs(logger='root', level='DEBUG') as log:
             r = self.client.post(
                 reverse('wca_api_key'),
@@ -203,7 +205,7 @@ def test_set_key_with_invalid_value(self, *args):
             )
             self.assertEqual(r.status_code, HTTPStatus.BAD_REQUEST)
             self.mock_secret_manager.save_secret.assert_not_called()
-            _assert_segment_log(self, log, "modelApiKeySet", "WcaTokenFailure")
+            _assert_segment_log(self, log, "modelApiKeySet", "WcaTokenFailureApiKeyError")
 
     @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
     def test_set_key_throws_secret_manager_exception(self, *args):
@@ -222,18 +224,17 @@ def test_set_key_throws_secret_manager_exception(self, *args):
             _assert_segment_log(self, log, "modelApiKeySet", "WcaSecretManagerError")
 
     @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
-    def test_set_key_throws_wca_client_exception(self, *args):
+    def test_set_key_throws_http_exception(self, *args):
         self.user.organization_id = '123'
         self.client.force_authenticate(user=self.user)
         self.mock_wca_client.get_token.side_effect = WcaTokenFailure()
-
         with self.assertLogs(logger='root', level='DEBUG') as log:
             r = self.client.post(
                 reverse('wca_api_key'),
                 data='{ "key": "a-new-key" }',
                 content_type='application/json',
             )
-            self.assertEqual(r.status_code, HTTPStatus.BAD_REQUEST)
+            self.assertEqual(r.status_code, HTTPStatus.SERVICE_UNAVAILABLE)
             _assert_segment_log(self, log, "modelApiKeySet", "WcaTokenFailure")
 
     @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
@@ -314,13 +315,37 @@ def _test_validate_key_with_valid_value(self, has_seat):
             self.assertEqual(r.status_code, HTTPStatus.OK)
             _assert_segment_log(self, log, "modelApiKeyValidate", None)
 
+    @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
+    def test_validate_key_with_missing_value(self, *args):
+        self.user.organization_id = '123'
+        self.client.force_authenticate(user=self.user)
+        self.mock_secret_manager.get_secret.return_value = None
+
+        with self.assertLogs(logger='root', level='DEBUG') as log:
+            r = self.client.get(reverse('wca_api_key_validator'))
+            self.assertEqual(r.status_code, HTTPStatus.BAD_REQUEST)
+            _assert_segment_log(self, log, "modelApiKeyValidate", None)
+
     @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
     def test_validate_key_with_invalid_value(self, *args):
         self.user.organization_id = '123'
         self.client.force_authenticate(user=self.user)
-        self.mock_wca_client.get_token.side_effect = WcaTokenFailure('Something went wrong')
+        self.mock_wca_client.get_token.side_effect = WcaTokenFailureApiKeyError(
+            'Something went wrong'
+        )
 
         with self.assertLogs(logger='root', level='DEBUG') as log:
             r = self.client.get(reverse('wca_api_key_validator'))
             self.assertEqual(r.status_code, HTTPStatus.BAD_REQUEST)
+            _assert_segment_log(self, log, "modelApiKeyValidate", "WcaTokenFailureApiKeyError")
+
+    @override_settings(SEGMENT_WRITE_KEY='DUMMY_KEY_VALUE')
+    def test_validate_key_throws_http_exception(self, *args):
+        self.user.organization_id = '123'
+        self.client.force_authenticate(user=self.user)
+        self.mock_wca_client.get_token.side_effect = WcaTokenFailure('Something went wrong')
+
+        with self.assertLogs(logger='root', level='DEBUG') as log:
+            r = self.client.get(reverse('wca_api_key_validator'))
+            self.assertEqual(r.status_code, HTTPStatus.SERVICE_UNAVAILABLE)
             _assert_segment_log(self, log, "modelApiKeyValidate", "WcaTokenFailure")

ansible_wisdom/ai/api/wca/tests/test_model_id_views.py

@@ -231,7 +231,7 @@ def test_set_model_id_throws_secret_manager_exception(self, *args):
                 data='{ "model_id": "secret_model_id" }',
                 content_type='application/json',
             )
-            self.assertEqual(r.status_code, HTTPStatus.INTERNAL_SERVER_ERROR)
+            self.assertEqual(r.status_code, HTTPStatus.SERVICE_UNAVAILABLE)
             self.assertInLog('ai.api.aws.exceptions.WcaSecretManagerError', log)
             _assert_segment_log(self, log, "modelIdSet", "WcaSecretManagerError")