From a556dcebaf4e520abbe396df2708129ea6196a18 Mon Sep 17 00:00:00 2001
From: Joseph Lombrozo <joe.lombrozo@zapier.com>
Date: Fri, 12 Apr 2024 21:38:06 -0400
Subject: [PATCH] trim ai summary when it puts it in a code block (#195)

---
 pkg/aisummary/diff_summary.go      |  2 +-
 pkg/checks/diff/ai_summary.go      | 27 +++++++++++++++++++-----
 pkg/checks/diff/ai_summary_test.go | 33 ++++++++++++++++++++++++++++++
 pkg/checks/diff/diff.go            |  2 +-
 4 files changed, 57 insertions(+), 7 deletions(-)
 create mode 100644 pkg/checks/diff/ai_summary_test.go

diff --git a/pkg/aisummary/diff_summary.go b/pkg/aisummary/diff_summary.go
index 996927fc..bf3387aa 100644
--- a/pkg/aisummary/diff_summary.go
+++ b/pkg/aisummary/diff_summary.go
@@ -13,7 +13,7 @@ import (
 var tracer = otel.Tracer("pkg/aisummary")
 
 // SummarizeDiff uses ChatGPT to summarize changes to a Kubernetes application.
-func (c *OpenAiClient) SummarizeDiff(ctx context.Context, appName string, manifests []string, diff string) (string, error) {
+func (c *OpenAiClient) SummarizeDiff(ctx context.Context, appName, diff string) (string, error) {
 	ctx, span := tracer.Start(ctx, "SummarizeDiff")
 	defer span.End()
 
diff --git a/pkg/checks/diff/ai_summary.go b/pkg/checks/diff/ai_summary.go
index 53969117..d8774d76 100644
--- a/pkg/checks/diff/ai_summary.go
+++ b/pkg/checks/diff/ai_summary.go
@@ -1,6 +1,8 @@
 package diff
 
 import (
+	"strings"
+
 	"github.com/rs/zerolog/log"
 	"go.opentelemetry.io/otel"
 	"golang.org/x/net/context"
@@ -14,7 +16,7 @@ import (
 
 var tracer = otel.Tracer("pkg/checks/diff")
 
-func aiDiffSummary(ctx context.Context, mrNote *msg.Message, cfg config.ServerConfig, name string, manifests []string, diff string) {
+func aiDiffSummary(ctx context.Context, mrNote *msg.Message, cfg config.ServerConfig, name, diff string) {
 	ctx, span := tracer.Start(ctx, "aiDiffSummary")
 	defer span.End()
 
@@ -23,7 +25,8 @@ func aiDiffSummary(ctx context.Context, mrNote *msg.Message, cfg config.ServerCo
 		return
 	}
 
-	aiSummary, err := aisummary.GetOpenAiClient(cfg.OpenAIAPIToken).SummarizeDiff(ctx, name, manifests, diff)
+	aiClient := aisummary.GetOpenAiClient(cfg.OpenAIAPIToken)
+	aiSummary, err := aiClient.SummarizeDiff(ctx, name, diff)
 	if err != nil {
 		telemetry.SetError(span, err, "OpenAI SummarizeDiff")
 		log.Error().Err(err).Msg("failed to summarize diff")
@@ -32,8 +35,22 @@ func aiDiffSummary(ctx context.Context, mrNote *msg.Message, cfg config.ServerCo
 		return
 	}
 
-	if aiSummary != "" {
-		cr := msg.Result{State: pkg.StateNone, Summary: "<b>Show AI Summary Diff</b>", Details: aiSummary}
-		mrNote.AddToAppMessage(ctx, name, cr)
+	aiSummary = cleanUpAiSummary(aiSummary)
+	if aiSummary == "" {
+		return
 	}
+
+	cr := msg.Result{State: pkg.StateNone, Summary: "<b>Show AI Summary Diff</b>", Details: aiSummary}
+	mrNote.AddToAppMessage(ctx, name, cr)
+}
+
+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, "```")
+
+	return strings.TrimSpace(aiSummary)
 }
diff --git a/pkg/checks/diff/ai_summary_test.go b/pkg/checks/diff/ai_summary_test.go
new file mode 100644
index 00000000..55d336c9
--- /dev/null
+++ b/pkg/checks/diff/ai_summary_test.go
@@ -0,0 +1,33 @@
+package diff
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCleanUpAiSummary(t *testing.T) {
+	t.Run("prefix", func(t *testing.T) {
+		input := "```markdown\nhello\nworld"
+		expected := "hello\nworld"
+
+		actual := cleanUpAiSummary(input)
+		assert.Equal(t, expected, actual)
+	})
+
+	t.Run("suffix", func(t *testing.T) {
+		input := "\nhello\nworld```"
+		expected := "hello\nworld"
+
+		actual := cleanUpAiSummary(input)
+		assert.Equal(t, expected, actual)
+	})
+
+	t.Run("prefix and suffix", func(t *testing.T) {
+		input := "```markdown\n\nhello\nworld```"
+		expected := "hello\nworld"
+
+		actual := cleanUpAiSummary(input)
+		assert.Equal(t, expected, actual)
+	})
+}
diff --git a/pkg/checks/diff/diff.go b/pkg/checks/diff/diff.go
index ed2b6e64..f4503896 100644
--- a/pkg/checks/diff/diff.go
+++ b/pkg/checks/diff/diff.go
@@ -131,7 +131,7 @@ func Check(ctx context.Context, request checks.Request) (msg.Result, error) {
 
 	cr.Details = fmt.Sprintf("```diff\n%s\n```", renderedDiff)
 
-	aiDiffSummary(ctx, request.Note, request.Container.Config, request.AppName, request.JsonManifests, renderedDiff)
+	aiDiffSummary(ctx, request.Note, request.Container.Config, request.AppName, renderedDiff)
 
 	return cr, nil
 }