/
suggest.go
80 lines (67 loc) · 1.88 KB
/
suggest.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package main
import (
"context"
"flag"
"fmt"
"github.com/khulnasoft-lab/go-vulndb/internal/derrors"
"github.com/khulnasoft-lab/go-vulndb/internal/palmapi"
"github.com/khulnasoft-lab/go-vulndb/internal/report"
)
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)")
)
func suggest(_ context.Context, filename string) (err error) {
defer derrors.Wrap(&err, "suggest(%q)", filename)
c := palmapi.NewDefaultClient()
r, err := report.Read(filename)
if err != nil {
return err
}
suggestions, err := c.Suggest(&palmapi.Input{
Module: r.Modules[0].Module,
Description: r.Description,
})
if err != nil {
return fmt.Errorf("PaLM API error: %s", err)
}
if len(suggestions) > *numSuggestions {
suggestions = suggestions[:*numSuggestions]
}
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)
for i, s := range suggestions {
outlog.Printf("\nSuggestion %d/%d\nsummary: %s\ndescription: %s\n",
i+1, found, s.Summary, s.Description)
// In interactive mode, allow user to accept the suggestion,
// see the next one, or quit.
if *interactive {
if i == found-1 {
outlog.Printf("\naccept or quit? (a=accept/Q=quit) ")
} else {
outlog.Printf("\naccept, see next suggestion, or quit? (a=accept/n=next/Q=quit) ")
}
var choice string
if _, err := fmt.Scanln(&choice); err != nil {
return err
}
switch choice {
case "a":
r.Summary = s.Summary
r.Description = s.Description
if err := r.Write(filename); err != nil {
errlog.Println(err)
}
return nil
case "n":
continue
default:
return nil
}
}
}
return nil
}