From cde1f25263221bd751dd28b7dac870f703474892 Mon Sep 17 00:00:00 2001
From: Manu Chandrasekhar <manuchn@amazon.com>
Date: Sat, 7 Sep 2024 01:39:44 -0400
Subject: [PATCH] v0.0.2

* fix: cleanup model respose, exception logic for ami

* fix: changed info to debug on few log statements
---
 VERSION                                       |  2 +-
 lambda/runtask_fulfillment/ai.py              | 24 ++++++++++++++-----
 .../tools/get_ami_releases.py                 |  8 +------
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/VERSION b/VERSION
index 95e94cd..90ab6e9 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-v0.0.1
\ No newline at end of file
+v0.0.2
\ No newline at end of file
diff --git a/lambda/runtask_fulfillment/ai.py b/lambda/runtask_fulfillment/ai.py
index 0a8fd4f..4ac72da 100644
--- a/lambda/runtask_fulfillment/ai.py
+++ b/lambda/runtask_fulfillment/ai.py
@@ -1,5 +1,6 @@
 import json
 import os
+import re
 
 import boto3
 import botocore
@@ -35,7 +36,7 @@ def eval(tf_plan_json):
     List the resources that will be created, modified or deleted in the following terraform plan using the following rules:
     1. Think step by step using the "thinking" json field
     2. For AMI changes, include the old and new AMI ID
-    3. Use the following schema
+    3. Use the following schema. Skip the preamble:
     <schema>
     {
         "$id": "https://example.com/arrays.schema.json",
@@ -85,9 +86,11 @@ def eval(tf_plan_json):
         bedrock_client, model_id, messages, system_text
     )
 
-    analysis_response_text = json.loads(analysis_response["content"][0]["text"])[
-        "resources"
-    ]
+    logger.debug("Analysis response: {}".format(analysis_response))
+
+    analysis_response_text= clean_response(analysis_response["content"][0]["text"])["resources"]
+
+    logger.debug("Analysis response Text: {}".format(analysis_response_text))
 
     #####################################################################
     ######## Secondly, evaluate AMIs per analysis                ########
@@ -150,9 +153,11 @@ def eval(tf_plan_json):
                     release_details = GetECSAmisReleases().execute(
                         tool["input"]["image_ids"]
                     )
+                    release_details_info = release_details if release_details else "No release notes were found the ami."
+
                     tool_result = {
                         "toolUseId": tool["toolUseId"],
-                        "content": [{"json": {"release_detail": release_details}}],
+                        "content": [{"json": {"release_detail": release_details_info}}],
                     }
 
                     tool_result_message = {
@@ -257,4 +262,11 @@ def guardrail_inspection(input_text, input_mode = 'OUTPUT'):
             return True, "No Guardrail action required"
 
     else:
-        return True, "Guardrail inspection skipped"
\ No newline at end of file
+        return True, "Guardrail inspection skipped"
+    
+def clean_response(json_str):
+    # Remove any tags in the format <tag> or </tag>
+    cleaned_str = re.sub(r'<\/?[\w\s]+>', '', json_str)
+    last_brace_index = cleaned_str.rfind('}')
+    cleaned_str = cleaned_str[:last_brace_index + 1]
+    return json.loads(cleaned_str)
diff --git a/lambda/runtask_fulfillment/tools/get_ami_releases.py b/lambda/runtask_fulfillment/tools/get_ami_releases.py
index 6ce0134..a60d34d 100644
--- a/lambda/runtask_fulfillment/tools/get_ami_releases.py
+++ b/lambda/runtask_fulfillment/tools/get_ami_releases.py
@@ -24,7 +24,6 @@ def get_ecs_amis_releases_info(self, ami_ids):
             ami_data.append({"name": self.get_ami_name_from_id(ami_id), "id": ami_id})
 
         for ami in ami_data:
-            found_release = False
             for release in response_json:
                 details = markdown_to_json.dictify(release["body"])
                 for os_name in details.keys():
@@ -34,7 +33,6 @@ def get_ecs_amis_releases_info(self, ami_ids):
                                 logger.info(
                                     f"Found release notes for {ami['id']}: {ami['name']}"
                                 )
-                                found_release = True
                                 releases_info.append(
                                     {
                                         "ami_id": ami["id"],
@@ -50,7 +48,6 @@ def get_ecs_amis_releases_info(self, ami_ids):
                                     logger.info(
                                         f"Found release notes for {ami['id']}: {ami['name']}"
                                     )
-                                    found_release = True
                                     releases_info.append(
                                         {
                                             "ami_id": ami["id"],
@@ -63,10 +60,7 @@ def get_ecs_amis_releases_info(self, ami_ids):
                                         }
                                     )
                                     break
-            if not found_release:
-                raise Exception(
-                    f"No release notes were found for {ami['id']}: {ami['name']}"
-                )
+
         return releases_info
 
     def get_ami_name_from_id(self, ami_id):