fix: explain result order (#98)

Author: miri-yehezkel

llm/api_client.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"strings"
 )
 
 var (
@@ -96,17 +97,28 @@ func (d *DeepCodeLLMBindingImpl) explainRequestBody(options *ExplainOptions) ([]
 	} else {
 		requestBody, marshalErr = json.Marshal(explainFixRequest{
 			RuleId:            options.RuleKey,
-			Diffs:             encodeDiffs(options.Diffs),
+			Diffs:             prepareDiffs(options.Diffs),
 			ExplanationLength: SHORT,
 		})
 		logger.Debug().Msg("payload for FixExplanation")
 	}
 	return requestBody, marshalErr
 }
 
-func encodeDiffs(diffs []string) []string {
-	var encodedDiffs []string
+func prepareDiffs(diffs []string) []string {
+	cleanedDiffs := make([]string, 0, len(diffs))
 	for _, diff := range diffs {
+		diffLines := strings.Split(diff, "\n")
+		cleanedLines := ""
+		for _, line := range diffLines {
+			if !strings.HasPrefix(line, "---") && !strings.HasPrefix(line, "+++") {
+				cleanedLines += line + "\n"
+			}
+		}
+		cleanedDiffs = append(cleanedDiffs, cleanedLines)
+	}
+	var encodedDiffs []string
+	for _, diff := range cleanedDiffs {
 		encodedDiffs = append(encodedDiffs, base64.StdEncoding.EncodeToString([]byte(diff)))
 	}
 	return encodedDiffs

llm/api_client_test.go

@@ -2,6 +2,7 @@ package llm
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -151,7 +152,7 @@ func TestDeepcodeLLMBinding_explainRequestBody(t *testing.T) {
 
 		assert.NotNil(t, request)
 		assert.Equal(t, "test-rule-key", request.RuleId)
-		expectedEncodedDiffs := encodeDiffs([]string{"test-Diffs"})
+		expectedEncodedDiffs := prepareDiffs([]string{"test-Diffs"})
 		assert.Equal(t, expectedEncodedDiffs, request.Diffs)
 		assert.Equal(t, SHORT, request.ExplanationLength)
 	})
@@ -206,6 +207,52 @@ func TestEndpoint(t *testing.T) {
 	}
 }
 
+func TestPrepareDiffs(t *testing.T) {
+	testCases := []struct {
+		name     string
+		input    []string
+		expected []string
+	}{
+		{
+			name: "Single diff with headers and content",
+			input: []string{
+				"--- a/file.txt\n" +
+					"+++ b/file.txt\n" +
+					"@@ -1,1 +1,1 @@\n" +
+					"-old line\n" +
+					"+new line\n",
+			},
+			expected: []string{
+				base64.StdEncoding.EncodeToString([]byte("@@ -1,1 +1,1 @@\n-old line\n+new line\n\n")),
+			},
+		},
+		{
+			name: "Multiple diffs",
+			input: []string{
+				"--- a/file1.txt\n" +
+					"+++ b/file1.txt\n" +
+					"-line 1\n" +
+					"+line 2\n",
+				"--- a/file2.txt\n" +
+					"+++ b/file2.txt\n" +
+					"content2a\n" +
+					"+content2b\n",
+			},
+			expected: []string{
+				base64.StdEncoding.EncodeToString([]byte("-line 1\n+line 2\n\n")),
+				base64.StdEncoding.EncodeToString([]byte("content2a\n+content2b\n\n")),
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actual := prepareDiffs(tc.input)
+			assert.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
 // Helper function for testing
 func testLogger(t *testing.T) *zerolog.Logger {
 	t.Helper()

llm/binding.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/url"
+	"slices"
 
 	"github.com/rs/zerolog"
 
@@ -45,7 +46,7 @@ type SnykLLMBindings interface {
 	// output - a channel that can be used to stream the results
 	Explain(ctx context.Context, input AIRequest, format OutputFormat, output chan<- string) error
 }
-type ExplainResult map[string]string
+type ExplainResult []string
 
 type DeepCodeLLMBinding interface {
 	SnykLLMBindings
@@ -69,15 +70,24 @@ func (d *DeepCodeLLMBindingImpl) ExplainWithOptions(ctx context.Context, options
 	if err != nil {
 		return explainResult, err
 	}
-	index := 0
-	for _, explanation := range response {
-		if index < len(options.Diffs) {
-			explainResult[options.Diffs[index]] = explanation
-		}
-		index++
+
+	orderedExplainResults := getOrderedResponse(response)
+
+	return orderedExplainResults, nil
+}
+
+func getOrderedResponse(explainResponse Explanations) []string {
+	explainMapKeys := make([]string, 0, len(explainResponse))
+	for k := range explainResponse {
+		explainMapKeys = append(explainMapKeys, k)
 	}
+	slices.Sort(explainMapKeys)
 
-	return explainResult, nil
+	orderedValues := make([]string, 0, len(explainResponse))
+	for _, key := range explainMapKeys {
+		orderedValues = append(orderedValues, explainResponse[key])
+	}
+	return orderedValues
 }
 
 func (d *DeepCodeLLMBindingImpl) PublishIssues(_ context.Context, _ []map[string]string) error {

llm/binding_test.go

@@ -29,8 +29,14 @@ func TestExplainWithOptions(t *testing.T) {
 		d, mockHTTPClient := getHTTPMockedBinding(t)
 
 		explainResponseJSON := explainResponse{
-			Status:      completeStatus,
-			Explanation: map[string]string{"explanation1": "This is the first explanation"},
+			Status: completeStatus,
+			Explanation: map[string]string{
+				"explanation1": "This is the first explanation",
+				"explanation2": "this is the second explanation",
+				"explanation3": "this is the third explanation",
+				"explanation4": "this is the fourth explanation",
+				"explanation5": "this is the fifth explanation",
+			},
 		}
 
 		expectedResponseBody, err := json.Marshal(explainResponseJSON)
@@ -49,7 +55,11 @@ func TestExplainWithOptions(t *testing.T) {
 		err = json.Unmarshal(expectedResponseBody, &exptectedExplanationsResponse)
 		assert.NoError(t, err)
 		expectedResExplanations := exptectedExplanationsResponse.Explanation
-		assert.Equal(t, expectedResExplanations["explanation1"], explanation[testDiff])
+		assert.Equal(t, expectedResExplanations["explanation1"], explanation[0])
+		assert.Equal(t, expectedResExplanations["explanation2"], explanation[1])
+		assert.Equal(t, expectedResExplanations["explanation3"], explanation[2])
+		assert.Equal(t, expectedResExplanations["explanation4"], explanation[3])
+		assert.Equal(t, expectedResExplanations["explanation5"], explanation[4])
 	})
 
 	t.Run("runExplain error", func(t *testing.T) {