chore(vertexai): consistent usage of context.Background (#3797)

GoogleCloudPlatform · Feb 22, 2024 · 275d117 · 275d117
1 parent 330296f
commit 275d117
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/vertexai/chat/chat.go b/vertexai/chat/chat.go
@@ -28,7 +28,8 @@ var region = "us-central1"
 var modelName = "gemini-1.0-pro-vision"
 
 func makeChatRequests(projectId string, region string, modelName string) error {
-	client, err := genai.NewClient(context.TODO(), projectId, region)
+	ctx := context.Background()
+	client, err := genai.NewClient(ctx, projectId, region)
 	if err != nil {
 		return fmt.Errorf("error creating client: %v", err)
 	}
@@ -38,7 +39,7 @@ func makeChatRequests(projectId string, region string, modelName string) error {
 	chat := gemini.StartChat()
 
 	r, err := chat.SendMessage(
-		context.TODO(),
+		ctx,
 		genai.Text("Hello"))
 	if err != nil {
 		return err
@@ -47,7 +48,7 @@ func makeChatRequests(projectId string, region string, modelName string) error {
 	fmt.Println(string(rb))
 
 	r, err = chat.SendMessage(
-		context.TODO(),
+		ctx,
 		genai.Text("What are all the colors in a rainbow?"))
 	if err != nil {
 		return err
@@ -56,7 +57,7 @@ func makeChatRequests(projectId string, region string, modelName string) error {
 	fmt.Println(string(rb))
 
 	r, err = chat.SendMessage(
-		context.TODO(),
+		ctx,
 		genai.Text("Why does it appear when it rains?"))
 	if err != nil {
 		return err

diff --git a/vertexai/snippets/try_gemini.go b/vertexai/snippets/try_gemini.go
@@ -28,15 +28,16 @@ var region = "us-central1"
 
 func tryGemini(w io.Writer, projectId string, region string, modelName string) error {
 
-	client, err := genai.NewClient(context.TODO(), projectId, region)
+	ctx := context.Background()
+	client, err := genai.NewClient(ctx, projectId, region)
 	gemini := client.GenerativeModel(modelName)
 
 	img := genai.FileData{
 		MIMEType: "image/jpeg",
 		FileURI:  "gs://generativeai-downloads/images/scones.jpg",
 	}
 	prompt := genai.Text("What is in this image?")
-	resp, err := gemini.GenerateContent(context.Background(), img, prompt)
+	resp, err := gemini.GenerateContent(ctx, img, prompt)
 	if err != nil {
 		return fmt.Errorf("error generating content: %w", err)
 	}