cmd/vulnreport: add option to use AI in vulnreport create

Add flag "-ai" that causes `vulnreport create` to attempt to auto-populate
an AI-generated draft of a summary and description for a report.

(This functionality can still be accessed post-create via command
`vulnreport suggest`.)

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

Author: tatianab

cmd/vulnreport/create.go

@@ -15,6 +15,7 @@ import (
 	"golang.org/x/vulndb/internal/cveclient"
 	"golang.org/x/vulndb/internal/cveschema5"
 	"golang.org/x/vulndb/internal/derrors"
+	"golang.org/x/vulndb/internal/genai"
 	"golang.org/x/vulndb/internal/genericosv"
 	"golang.org/x/vulndb/internal/ghsa"
 	"golang.org/x/vulndb/internal/gitrepo"
@@ -29,6 +30,7 @@ var (
 	closedOk  = flag.Bool("closed-ok", false, "for create & create-excluded, allow closed issues to be created")
 	graphQL   = flag.Bool("graphql", false, "for create, fetch GHSAs from the Github GraphQL API instead of the OSV database")
 	issueRepo = flag.String("issue-repo", "github.com/golang/vulndb", "for create, repo locate Github issues")
+	useAI     = flag.Bool("ai", false, "for create, use AI to write draft summary and description when creating report")
 )
 
 func create(ctx context.Context, issueNumber int, cfg *createCfg) (err error) {
@@ -126,6 +128,7 @@ type createCfg struct {
 	existingByFile  map[string]*report.Report
 	existingByIssue map[int]*report.Report
 	allowClosed     bool
+	aiClient        *genai.GeminiClient
 }
 
 func setupCreate(ctx context.Context, args []string) ([]int, *createCfg, error) {
@@ -148,13 +151,21 @@ func setupCreate(ctx context.Context, args []string) ([]int, *createCfg, error)
 	if err != nil {
 		return nil, nil, err
 	}
+	var aiClient *genai.GeminiClient
+	if *useAI {
+		aiClient, err = genai.NewGeminiClient(ctx)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
 	return githubIDs, &createCfg{
 		issuesClient:    issues.NewClient(ctx, &issues.Config{Owner: owner, Repo: repoName, Token: *githubToken}),
 		ghsaClient:      ghsa.NewClient(ctx, *githubToken),
 		proxyClient:     proxy.NewDefaultClient(),
 		existingByFile:  existingByFile,
 		existingByIssue: existingByIssue,
 		allowClosed:     *closedOk,
+		aiClient:        aiClient,
 	}, nil
 }
 
@@ -244,6 +255,18 @@ func createReport(ctx context.Context, cfg *createCfg, iss *issues.Issue) (r *re
 	// Find any additional aliases referenced by the source aliases.
 	addMissingAliases(ctx, r, cfg.ghsaClient)
 
+	if cfg.aiClient != nil {
+		suggestions, err := suggest(ctx, cfg.aiClient, r, 1)
+		if err != nil {
+			warnlog.Printf("failed to get AI-generated suggestions for %s: %v\n", r.ID, err)
+		} else if len(suggestions) == 0 {
+			warnlog.Printf("failed to get AI-generated suggestions for %s (none generated)\n", r.ID)
+		} else {
+			infolog.Printf("applying AI-generated suggestion for %s", r.ID)
+			applySuggestion(r, suggestions[0])
+		}
+	}
+
 	addTODOs(r)
 	return r, nil
 }

cmd/vulnreport/main.go

@@ -128,7 +128,7 @@ func main() {
 	case "lint":
 		cmdFunc = func(ctx context.Context, name string) error { return lint(ctx, name, pc) }
 	case "suggest":
-		cmdFunc = func(ctx context.Context, name string) error { return suggest(ctx, name) }
+		cmdFunc = func(ctx context.Context, name string) error { return suggestCmd(ctx, name) }
 	case "commit":
 		cmdFunc = func(ctx context.Context, name string) error { return commit(ctx, name, ghsaClient, pc, *force) }
 	case "cve":

cmd/vulnreport/suggest.go

@@ -20,7 +20,7 @@ var (
 	palm           = flag.Bool("palm", false, "use the legacy PaLM API instead of the Gemini API")
 )
 
-func suggest(ctx context.Context, filename string) (err error) {
+func suggestCmd(ctx context.Context, filename string) (err error) {
 	defer derrors.Wrap(&err, "suggest(%q)", filename)
 
 	r, err := report.Read(filename)
@@ -40,21 +40,11 @@ func suggest(ctx context.Context, filename string) (err error) {
 		}
 	}
 
-	suggestions, err := genai.Suggest(ctx, c, &genai.Input{
-		Module:      r.Modules[0].Module,
-		Description: r.Description.String(),
-	})
+	suggestions, err := suggest(ctx, c, r, *numSuggestions)
 	if err != nil {
-		return fmt.Errorf("GenAI API error: %s", err)
-	}
-	if len(suggestions) > *numSuggestions {
-		suggestions = suggestions[:*numSuggestions]
+		return err
 	}
-
 	found := len(suggestions)
-	if found == 0 {
-		return fmt.Errorf("could not generate any valid suggestions for report %s (try again?)", r.ID)
-	}
 
 	outlog.Printf("== AI-generated suggestions for report %s ==\n", r.ID)
 
@@ -77,8 +67,7 @@ func suggest(ctx context.Context, filename string) (err error) {
 			}
 			switch choice {
 			case "a":
-				r.Summary = report.Summary(s.Summary)
-				r.Description = report.Description(s.Description)
+				applySuggestion(r, s)
 				if err := r.Write(filename); err != nil {
 					errlog.Println(err)
 				}
@@ -93,3 +82,29 @@ func suggest(ctx context.Context, filename string) (err error) {
 
 	return nil
 }
+
+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)
+	}
+	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 suggestions, nil
+}
+
+func applySuggestion(r *report.Report, s *genai.Suggestion) {
+	r.Summary = report.Summary(s.Summary)
+	r.Description = report.Description(s.Description)
+	r.Fix(nil)
+}