Lint the results from role generation (#1513)

Author: mabashian

ansible_ai_connect/ai/api/model_pipelines/dummy/pipelines.py

@@ -181,7 +181,7 @@ def __init__(self, config: DummyConfiguration):
 
     def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse:
         create_outline = params.create_outline
-        return "install_nginx", ROLE_FILES, ROLE_OUTLINE.strip() if create_outline else ""
+        return "install_nginx", ROLE_FILES, ROLE_OUTLINE.strip() if create_outline else "", []
 
     def self_test(self) -> Optional[HealthCheckSummary]:
         raise NotImplementedError

ansible_ai_connect/ai/api/model_pipelines/langchain/pipelines.py

@@ -374,7 +374,7 @@ def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse:
         else:
             files[1]["content"] = unwrap_message_with_yaml_answer(content)
 
-        return role, files, outline
+        return role, files, outline, []
 
     def self_test(self) -> Optional[HealthCheckSummary]:
         raise NotImplementedError

ansible_ai_connect/ai/api/model_pipelines/langchain/tests/test_pipeline.py

@@ -316,7 +316,7 @@ def fake_get_chat_mode(model_id=None):
         self.my_client.get_chat_model = fake_get_chat_mode
 
     def test_generate_role(self):
-        role, files, outline = self.my_client.invoke(
+        role, files, outline, warnings = self.my_client.invoke(
             RoleGenerationParameters.init(
                 request=Mock(),
                 text="foo",
@@ -326,9 +326,10 @@ def test_generate_role(self):
         self.assertEqual(role, "vpc_subnet_ec2")
         self.assertEqual(files[1]["content"], TestUnwrapRoleAnswer._expected_second_request_content)
         self.assertEqual(outline, "")
+        self.assertEqual(warnings, [])
 
     def test_generate_role_with_outline(self):
-        role, files, outline = self.my_client.invoke(
+        role, files, outline, warnings = self.my_client.invoke(
             RoleGenerationParameters.init(
                 request=Mock(),
                 text="foo",
@@ -343,6 +344,7 @@ def test_generate_role_with_outline(self):
 2. Display EC2 Instance ID
 """,
         )
+        self.assertEqual(warnings, [])
 
 
 class TestLangChainPlaybookExplanationPipeline(WisdomServiceLogAwareTestCase):

ansible_ai_connect/ai/api/model_pipelines/pipelines.py

@@ -152,7 +152,7 @@ def init(
         )
 
 
-RoleGenerationResponse = tuple[str, list, str]
+RoleGenerationResponse = tuple[str, list, str, list]
 
 
 @define

ansible_ai_connect/ai/api/model_pipelines/tests/test_wca_client.py

@@ -53,6 +53,7 @@
     ContentMatchParameters,
     PlaybookExplanationParameters,
     PlaybookGenerationParameters,
+    RoleGenerationParameters,
 )
 from ansible_ai_connect.ai.api.model_pipelines.tests import mock_pipeline_config
 from ansible_ai_connect.ai.api.model_pipelines.wca.pipelines_base import (
@@ -63,6 +64,7 @@
     wca_codegen_playbook_hist,
     wca_codegen_playbook_retry_counter,
     wca_codegen_retry_counter,
+    wca_codegen_role_hist,
     wca_codematch_hist,
     wca_codematch_retry_counter,
     wca_explain_playbook_hist,
@@ -77,6 +79,7 @@
     WCASaaSContentMatchPipeline,
     WCASaaSPlaybookExplanationPipeline,
     WCASaaSPlaybookGenerationPipeline,
+    WCASaaSRoleGenerationPipeline,
 )
 from ansible_ai_connect.test_utils import (
     WisdomAppsBackendMocking,
@@ -335,7 +338,7 @@ def test_get_api_key_with_wca_configured(self):
 
 @override_settings(WCA_SECRET_BACKEND_TYPE="dummy")
 @override_settings(ENABLE_ANSIBLE_LINT_POSTPROCESS=False)
-class TestWCAClientGeneration(WisdomAppsBackendMocking, WisdomServiceLogAwareTestCase):
+class TestWCAClientPlaybookGeneration(WisdomAppsBackendMocking, WisdomServiceLogAwareTestCase):
     def setUp(self):
         super().setUp()
         wca_client = WCASaaSPlaybookGenerationPipeline(
@@ -510,6 +513,64 @@ def test_playbook_gen_request_id_correlation_failure(self):
             )
 
 
+@override_settings(WCA_SECRET_BACKEND_TYPE="dummy")
+@override_settings(ENABLE_ANSIBLE_LINT_POSTPROCESS=False)
+class TestWCAClientRoleGeneration(WisdomAppsBackendMocking, WisdomServiceLogAwareTestCase):
+    def setUp(self):
+        super().setUp()
+        wca_client = WCASaaSRoleGenerationPipeline(
+            mock_pipeline_config("wca", api_key=None, model_id=None)
+        )
+        wca_client.get_api_key = Mock(return_value="some-key")
+        wca_client.get_token = Mock(return_value={"access_token": "a-token"})
+        wca_client.get_model_id = Mock(return_value="a-random-model")
+        wca_client.session = Mock()
+        response = Mock
+        response.text = (
+            '{"name": "foo_bar", "outline": "Ahh!", "files": [{"path": '
+            '"roles/foo_bar/tasks/main.yml", "content": "some content", '
+            '"file_type": "task"}, {"path": "roles/foo_bar/defaults/main.yml", '
+            '"content": "some content", "file_type": "default"}], "warnings": []}'
+        )
+        response.status_code = 200
+        response.headers = {WCA_REQUEST_ID_HEADER: WCA_REQUEST_ID_HEADER}
+        response.raise_for_status = Mock()
+        wca_client.session.post.return_value = response
+        self.wca_client = wca_client
+
+    @assert_call_count_metrics(metric=wca_codegen_role_hist)
+    @override_settings(ENABLE_ANSIBLE_LINT_POSTPROCESS=True)
+    def test_role_gen_with_lint(self):
+        fake_linter = Mock()
+        fake_linter.run_linter.return_value = "I'm super fake!"
+        self.mock_ansible_lint_caller_with(fake_linter)
+        name, files, outline, warnings = self.wca_client.invoke(
+            RoleGenerationParameters.init(
+                request=Mock(), text="Install Wordpress", create_outline=True
+            )
+        )
+        self.assertEqual(name, "foo_bar")
+        self.assertEqual(outline, "Ahh!")
+        self.assertEqual(warnings, [])
+        for file in files:
+            self.assertEqual(file["content"], "I'm super fake!")
+
+    @assert_call_count_metrics(metric=wca_codegen_role_hist)
+    @override_settings(ENABLE_ANSIBLE_LINT_POSTPROCESS=True)
+    def test_role_gen_when_is_not_initialized(self):
+        self.mock_ansible_lint_caller_with(None)
+        name, files, outline, warnings = self.wca_client.invoke(
+            RoleGenerationParameters.init(
+                request=Mock(), text="Install Wordpress", create_outline=True
+            )
+        )
+        self.assertEqual(name, "foo_bar")
+        self.assertEqual(outline, "Ahh!")
+        self.assertEqual(warnings, [])
+        for file in files:
+            self.assertEqual(file["content"], "some content")
+
+
 @override_settings(WCA_SECRET_BACKEND_TYPE="dummy")
 @override_settings(ENABLE_ANSIBLE_LINT_POSTPROCESS=False)
 class TestWCAClientExplanation(WisdomAppsBackendMocking, WisdomServiceLogAwareTestCase):

ansible_ai_connect/ai/api/model_pipelines/wca/pipelines_base.py

@@ -120,6 +120,12 @@
     namespace=NAMESPACE,
     buckets=DEFAULT_LATENCY_BUCKETS,
 )
+wca_codegen_role_hist = Histogram(
+    "wca_codegen_role_latency_seconds",
+    "Histogram of WCA codegen-role API processing time",
+    namespace=NAMESPACE,
+    buckets=DEFAULT_LATENCY_BUCKETS,
+)
 wca_explain_playbook_hist = Histogram(
     "wca_explain_playbook_latency_seconds",
     "Histogram of WCA explain-playbook API processing time",
@@ -147,6 +153,11 @@
     "Counter of WCA codegen-playbook API invocation retries",
     namespace=NAMESPACE,
 )
+wca_codegen_role_retry_counter = Counter(
+    "wca_codegen_role_retries",
+    "Counter of WCA codegen-role API invocation retries",
+    namespace=NAMESPACE,
+)
 wca_explain_playbook_retry_counter = Counter(
     "wca_explain_playbook_retries",
     "Counter of WCA explain-playbook API invocation retries",
@@ -241,6 +252,11 @@ def on_backoff_codegen_playbook(details):
         WCABasePipeline.log_backoff_exception(details)
         wca_codegen_playbook_retry_counter.inc()
 
+    @staticmethod
+    def on_backoff_codegen_role(details):
+        WCABasePipeline.log_backoff_exception(details)
+        wca_codegen_role_retry_counter.inc()
+
     @staticmethod
     def on_backoff_explain_playbook(details):
         WCABasePipeline.log_backoff_exception(details)

ansible_ai_connect/ai/api/model_pipelines/wca/pipelines_dummy.py

@@ -98,7 +98,7 @@ def __init__(self, config: WCADummyConfiguration):
         super().__init__(config=config)
 
     def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse:
-        return "wca_dummy_role", [], "wca_dummy_outline"
+        return "wca_dummy_role", [], "wca_dummy_outline", []
 
     def self_test(self) -> Optional[HealthCheckSummary]:
         raise NotImplementedError

ansible_ai_connect/ai/api/model_pipelines/wca/pipelines_saas.py

@@ -12,9 +12,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import json
 import logging
 from abc import ABCMeta
-from typing import TYPE_CHECKING, Generic, Optional
+from typing import TYPE_CHECKING, Generic, Optional, cast
 
 import backoff
 from django.apps import apps
@@ -32,6 +33,7 @@
     WcaKeyNotFound,
     WcaModelIdNotFound,
     WcaNoDefaultModelId,
+    WcaRequestIdCorrelationFailure,
     WcaTokenFailure,
 )
 from ansible_ai_connect.ai.api.model_pipelines.pipelines import (
@@ -65,8 +67,11 @@
     WcaModelRequestException,
     WcaTokenRequestException,
     ibm_cloud_identity_token_hist,
+    wca_codegen_role_hist,
 )
 from ansible_ai_connect.ai.api.model_pipelines.wca.wca_utils import (
+    Context,
+    InferenceResponseChecks,
     TokenContext,
     TokenResponseChecks,
 )
@@ -336,6 +341,73 @@ class WCASaaSRoleGenerationPipeline(
     def __init__(self, config: WCASaaSConfiguration):
         super().__init__(config=config)
 
+    # This should be moved to the base WCA class when it becomes available on-prem
+    def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse:
+        request = params.request
+        text = params.text
+        create_outline = params.create_outline
+        outline = params.outline
+        model_id = params.model_id
+        generation_id = params.generation_id
+
+        organization_id = request.user.organization.id if request.user.organization else None
+        api_key = self.get_api_key(request.user, organization_id)
+        model_id = self.get_model_id(request.user, organization_id, model_id)
+
+        headers = self.get_request_headers(api_key, generation_id)
+        data = {
+            "model_id": model_id,
+            "text": text,
+            "create_outline": create_outline,
+        }
+        if outline:
+            data["outline"] = outline
+
+        @backoff.on_exception(
+            backoff.expo,
+            Exception,
+            max_tries=self.retries + 1,
+            giveup=self.fatal_exception,
+            on_backoff=self.on_backoff_codegen_role,
+        )
+        @wca_codegen_role_hist.time()
+        def post_request():
+            return self.session.post(
+                f"{self.config.inference_url}/v1/wca/codegen/ansible/roles",
+                headers=headers,
+                json=data,
+                verify=self.config.verify_ssl,
+            )
+
+        result = post_request()
+
+        x_request_id = result.headers.get(WCA_REQUEST_ID_HEADER)
+        if generation_id and x_request_id:
+            # request/payload suggestion_id is a UUID not a string whereas
+            # HTTP headers are strings.
+            if x_request_id != str(generation_id):
+                raise WcaRequestIdCorrelationFailure(model_id=model_id, x_request_id=x_request_id)
+
+        context = Context(model_id, result, False)
+        InferenceResponseChecks().run_checks(context)
+        result.raise_for_status()
+
+        response = json.loads(result.text)
+
+        name = response["name"]
+        files = response["files"]
+        outline = response["outline"]
+        warnings = response["warnings"] if "warnings" in response else []
+
+        from ansible_ai_connect.ai.apps import AiConfig
+
+        ai_config = cast(AiConfig, apps.get_app_config("ai"))
+        if ansible_lint_caller := ai_config.get_ansible_lint_caller():
+            for file in files:
+                file["content"] = ansible_lint_caller.run_linter(file["content"])
+
+        return name, files, outline, warnings
+
     def self_test(self) -> Optional[HealthCheckSummary]:
         raise NotImplementedError
 

ansible_ai_connect/ai/api/tests/test_role_generation_view.py

@@ -47,14 +47,26 @@ def __init__(self):
 
 class MockedPipelineRoleGeneration(ModelPipelineRoleGeneration[MockedConfig]):
 
-    def __init__(self, response_roles: str, response_files: list, response_outline: str):
+    def __init__(
+        self,
+        response_roles: str,
+        response_files: list,
+        response_outline: str,
+        response_warnings: list,
+    ):
         super().__init__(MockedConfig())
         self.response_roles = response_roles
         self.response_files = response_files
         self.response_outline = response_outline
+        self.response_warnings = response_warnings
 
     def invoke(self, params: RoleGenerationParameters) -> RoleGenerationResponse:
-        return self.response_roles, self.response_files, self.response_outline
+        return (
+            self.response_roles,
+            self.response_files,
+            self.response_outline,
+            self.response_warnings,
+        )
 
     def self_test(self) -> Optional[HealthCheckSummary]:
         raise NotImplementedError
@@ -127,6 +139,7 @@ def test_with_anonymized_response(self):
                         }
                     ],
                     "Install mysql and email [email protected]",
+                    [],
                 )
             ),
         ):

ansible_ai_connect/ai/api/views.py

@@ -933,7 +933,7 @@ def post(self, request) -> Response:
             ModelPipelineRoleGeneration
         )
 
-        roles, files, outline = llm.invoke(
+        roles, files, outline, warnings = llm.invoke(
             RoleGenerationParameters.init(
                 request=request,
                 text=self.validated_data["text"],
@@ -964,6 +964,7 @@ def post(self, request) -> Response:
             "files": anonymized_files,
             "format": "plaintext",
             "generationId": self.validated_data["generationId"],
+            "warnings": warnings,
         }
 
         return Response(