feat(vertexai): Vertex AI for go (#9095)

go.work

@@ -149,6 +149,7 @@ use (
 	./tpu
 	./trace
 	./translate
+	./vertexai
 	./video
 	./videointelligence
 	./vision

vertexai/README.md

@@ -0,0 +1,13 @@
+# Vertex AI Go SDK
+
+[!IMPORTANT]
+Thanks for your interest in the Vertex AI SDKs! **You can start using this SDK
+and its samples on December 13, 2023.** Until then, check out our [blog
+post](https://blog.google/technology/ai/google-gemini-ai/) to learn more about
+Google's Gemini multimodal model.
+
+The Vertex AI Go SDK enables developers to use Google's state-of-the-art
+generative AI models (like Gemini) to build AI-powered features and
+applications.
+
+*More details and information coming soon!*

vertexai/cmd/gen/gen.go

@@ -0,0 +1,106 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"strings"
+
+	"cloud.google.com/go/vertexai/genai"
+
+	"google.golang.org/api/iterator"
+)
+
+var (
+	project   = flag.String("project", "", "project ID")
+	location  = flag.String("location", "", "location")
+	temp      = flag.Float64("t", 0.2, "temperature")
+	model     = flag.String("model", "", "model")
+	streaming = flag.Bool("stream", false, "using the streaming API")
+)
+
+func main() {
+	flag.Parse()
+	if *project == "" || *location == "" || *model == "" {
+		log.Fatal("need -project, -location, and -model")
+	}
+
+	ctx := context.Background()
+	client, err := genai.NewClient(ctx, *project, *location)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer client.Close()
+	model := client.GenerativeModel(*model)
+	model.Temperature = float32(*temp)
+
+	model.SafetySettings = []*genai.SafetySetting{
+		{
+			Category:  genai.HarmCategorySexuallyExplicit,
+			Threshold: genai.HarmBlockLowAndAbove,
+		},
+		{
+			Category:  genai.HarmCategoryDangerousContent,
+			Threshold: genai.HarmBlockLowAndAbove,
+		},
+	}
+
+	text := strings.Join(flag.Args(), " ")
+	if *streaming {
+		iter := model.GenerateContentStream(ctx, genai.Text(text))
+		for {
+			res, err := iter.Next()
+			if err == iterator.Done {
+				break
+			}
+			if err != nil {
+				showError(err)
+			}
+			showJSON(res)
+			fmt.Println("---")
+		}
+	} else {
+		res, err := model.GenerateContent(ctx, genai.Text(text))
+		if err != nil {
+			showError(err)
+		}
+		showJSON(res)
+	}
+}
+
+func showJSON(x any) {
+	bs, err := json.MarshalIndent(x, "", "    ")
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Printf("%s\n", bs)
+}
+
+func showError(err error) {
+	var berr *genai.BlockedError
+	if errors.As(err, &berr) {
+		fmt.Println("ERROR:")
+		showJSON(err)
+	} else {
+		fmt.Printf("ERROR: %s\n", err)
+	}
+	os.Exit(1)
+}

vertexai/genai/chat.go

@@ -0,0 +1,70 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package genai
+
+import (
+	"context"
+)
+
+// A ChatSession provides interactive chat.
+type ChatSession struct {
+	m       *GenerativeModel
+	History []*Content
+}
+
+// StartChat starts a chat session.
+func (m *GenerativeModel) StartChat() *ChatSession {
+	return &ChatSession{m: m}
+}
+
+// SendMessage sends a request to the model as part of a chat session.
+func (cs *ChatSession) SendMessage(ctx context.Context, parts ...Part) (*GenerateContentResponse, error) {
+	// Call the underlying client with the entire history plus the argument Content.
+	cs.History = append(cs.History, newUserContent(parts))
+	req := cs.m.newRequest(cs.History...)
+	cc := int32(1)
+	req.GenerationConfig.CandidateCount = &cc
+	resp, err := cs.m.generateContent(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	cs.addToHistory(resp.Candidates)
+	return resp, nil
+}
+
+// SendMessageStream is like SendMessage, but with a streaming request.
+func (cs *ChatSession) SendMessageStream(ctx context.Context, parts ...Part) *GenerateContentResponseIterator {
+	cs.History = append(cs.History, newUserContent(parts))
+	req := cs.m.newRequest(cs.History...)
+	var cc int32 = 1
+	req.GenerationConfig.CandidateCount = &cc
+	streamClient, err := cs.m.c.c.StreamGenerateContent(ctx, req)
+	return &GenerateContentResponseIterator{
+		sc:  streamClient,
+		err: err,
+		cs:  cs,
+	}
+}
+
+// By default, use the first candidate for history. The user can modify that if they want.
+func (cs *ChatSession) addToHistory(cands []*Candidate) bool {
+	if len(cands) > 0 {
+		c := cands[0].Content
+		c.Role = roleModel
+		cs.History = append(cs.History, c)
+		return true
+	}
+	return false
+}