AAP-36665: Telemetry events for Roles

Also, move the view to `AACSAPIView`.
ansible · Jan 13, 2025 · 6047ba6 · 6047ba6
1 parent 845ffef
commit 6047ba6
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 88 deletions.
diff --git a/ansible_ai_connect/ai/api/serializers.py b/ansible_ai_connect/ai/api/serializers.py
@@ -464,17 +464,6 @@ class ExplanationResponseSerializer(serializers.Serializer):
 
 
 class GenerationPlaybookRequestSerializer(serializers.Serializer):
-    class Meta:
-        fields = [
-            "text",
-            "generationId",
-            "wizardId",
-            "createOutline",
-            "customPrompt",
-            "ansibleExtensionVersion",
-            "outline",
-        ]
-
     text = AnonymizedCharField(
         required=True,
         label="Description content",
@@ -539,18 +528,6 @@ def validate(self, data):
 
 
 class GenerationRoleRequestSerializer(serializers.Serializer):
-    class Meta:
-        fields = [
-            "text",
-            "outline",
-            "createOutline",
-            "additionalContext",
-            "fileTypes",
-            "generationId",
-            "wizardId",
-            "ansibleExtensionVersion",
-        ]
-
     text = AnonymizedCharField(
         required=True,
         label="the goal of the role",
@@ -563,6 +540,7 @@ class Meta:
         required=False,
         label="an outline of the role",
         help_text="An outline of the role should be a numbered list.",
+        default="",
     )
     createOutline = serializers.BooleanField(
         required=False,
@@ -588,6 +566,7 @@ class Meta:
             "The file type name is based on the inner role directories, "
             "without the trailing 's'"
         ),
+        default=['task', 'default'],
     )
     generationId = serializers.UUIDField(
         format="hex_verbose",

diff --git a/ansible_ai_connect/ai/api/telemetry/schema1.py b/ansible_ai_connect/ai/api/telemetry/schema1.py
@@ -173,6 +173,18 @@ class GenerationPlaybookEvent(Schema1Event):
     )
 
 
+@define
+class GenerationRoleEvent(Schema1Event):
+    event_name: str = "codegenRole"
+    tasks_file_length: int = field(validator=validators.instance_of(int), default=0)
+    defaults_file_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/views.py b/ansible_ai_connect/ai/api/views.py
@@ -870,7 +870,7 @@ def post(self, request) -> Response:
         )
 
 
-class GenerationRole(APIView):
+class GenerationRole(AACSAPIView):
     """
     Returns a role based on a text input.
     """
@@ -883,7 +883,8 @@ class GenerationRole(APIView):
         IsAuthenticatedOrTokenHasScope,
     ]
     required_scopes = ["read", "write"]
-
+    schema1_event = schema1.GenerationRoleEvent
+    request_serializer_class = GenerationRoleRequestSerializer
     throttle_cache_key_suffix = "_generation_role"
 
     @extend_schema(
@@ -895,73 +896,44 @@ class GenerationRole(APIView):
         summary="Inline code suggestions",
     )
     def post(self, request) -> Response:
-        # Declaring but commenting out variables to define them but also comply pre-commit
-        text = ""
-        outline = ""
-        create_outline = False
-        additional_context = {}
-        file_types = ["task", "default"]
-        generation_id = ""
-        wizard_id = ""
-        request_serializer = GenerationRoleRequestSerializer(data=request.data)
-        answer = {}
-        model_id = ""
-        try:
-            request_serializer.is_valid(raise_exception=True)
-
-            text = request_serializer.validated_data["text"]
-            outline = str(request_serializer.validated_data.get("outline", outline))
-            create_outline = request_serializer.validated_data["createOutline"]
-            additional_context = request_serializer.validated_data.get(
-                "additionalContext", additional_context
-            )
-            file_types = request_serializer.validated_data.get("fileTypes", file_types)
-            generation_id = str(
-                request_serializer.validated_data.get("generationId", generation_id)
-            )
-            wizard_id = str(request_serializer.validated_data.get("wizardId", wizard_id))
-            model_id = str(request_serializer.validated_data.get("model", ""))
-
-            llm: ModelPipelineRoleGeneration = apps.get_app_config("ai").get_model_pipeline(
-                ModelPipelineRoleGeneration
-            )
+        self.event.create_outline = self.validated_data["createOutline"]
+        self.event.generationId = self.validated_data["generationId"]
+        self.event.wizardId = self.validated_data["wizardId"]
+        llm: ModelPipelineRoleGeneration = apps.get_app_config("ai").get_model_pipeline(
+            ModelPipelineRoleGeneration
+        )
 
-            roles, files, outline = llm.invoke(
-                RoleGenerationParameters.init(
-                    request=request,
-                    text=text,
-                    outline=outline,
-                    model_id=model_id,
-                    create_outline=create_outline,
-                    additional_context=additional_context,
-                    file_types=file_types,
-                    generation_id=generation_id,
-                )
+        roles, files, outline = llm.invoke(
+            RoleGenerationParameters.init(
+                request=request,
+                text=self.validated_data["text"],
+                outline=self.validated_data["outline"],
+                model_id=self.req_model_id,
+                create_outline=self.validated_data["createOutline"],
+                additional_context=self.validated_data.get(
+                    "additionalContext", {}
+                ),
+                file_types=self.validated_data["fileTypes"],
+                generation_id=self.validated_data.get("generationId", "")
             )
+        )
 
-            # Anonymize responses
-            # Anonymized in the View to be consistent with where Completions are anonymized
-            anonymized_role = anonymizer.anonymize_struct(
-                roles, value_template=Template("{{ _${variable_name}_ }}")
-            )
-            anonymized_outline = anonymizer.anonymize_struct(
-                outline, value_template=Template("{{ _${variable_name}_ }}")
-            )
+        # Anonymize responses
+        # Anonymized in the View to be consistent with where Completions are anonymized
+        anonymized_role = anonymizer.anonymize_struct(
+            roles, value_template=Template("{{ _${variable_name}_ }}")
+        )
+        anonymized_outline = anonymizer.anonymize_struct(
+            outline, value_template=Template("{{ _${variable_name}_ }}")
+        )
 
-            answer = {
-                "role": anonymized_role,
-                "outline": anonymized_outline,
-                "files": files,
-                "format": "plaintext",
-                "generationId": generation_id,
-            }
-            ##################################################
-        except Exception as e:
-            logger.exception(f"An exception {e.__class__} occurred during a role generation")
-            raise
-        finally:
-            # implement write to segment there.
-            pass
+        answer = {
+            "role": anonymized_role,
+            "outline": anonymized_outline,
+            "files": files,
+            "format": "plaintext",
+            "generationId": generation_id,
+        }
 
         return Response(
             answer,