add googleai

Author: batmac

go.mod

@@ -25,6 +25,7 @@ require (
 	github.com/gabriel-vasile/mimetype v1.4.5
 	github.com/gage-technologies/mistral-go v1.1.0
 	github.com/gdamore/tcell/v2 v2.7.4
+	github.com/google/generative-ai-go v0.18.0
 	github.com/google/renameio/v2 v2.0.0
 	github.com/gowebpki/jcs v1.0.1
 	github.com/hbollon/go-edlib v1.6.0
@@ -54,6 +55,11 @@ require (
 	sigs.k8s.io/yaml v1.4.0
 )
 
+require (
+	cloud.google.com/go/ai v0.8.0 // indirect
+	cloud.google.com/go/longrunning v0.5.9 // indirect
+)
+
 require (
 	cloud.google.com/go v0.115.0 // indirect
 	cloud.google.com/go/auth v0.7.1 // indirect
@@ -131,7 +137,7 @@ require (
 	golang.org/x/sys v0.25.0 // indirect
 	golang.org/x/text v0.18.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
-	google.golang.org/api v0.188.0 // indirect
+	google.golang.org/api v0.188.0
 	google.golang.org/genproto v0.0.0-20240711142825-46eb208f015d // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240711142825-46eb208f015d // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240711142825-46eb208f015d // indirect

go.sum

@@ -1,6 +1,8 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.115.0 h1:CnFSK6Xo3lDYRoBKEcAtia6VSC837/ZkJuRduSFnr14=
 cloud.google.com/go v0.115.0/go.mod h1:8jIM5vVgoAEoiVxQ/O4BFTfHqulPZgs/ufEzMcFMdWU=
+cloud.google.com/go/ai v0.8.0 h1:rXUEz8Wp2OlrM8r1bfmpF2+VKqc1VJpafE3HgzRnD/w=
+cloud.google.com/go/ai v0.8.0/go.mod h1:t3Dfk4cM61sytiggo2UyGsDVW3RF1qGZaUKDrZFyqkE=
 cloud.google.com/go/auth v0.7.1 h1:Iv1bbpzJ2OIg16m94XI9/tlzZZl3cdeR3nGVGj78N7s=
 cloud.google.com/go/auth v0.7.1/go.mod h1:VEc4p5NNxycWQTMQEDQF0bd6aTMb6VgYDXEwiJJQAbs=
 cloud.google.com/go/auth/oauth2adapt v0.2.3 h1:MlxF+Pd3OmSudg/b1yZ5lJwoXCEaeedAguodky1PcKI=
@@ -155,6 +157,8 @@ github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QD
 github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/google/generative-ai-go v0.18.0 h1:6ybg9vOCLcI/UpBBYXOTVgvKmcUKFRNj+2Cj3GnebSo=
+github.com/google/generative-ai-go v0.18.0/go.mod h1:JYolL13VG7j79kM5BtHz4qwONHkeJQzOCkKXnpqtS/E=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=

pkg/mutators/single/googleai.go

@@ -0,0 +1,80 @@
+package mutators
+
+import (
+	"context"
+	"errors"
+	"io"
+	"os"
+
+	"github.com/batmac/ccat/pkg/log"
+	"github.com/google/generative-ai-go/genai"
+	"google.golang.org/api/iterator"
+	"google.golang.org/api/option"
+)
+
+func init() {
+	singleRegister("googleai", googleai,
+		withDescription("googleai, X:gemini-1.5-flash is the model (Requires a valid key in $GOOGLE_API_KEY)"),
+		withAliases("gai"),
+		withHintSlow(), // output asap (when no other mutator is used)
+		withCategory("external APIs"),
+		withConfigBuilder(stdConfigStringWithDefault("gemini-1.5-flash")),
+	)
+}
+
+func googleai(w io.WriteCloser, r io.ReadCloser, config any) (int64, error) {
+	modelname := config.(string)
+
+	ctx := context.Background()
+
+	apikey := os.Getenv("GOOGLE_API_KEY")
+	log.Debugln("masked apikey: ", mask(apikey))
+	client, err := genai.NewClient(ctx, option.WithAPIKey(apikey))
+	if err != nil {
+		return 0, err
+	}
+	defer client.Close()
+
+	model := client.GenerativeModel(modelname)
+
+	prompt, err := io.ReadAll(r)
+	if err != nil {
+		return 0, err
+	}
+
+	iter := model.GenerateContentStream(ctx, genai.Text(prompt))
+
+	totalWritten := int64(0)
+	for {
+		resp, err := iter.Next()
+		if errors.Is(err, iterator.Done) {
+			break
+		}
+		if err != nil {
+			log.Debugln("error: ", err)
+			return 0, err
+		}
+
+		p := ""
+		for i, c := range resp.Candidates[0].Content.Parts {
+			log.Debugln("part ", i, ": ", c)
+			if t, ok := c.(genai.Text); ok {
+				p += string(t)
+			}
+		}
+
+		n, err := io.WriteString(w, p)
+		if err != nil {
+			return 0, err
+		}
+		totalWritten += int64(n)
+	}
+	return totalWritten, nil
+}
+
+func mask(s string) string {
+	if len(s) <= 8 {
+		return "********"
+	}
+	return "****..." + s[len(s)-2:]
+}