cmd/vulnreport: fix handling of numSuggestions for Gemini

Unlike the PaLM API, which could generate multiple candidate responses
at once, the Gemini API only gives one response at a time (at least for
now).

To accommodate this, modify suggest to make a separate request to the
API to generate each suggestion.

Change-Id: I5c3adce302abdb5228aef7a8a94d8b3ee5881371
Reviewed-on: https://go-review.googlesource.com/c/vulndb/+/559599
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Damien Neil <[email protected]>

Author: tatianab

cmd/vulnreport/suggest.go

@@ -16,7 +16,7 @@ import (
 
 var (
 	interactive    = flag.Bool("i", false, "for suggest, interactive mode")
-	numSuggestions = flag.Int("n", 4, "for suggest, the number of suggestions to attempt to generate (max is 8)")
+	numSuggestions = flag.Int("n", 1, "for suggest, the number of suggestions to generate (>1 can be slow)")
 )
 
 func suggestCmd(ctx context.Context, filename string) (err error) {
@@ -47,6 +47,8 @@ func suggestCmd(ctx context.Context, filename string) (err error) {
 
 		// In interactive mode, allow user to accept the suggestion,
 		// see the next one, or quit.
+		// TODO(tatianabradley): In interactive mode, call the API as requested
+		// instead of upfront.
 		if *interactive {
 			if i == found-1 {
 				outlog.Printf("\naccept or quit? (a=accept/Q=quit) ")
@@ -77,20 +79,27 @@ func suggestCmd(ctx context.Context, filename string) (err error) {
 }
 
 func suggest(ctx context.Context, c genai.Client, r *report.Report, max int) (suggestions []*genai.Suggestion, err error) {
-	suggestions, err = genai.Suggest(ctx, c, &genai.Input{
-		Module:      r.Modules[0].Module,
-		Description: r.Description.String(),
-	})
-	if err != nil {
-		return nil, fmt.Errorf("GenAI API error: %s", err)
+	attempts := 0
+	maxAttempts := max + 2
+	for len(suggestions) < max && attempts < maxAttempts {
+		s, err := genai.Suggest(ctx, c, &genai.Input{
+			Module:      r.Modules[0].Module,
+			Description: r.Description.String(),
+		})
+		if err != nil {
+			return nil, fmt.Errorf("GenAI API error: %s", err)
+		}
+		suggestions = append(suggestions, s...)
+		attempts++
 	}
+
 	if len(suggestions) > max {
 		suggestions = suggestions[:max]
 	}
 
 	found := len(suggestions)
 	if found == 0 {
-		return nil, fmt.Errorf("could not generate any valid suggestions for report %s (try again?)", r.ID)
+		return nil, fmt.Errorf("could not generate any valid suggestions for report %s after %d attempts", r.ID, attempts)
 	}
 
 	return suggestions, nil