From 617c1790889c3bae54ab495545f33f047af2c378 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gon=C3=A9ri=20Le=20Bouder?= Date: Thu, 9 Jan 2025 07:53:45 -0600 Subject: [PATCH] AAP-30619: migrate PlaybookGeneration to the new Schema1 object (#1485) Use an object defined in ansible_ai_connect/ai/api/telemetry/schema1.py to prepare the Schema1 payload. --- .../ai/api/telemetry/schema1.py | 11 + .../ai/api/telemetry/test_schema1.py | 16 +- ansible_ai_connect/ai/api/tests/test_views.py | 22 +- ansible_ai_connect/ai/api/views.py | 223 ++++-------------- 4 files changed, 86 insertions(+), 186 deletions(-) diff --git a/ansible_ai_connect/ai/api/telemetry/schema1.py b/ansible_ai_connect/ai/api/telemetry/schema1.py index 7d80bba57..e1b1673fd 100644 --- a/ansible_ai_connect/ai/api/telemetry/schema1.py +++ b/ansible_ai_connect/ai/api/telemetry/schema1.py @@ -162,6 +162,17 @@ class ExplainPlaybookEvent(Schema1Event): explanationId: str = field(validator=validators.instance_of(str), converter=str, default="") +@define +class GenerationPlaybookEvent(Schema1Event): + event_name: str = "codegenPlaybook" + playbook_length: int = field(validator=validators.instance_of(int), default=0) + generationId: str = field(validator=validators.instance_of(str), converter=str, default="") + wizardId: str = field(validator=validators.instance_of(str), converter=str, default="") + create_outline: bool = field( + validator=validators.instance_of(bool), converter=bool, default=False + ) + + @define class ChatBotResponseDocsReferences: docs_url: str = field(validator=validators.instance_of(str), converter=str, default="") diff --git a/ansible_ai_connect/ai/api/telemetry/test_schema1.py b/ansible_ai_connect/ai/api/telemetry/test_schema1.py index cc436c383..b19a61c67 100644 --- a/ansible_ai_connect/ai/api/telemetry/test_schema1.py +++ b/ansible_ai_connect/ai/api/telemetry/test_schema1.py @@ -21,7 +21,12 @@ from ansible_ai_connect.test_utils import WisdomServiceAPITestCaseBaseOIDC from ansible_ai_connect.users.models import Plan -from .schema1 import ExplainPlaybookEvent, OneClickTrialStartedEvent, Schema1Event +from .schema1 import ( + ExplainPlaybookEvent, + GenerationPlaybookEvent, + OneClickTrialStartedEvent, + Schema1Event, +) @override_settings(AUTHZ_BACKEND_TYPE="dummy") @@ -107,3 +112,12 @@ class TestExplainPlaybookEvent(WisdomServiceAPITestCaseBaseOIDC): def test_base(self): event1 = ExplainPlaybookEvent() self.assertEqual(event1.event_name, "explainPlaybook") + + +@override_settings(AUTHZ_BACKEND_TYPE="dummy") +@override_settings(WCA_SECRET_DUMMY_SECRETS="1981:valid") +class TestGenerationPlaybookEvent(WisdomServiceAPITestCaseBaseOIDC): + def test_base(self): + event1 = GenerationPlaybookEvent() + self.assertEqual(event1.event_name, "codegenPlaybook") + self.assertFalse(event1.create_outline) diff --git a/ansible_ai_connect/ai/api/tests/test_views.py b/ansible_ai_connect/ai/api/tests/test_views.py index 6006ec7f7..7a10d71d6 100644 --- a/ansible_ai_connect/ai/api/tests/test_views.py +++ b/ansible_ai_connect/ai/api/tests/test_views.py @@ -3604,7 +3604,7 @@ def test_bad_wca_request(self): model_client, HTTPStatus.NO_CONTENT, WcaBadRequestException, - "bad request for playbook generation", + "WCA returned a bad request response", ) def test_missing_api_key(self): @@ -3616,7 +3616,7 @@ def test_missing_api_key(self): model_client, HTTPStatus.FORBIDDEN, WcaKeyNotFoundException, - "A WCA Api Key was expected but not found for playbook generation", + "A WCA Api Key was expected but not found", ) def test_missing_model_id(self): @@ -3628,7 +3628,7 @@ def test_missing_model_id(self): model_client, HTTPStatus.FORBIDDEN, WcaModelIdNotFoundException, - "A WCA Model ID was expected but not found for playbook generation", + "A WCA Model ID was expected but not found", ) def test_missing_default_model_id(self): @@ -3640,7 +3640,7 @@ def test_missing_default_model_id(self): model_client, HTTPStatus.FORBIDDEN, WcaNoDefaultModelIdException, - "A default WCA Model ID was expected but not found for playbook generation", + "No default WCA Model ID was found", ) def test_request_id_correlation_failure(self): @@ -3656,7 +3656,7 @@ def test_request_id_correlation_failure(self): model_client, HTTPStatus.INTERNAL_SERVER_ERROR, WcaRequestIdCorrelationFailureException, - "WCA Request/Response GenerationId correlation failed", + "WCA Request/Response Request Id correlation failed", ) def test_invalid_model_id(self): @@ -3669,7 +3669,7 @@ def test_invalid_model_id(self): model_client, HTTPStatus.FORBIDDEN, WcaInvalidModelIdException, - "WCA Model ID is invalid for playbook generation", + "WCA Model ID is invalid", ) def test_empty_response(self): @@ -3680,7 +3680,7 @@ def test_empty_response(self): model_client, HTTPStatus.NO_CONTENT, WcaEmptyResponseException, - "WCA returned an empty response for playbook generation", + "WCA returned an empty response", ) def test_cloudflare_rejection(self): @@ -3689,7 +3689,7 @@ def test_cloudflare_rejection(self): model_client, HTTPStatus.BAD_REQUEST, WcaCloudflareRejectionException, - "Cloudflare rejected the request for playbook generation", + "Cloudflare rejected the request", ) def test_hap_filter(self): @@ -3704,7 +3704,7 @@ def test_hap_filter(self): model_client, HTTPStatus.BAD_REQUEST, WcaHAPFilterRejectionException, - "WCA Hate, Abuse, and Profanity filter rejected the request for playbook generation", + "WCA Hate, Abuse, and Profanity filter rejected the request", ) def test_user_trial_expired(self): @@ -3716,7 +3716,7 @@ def test_user_trial_expired(self): model_client, HTTPStatus.FORBIDDEN, WcaUserTrialExpiredException, - "User trial expired, when requesting playbook generation", + "User trial expired", ) def test_wca_instance_deleted(self): @@ -3733,7 +3733,7 @@ def test_wca_instance_deleted(self): model_client, HTTPStatus.IM_A_TEAPOT, WcaInstanceDeletedException, - "WCA Instance has been deleted when requesting playbook generation", + "The WCA instance associated with the Model ID has been deleted", ) def test_wca_request_with_model_id_given(self): diff --git a/ansible_ai_connect/ai/api/views.py b/ansible_ai_connect/ai/api/views.py index 5d82bf116..006fccaed 100644 --- a/ansible_ai_connect/ai/api/views.py +++ b/ansible_ai_connect/ai/api/views.py @@ -194,7 +194,7 @@ def initialize_request(self, request, *args, **kwargs): # should exposed as self.validated_data.model, or using a special # method. # See: https://github.com/ansible/ansible-ai-connect-service/pull/1147/files#diff-ecfb6919dfd8379aafba96af7457b253e4dce528897dfe6bfc207ca2b3b2ada9R143-R151 # noqa: E501 - self.model_id: str = "" + self.req_model_id: str = "" return initialised_request @@ -232,7 +232,9 @@ def finalize_response(self, request, response, *args, **kwargs): model_meta_data: MetaData = apps.get_app_config("ai").get_model_pipeline(MetaData) user = request.user org_id = hasattr(user, "organization") and user.organization and user.organization.id - self.event.modelName = model_meta_data.get_model_id(request.user, org_id, self.model_id) + self.event.modelName = self.event.modelName or model_meta_data.get_model_id( + request.user, org_id, self.req_model_id + ) except (WcaNoDefaultModelId, WcaModelIdNotFound, WcaSecretManagerError): pass self.event.set_response(response) @@ -771,7 +773,7 @@ def post(self, request) -> Response: explanation_id = str(request_serializer.validated_data.get("explanationId", "")) playbook = request_serializer.validated_data.get("content") custom_prompt = str(request_serializer.validated_data.get("customPrompt", "")) - self.model_id = str(request_serializer.validated_data.get("model", "")) + self.req_model_id = str(request_serializer.validated_data.get("model", "")) llm: ModelPipelinePlaybookExplanation = apps.get_app_config("ai").get_model_pipeline( ModelPipelinePlaybookExplanation @@ -806,13 +808,14 @@ def post(self, request) -> Response: ) -class GenerationPlaybook(APIView): +class GenerationPlaybook(AACSAPIView): """ Returns a playbook based on a text input. """ permission_classes = PERMISSIONS_MAP.get(settings.DEPLOYMENT_MODE) required_scopes = ["read", "write"] + schema1_event = schema1.GenerationPlaybookEvent throttle_cache_key_suffix = "_generation_playbook" @@ -829,195 +832,67 @@ class GenerationPlaybook(APIView): summary="Inline code suggestions", ) def post(self, request) -> Response: - exception = None generation_id = None wizard_id = None - duration = None create_outline = None anonymized_playbook = "" playbook = "" - request_serializer = GenerationPlaybookRequestSerializer(data=request.data) answer = {} model_id = "" - try: - request_serializer.is_valid(raise_exception=True) - generation_id = str(request_serializer.validated_data.get("generationId", "")) - create_outline = request_serializer.validated_data["createOutline"] - outline = str(request_serializer.validated_data.get("outline", "")) - text = request_serializer.validated_data["text"] - custom_prompt = str(request_serializer.validated_data.get("customPrompt", "")) - wizard_id = str(request_serializer.validated_data.get("wizardId", "")) - model_id = str(request_serializer.validated_data.get("model", "")) - - llm: ModelPipelinePlaybookGeneration = apps.get_app_config("ai").get_model_pipeline( - ModelPipelinePlaybookGeneration - ) - start_time = time.time() - playbook, outline, warnings = llm.invoke( - PlaybookGenerationParameters.init( - request=request, - text=text, - custom_prompt=custom_prompt, - create_outline=create_outline, - outline=outline, - generation_id=generation_id, - model_id=model_id, - ) - ) - duration = round((time.time() - start_time) * 1000, 2) - - # Anonymize responses - # Anonymized in the View to be consistent with where Completions are anonymized - anonymized_playbook = anonymizer.anonymize_struct( - playbook, value_template=Template("{{ _${variable_name}_ }}") - ) - anonymized_outline = anonymizer.anonymize_struct( - outline, value_template=Template("{{ _${variable_name}_ }}") - ) - - answer = { - "playbook": anonymized_playbook, - "outline": anonymized_outline, - "warnings": warnings, - "format": "plaintext", - "generationId": generation_id, - } - - except WcaBadRequest as e: - exception = e - logger.exception(f"bad request for playbook generation {generation_id}") - raise WcaBadRequestException(cause=e) - - except WcaInvalidModelId as e: - exception = e - logger.exception(f"WCA Model ID is invalid for playbook generation {generation_id}") - raise WcaInvalidModelIdException(cause=e) - - except WcaKeyNotFound as e: - exception = e - logger.exception( - f"A WCA Api Key was expected but not found for " - f"playbook generation {generation_id}" - ) - raise WcaKeyNotFoundException(cause=e) - - except WcaModelIdNotFound as e: - exception = e - logger.exception( - f"A WCA Model ID was expected but not found for " - f"playbook generation {generation_id}" - ) - raise WcaModelIdNotFoundException(cause=e) - - except WcaNoDefaultModelId as e: - exception = e - logger.exception( - "A default WCA Model ID was expected but not found for " - f"playbook generation {generation_id}" - ) - raise WcaNoDefaultModelIdException(cause=e) - - except WcaRequestIdCorrelationFailure as e: - exception = e - logger.exception( - f"WCA Request/Response GenerationId correlation failed " - f"for suggestion {generation_id}" - ) - raise WcaRequestIdCorrelationFailureException(cause=e) - except WcaEmptyResponse as e: - exception = e - logger.exception( - f"WCA returned an empty response for playbook generation {generation_id}" - ) - raise WcaEmptyResponseException(cause=e) - - except WcaCloudflareRejection as e: - exception = e - logger.exception( - f"Cloudflare rejected the request for playbook generation {generation_id}" - ) - raise WcaCloudflareRejectionException(cause=e) - - except WcaHAPFilterRejection as e: - exception = e - logger.exception( - f"WCA Hate, Abuse, and Profanity filter rejected " - f"the request for playbook generation {generation_id}" - ) - raise WcaHAPFilterRejectionException(cause=e) + request_serializer = GenerationPlaybookRequestSerializer(data=request.data) + request_serializer.is_valid(raise_exception=True) + generation_id = str(request_serializer.validated_data.get("generationId", "")) + create_outline = request_serializer.validated_data["createOutline"] + outline = str(request_serializer.validated_data.get("outline", "")) + text = request_serializer.validated_data["text"] + custom_prompt = str(request_serializer.validated_data.get("customPrompt", "")) + wizard_id = str(request_serializer.validated_data.get("wizardId", "")) + self.req_model_id = str(request_serializer.validated_data.get("model", "")) - except WcaUserTrialExpired as e: - exception = e - logger.exception( - f"User trial expired, when requesting playbook generation {generation_id}" - ) - raise WcaUserTrialExpiredException(cause=e) + self.event.generationId = generation_id + self.event.wizardId = wizard_id + self.event.modelName = model_id - except WcaInstanceDeleted as e: - exception = e - logger.exception( - "WCA Instance has been deleted when requesting playbook generation " - f"{generation_id} for model {e.model_id}" + self.event.create_outline = create_outline + llm: ModelPipelinePlaybookGeneration = apps.get_app_config("ai").get_model_pipeline( + ModelPipelinePlaybookGeneration + ) + playbook, outline, warnings = llm.invoke( + PlaybookGenerationParameters.init( + request=request, + text=text, + custom_prompt=custom_prompt, + create_outline=create_outline, + outline=outline, + generation_id=generation_id, + model_id=model_id, ) - raise WcaInstanceDeletedException(cause=e) + ) - except Exception as exc: - exception = exc - logger.exception(f"An exception {exc.__class__} occurred during a playbook generation") - raise + # Anonymize responses + # Anonymized in the View to be consistent with where Completions are anonymized + anonymized_playbook = anonymizer.anonymize_struct( + playbook, value_template=Template("{{ _${variable_name}_ }}") + ) + anonymized_outline = anonymizer.anonymize_struct( + outline, value_template=Template("{{ _${variable_name}_ }}") + ) + self.event.playbook_length = len(anonymized_playbook) - finally: - self.write_to_segment( - request.user, - generation_id, - wizard_id, - exception, - duration, - create_outline, - playbook_length=len(anonymized_playbook), - model_id=model_id, - ) + answer = { + "playbook": anonymized_playbook, + "outline": anonymized_outline, + "warnings": warnings, + "format": "plaintext", + "generationId": generation_id, + } return Response( answer, status=rest_framework_status.HTTP_200_OK, ) - def write_to_segment( - self, - user, - generation_id, - wizard_id, - exception, - duration, - create_outline, - playbook_length, - model_id, - ): - model_name = "" - try: - model_meta_data: MetaData = apps.get_app_config("ai").get_model_pipeline(MetaData) - model_name = model_meta_data.get_model_id(user, user.org_id, model_id) - except (WcaNoDefaultModelId, WcaModelIdNotFound, WcaSecretManagerError): - pass - event = { - "create_outline": create_outline, - "duration": duration, - "exception": exception is not None, - "generationId": generation_id, - "modelName": model_name, - "playbook_length": playbook_length, - "wizardId": wizard_id, - } - if exception: - event["response"] = ( - { - "exception": str(exception), - }, - ) - send_segment_event(event, "codegenPlaybook", user) - class GenerationRole(APIView): """