internal/palmapi: rename package palmapi to genai
Rename palmapi package to "genai" and rename Client to PaLMClient.
This will allow us to more easily migrate to the new Gemini API.
Change-Id: I027ce9f6602b4e0e0b1781999520f539fcc4cf8a
Reviewed-on: https://go-review.googlesource.com/c/vulndb/+/552035
Reviewed-by: Damien Neil <dneil@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/cmd/vulnreport/suggest.go b/cmd/vulnreport/suggest.go
index 54be40b..1d9d39b 100644
--- a/cmd/vulnreport/suggest.go
+++ b/cmd/vulnreport/suggest.go
@@ -10,7 +10,7 @@
"fmt"
"golang.org/x/vulndb/internal/derrors"
- "golang.org/x/vulndb/internal/palmapi"
+ "golang.org/x/vulndb/internal/genai"
"golang.org/x/vulndb/internal/report"
)
@@ -22,14 +22,14 @@
func suggest(_ context.Context, filename string) (err error) {
defer derrors.Wrap(&err, "suggest(%q)", filename)
- c := palmapi.NewDefaultClient()
+ c := genai.NewDefaultPaLMClient()
r, err := report.Read(filename)
if err != nil {
return err
}
- suggestions, err := c.Suggest(&palmapi.Input{
+ suggestions, err := c.Suggest(&genai.Input{
Module: r.Modules[0].Module,
Description: r.Description.String(),
})
diff --git a/internal/palmapi/data/examples.csv b/internal/genai/data/examples.csv
similarity index 100%
rename from internal/palmapi/data/examples.csv
rename to internal/genai/data/examples.csv
diff --git a/internal/palmapi/data/examples.json b/internal/genai/data/examples.json
similarity index 100%
rename from internal/palmapi/data/examples.json
rename to internal/genai/data/examples.json
diff --git a/internal/palmapi/data/prompt.txt b/internal/genai/data/prompt.txt
similarity index 100%
rename from internal/palmapi/data/prompt.txt
rename to internal/genai/data/prompt.txt
diff --git a/internal/palmapi/examples.go b/internal/genai/examples.go
similarity index 99%
rename from internal/palmapi/examples.go
rename to internal/genai/examples.go
index 3aad439..bdf855b 100644
--- a/internal/palmapi/examples.go
+++ b/internal/genai/examples.go
@@ -2,7 +2,7 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-package palmapi
+package genai
import (
"encoding/csv"
diff --git a/internal/palmapi/examples_test.go b/internal/genai/examples_test.go
similarity index 99%
rename from internal/palmapi/examples_test.go
rename to internal/genai/examples_test.go
index 3635481..f5e1626 100644
--- a/internal/palmapi/examples_test.go
+++ b/internal/genai/examples_test.go
@@ -2,7 +2,7 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-package palmapi
+package genai
import (
"bytes"
diff --git a/internal/palmapi/gen_examples/main.go b/internal/genai/gen_examples/main.go
similarity index 90%
rename from internal/palmapi/gen_examples/main.go
rename to internal/genai/gen_examples/main.go
index 4fce1c8..12ae903 100644
--- a/internal/palmapi/gen_examples/main.go
+++ b/internal/genai/gen_examples/main.go
@@ -17,17 +17,17 @@
"strings"
"golang.org/x/exp/maps"
+ "golang.org/x/vulndb/internal/genai"
"golang.org/x/vulndb/internal/genericosv"
"golang.org/x/vulndb/internal/ghsarepo"
"golang.org/x/vulndb/internal/gitrepo"
- "golang.org/x/vulndb/internal/palmapi"
"golang.org/x/vulndb/internal/report"
"golang.org/x/vulndb/internal/stdlib"
)
var (
localGHSA = flag.String("lg", os.Getenv("LOCAL_GHSA_DB"), "path to local GHSA repo, instead of cloning remote")
- outFolder = flag.String("out", filepath.Join("internal", "palmapi"), "folder to write files to")
+ outFolder = flag.String("out", filepath.Join("internal", "genai"), "folder to write files to")
)
func main() {
@@ -120,18 +120,18 @@
return false
}
-func toExamples(vs []*vuln) (palmapi.Examples, error) {
- var es palmapi.Examples
+func toExamples(vs []*vuln) (genai.Examples, error) {
+ var es genai.Examples
for _, v := range vs {
if v.r == nil || v.ghsa == nil {
return nil, errors.New("invalid example")
}
- ex := &palmapi.Example{
- Input: palmapi.Input{
+ ex := &genai.Example{
+ Input: genai.Input{
Module: v.r.Modules[0].Module,
Description: v.ghsa.Details,
},
- Suggestion: palmapi.Suggestion{
+ Suggestion: genai.Suggestion{
Summary: removeNewlines(v.r.Summary.String()),
Description: removeNewlines(v.r.Description.String()),
},
diff --git a/internal/palmapi/client.go b/internal/genai/palmclient.go
similarity index 91%
rename from internal/palmapi/client.go
rename to internal/genai/palmclient.go
index 6d812a6..f33337e 100644
--- a/internal/palmapi/client.go
+++ b/internal/genai/palmclient.go
@@ -2,9 +2,9 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-// Package palmapi provides a client and utilities for interacting with
-// the PaLM API (https://developers.generativeai.google/guide/palm_api_overview).
-package palmapi
+// Package genai provides a client and utilities for interacting with
+// Google's generative AI libraries.
+package genai
import (
"bytes"
@@ -15,15 +15,15 @@
"os"
)
-type Client struct {
+type PaLMClient struct {
c *http.Client
url string
getAPIKey func() (string, error)
}
-// NewDefaultClient returns a new default client for the PaLM API that reads
+// NewDefaultPaLMClient returns a new default client for the PaLM API that reads
// an API key from the environment variable "PALM_API_KEY".
-func NewDefaultClient() *Client {
+func NewDefaultPaLMClient() *PaLMClient {
const (
defaultURL = `https://generativelanguage.googleapis.com`
apiKeyEnv = "PALM_API_KEY"
@@ -37,8 +37,8 @@
})
}
-func NewClient(httpClient *http.Client, url string, getAPIKey func() (string, error)) *Client {
- return &Client{
+func NewClient(httpClient *http.Client, url string, getAPIKey func() (string, error)) *PaLMClient {
+ return &PaLMClient{
c: httpClient,
url: url,
getAPIKey: getAPIKey}
@@ -49,7 +49,7 @@
// GenerateText is a wrapper for the PaLM API "generateText" endpoint.
// See https://developers.generativeai.google/api/rest/generativelanguage/models/generateText.
-func (c *Client) GenerateText(prompt string) (*GenerateTextResponse, error) {
+func (c *PaLMClient) GenerateText(prompt string) (*GenerateTextResponse, error) {
reqBody, err := toRequestBody(prompt)
if err != nil {
return nil, err
diff --git a/internal/palmapi/client_test.go b/internal/genai/palmclient_test.go
similarity index 90%
rename from internal/palmapi/client_test.go
rename to internal/genai/palmclient_test.go
index 21be679..2d107bb 100644
--- a/internal/palmapi/client_test.go
+++ b/internal/genai/palmclient_test.go
@@ -2,7 +2,7 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-package palmapi
+package genai
import (
"context"
@@ -16,7 +16,7 @@
"strings"
"testing"
- "github.com/google/generative-ai-go/genai"
+ gemini "github.com/google/generative-ai-go/genai"
"google.golang.org/api/option"
)
@@ -98,7 +98,7 @@
}
}
-func testClient(endpoint, prompt string, response *GenerateTextResponse) (c *Client, cleanup func(), err error) {
+func testClient(endpoint, prompt string, response *GenerateTextResponse) (c *PaLMClient, cleanup func(), err error) {
rBytes, err := json.Marshal(response)
if err != nil {
return nil, nil, err
@@ -135,7 +135,7 @@
return NewClient(s.Client(), s.URL, getTestAPIKey), func() { s.Close() }, nil
}
-func testClientErr(endpoint, prompt string, errMsg string) (c *Client, cleanup func(), err error) {
+func testClientErr(endpoint, prompt string, errMsg string) (c *PaLMClient, cleanup func(), err error) {
handler := func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusBadRequest)
errJSON := fmt.Sprintf(`{"error":{"message":"%s"}}`, errMsg)
@@ -154,13 +154,13 @@
func TestGemini(t *testing.T) {
t.SkipNow()
ctx := context.Background()
- client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+ client, err := gemini.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
if err != nil {
t.Fatal(err)
}
defer client.Close()
m := client.GenerativeModel("gemini-pro")
- res, err := m.GenerateContent(ctx, genai.Text("something"))
+ res, err := m.GenerateContent(ctx, gemini.Text("something"))
if err != nil {
t.Fatal(err)
}
diff --git a/internal/palmapi/suggest.go b/internal/genai/suggest.go
similarity index 91%
rename from internal/palmapi/suggest.go
rename to internal/genai/suggest.go
index f4ed086..fcd0d0b 100644
--- a/internal/palmapi/suggest.go
+++ b/internal/genai/suggest.go
@@ -2,7 +2,7 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-package palmapi
+package genai
import (
"encoding/json"
@@ -28,12 +28,13 @@
Description string
}
-// Suggest uses the PaLM API to generate suggestions for vulnerability
+// Suggest uses generative AI to generate suggestions for vulnerability
// reports based on the input.
+//
// This function must be called from the root of the vulndb repo,
// as it accesses a specific file to read examples.
-func (c *Client) Suggest(in *Input) ([]*Suggestion, error) {
- examples, err := readFile(filepath.Join("internal", "palmapi", dataFolder, jsonFile))
+func (c *PaLMClient) Suggest(in *Input) ([]*Suggestion, error) {
+ examples, err := readFile(filepath.Join("internal", "genai", dataFolder, jsonFile))
if err != nil {
return nil, err
}
@@ -44,7 +45,7 @@
return c.suggest(prompt)
}
-func (c *Client) suggest(prompt string) ([]*Suggestion, error) {
+func (c *PaLMClient) suggest(prompt string) ([]*Suggestion, error) {
response, err := c.GenerateText(prompt)
if err != nil {
return nil, err
diff --git a/internal/palmapi/suggest_test.go b/internal/genai/suggest_test.go
similarity index 99%
rename from internal/palmapi/suggest_test.go
rename to internal/genai/suggest_test.go
index 60b4e3d..13fd818 100644
--- a/internal/palmapi/suggest_test.go
+++ b/internal/genai/suggest_test.go
@@ -2,7 +2,7 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
-package palmapi
+package genai
import (
"strings"
