From 7bdb545d52fa014c86bd9410c8a5a03d4a86aa64 Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe.lombrozo@zapier.com>
Date: Wed, 24 Apr 2024 23:07:48 -0400
Subject: [PATCH] clean up more AI output (#203)

---
 pkg/argo_client/applications.go    |  7 +++++++
 pkg/checks/diff/ai_summary.go      | 20 ++++++++++++++++++--
 pkg/checks/diff/ai_summary_test.go |  8 ++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/pkg/argo_client/applications.go b/pkg/argo_client/applications.go
index e047c678..e2b614fb 100644
--- a/pkg/argo_client/applications.go
+++ b/pkg/argo_client/applications.go
@@ -19,6 +19,8 @@ import (
 
 var tracer = otel.Tracer("pkg/argo_client")
 
+var ErrNoVersionFound = errors.New("no kubernetes version found")
+
 // GetApplicationByName takes a context and a name, then queries the Argo Application client to retrieve the Application with the specified name.
 // It returns the found Application and any error encountered during the process.
 // If successful, the Application client connection is closed before returning.
@@ -72,6 +74,11 @@ func (argo *ArgoClient) GetKubernetesVersionByApplication(ctx context.Context, a
 	// cleanup trailing "+"
 	version = strings.TrimSuffix(version, "+")
 
+	version = strings.TrimSpace(version)
+	if version == "" {
+		return "", ErrNoVersionFound
+	}
+
 	return version, nil
 }
 
diff --git a/pkg/checks/diff/ai_summary.go b/pkg/checks/diff/ai_summary.go
index d8774d76..31cf1f8d 100644
--- a/pkg/checks/diff/ai_summary.go
+++ b/pkg/checks/diff/ai_summary.go
@@ -1,6 +1,7 @@
 package diff
 
 import (
+	"regexp"
 	"strings"
 
 	"github.com/rs/zerolog/log"
@@ -44,13 +45,28 @@ func aiDiffSummary(ctx context.Context, mrNote *msg.Message, cfg config.ServerCo
 	mrNote.AddToAppMessage(ctx, name, cr)
 }
 
+var codeBlockBegin = regexp.MustCompile("^```(\\w+)")
+
 func cleanUpAiSummary(aiSummary string) string {
 	aiSummary = strings.TrimSpace(aiSummary)
 
 	// occasionally the model thinks it should wrap it in a code block.
 	// comments do not need this, as they are already rendered as markdown.
-	aiSummary = strings.TrimPrefix(aiSummary, "```markdown")
-	aiSummary = strings.TrimSuffix(aiSummary, "```")
+	for {
+		newSummary := aiSummary
+
+		newSummary = codeBlockBegin.ReplaceAllString(newSummary, "")
+		newSummary = strings.TrimPrefix(newSummary, "#***")
+		newSummary = strings.TrimSuffix(newSummary, "```")
+		newSummary = strings.TrimSuffix(newSummary, "#***")
+		newSummary = strings.TrimSpace(newSummary)
+
+		if newSummary == aiSummary {
+			break
+		}
+
+		aiSummary = newSummary
+	}
 
 	return strings.TrimSpace(aiSummary)
 }
diff --git a/pkg/checks/diff/ai_summary_test.go b/pkg/checks/diff/ai_summary_test.go
index 55d336c9..7a048a6b 100644
--- a/pkg/checks/diff/ai_summary_test.go
+++ b/pkg/checks/diff/ai_summary_test.go
@@ -30,4 +30,12 @@ func TestCleanUpAiSummary(t *testing.T) {
 		actual := cleanUpAiSummary(input)
 		assert.Equal(t, expected, actual)
 	})
+
+	t.Run("weird prefix and suffix", func(t *testing.T) {
+		input := "```plaintext\n#***\n- Added environment variables FF_TIMESTAMPS and FF_SCRIPT_SECTIONS\n#***"
+		expected := "- Added environment variables FF_TIMESTAMPS and FF_SCRIPT_SECTIONS"
+
+		actual := cleanUpAiSummary(input)
+		assert.Equal(t, expected, actual)
+	})
 }