Skip to content
Open
99 changes: 99 additions & 0 deletions genai/batch_prediction/batchpredict_embeddings_with_gcs.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
// Copyright 2025 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
//
// https://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 batch_prediction shows how to use the GenAI SDK to batch prediction.

package batch_prediction

// [START googlegenaisdk_batchpredict_embeddings_with_gcs]
import (
"context"
"fmt"
"io"
"time"

"google.golang.org/genai"
)

// generateBatchEmbeddings shows how to run a batch embeddings prediction job.
func generateBatchEmbeddings(w io.Writer, outputURI string) error {
// outputURI = "gs://your-bucket/your-prefix"
ctx := context.Background()

client, err := genai.NewClient(ctx, &genai.ClientConfig{
HTTPOptions: genai.HTTPOptions{APIVersion: "v1"},
})
if err != nil {
return fmt.Errorf("failed to create genai client: %w", err)
}
modelName := "text-embedding-005"
// See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.batches.Batches.create
job, err := client.Batches.Create(ctx,
modelName,
&genai.BatchJobSource{
Format: "jsonl",
// Source link: https://storage.cloud.google.com/cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl
GCSURI: []string{"gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl"},
},
&genai.CreateBatchJobConfig{
Dest: &genai.BatchJobDestination{
Format: "jsonl",
GCSURI: outputURI,
},
},
)
if err != nil {
return fmt.Errorf("failed to create batch job: %w", err)
}

fmt.Fprintf(w, "Job name: %s\n", job.Name)
fmt.Fprintf(w, "Job state: %s\n", job.State)
// Example response:
// Job name: projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/9876453210000000000
// Job state: JOB_STATE_PENDING

// See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
completedStates := map[genai.JobState]bool{
genai.JobStateSucceeded: true,
genai.JobStateFailed: true,
genai.JobStateCancelled: true,
genai.JobStatePaused: true,
}

// Poll until job finishes
for !completedStates[job.State] {
time.Sleep(30 * time.Second)
job, err = client.Batches.Get(ctx, job.Name, nil)
if err != nil {
return fmt.Errorf("failed to get batch job: %w", err)
}
fmt.Fprintf(w, "Job state: %s\n", job.State)

if job.State == genai.JobStateFailed {
fmt.Fprintf(w, "Error: %+v\n", job.Error)
break
}
}

// Example response:
// Job state: JOB_STATE_PENDING
// Job state: JOB_STATE_RUNNING
// Job state: JOB_STATE_RUNNING
// ...
// Job state: JOB_STATE_SUCCEEDED

return nil
}

// [END googlegenaisdk_batchpredict_embeddings_with_gcs]
96 changes: 96 additions & 0 deletions genai/batch_prediction/batchpredict_with_bq.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Copyright 2025 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
//
// https://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 batch_prediction shows how to use the GenAI SDK to batch prediction.

package batch_prediction

// [START googlegenaisdk_batchpredict_with_bq]
import (
"context"
"fmt"
"io"
"time"

"google.golang.org/genai"
)

// generateBatchPredictWithBQ shows how to run a batch prediction job with BigQuery input/output.
func generateBatchPredictWithBQ(w io.Writer, outputURI string) error {
// outputURI = "bq://your-project.your_dataset.your_table"
ctx := context.Background()

client, err := genai.NewClient(ctx, &genai.ClientConfig{
HTTPOptions: genai.HTTPOptions{APIVersion: "v1"},
})
if err != nil {
return fmt.Errorf("failed to create genai client: %w", err)
}

// BigQuery input
src := &genai.BatchJobSource{
Format: "bigquery",
BigqueryURI: "bq://storage-samples.generative_ai.batch_requests_for_multimodal_input",
}

// BigQuery output
config := &genai.CreateBatchJobConfig{
Dest: &genai.BatchJobDestination{
Format: "bigquery",
BigqueryURI: outputURI,
},
}

// To use a tuned model, set the model param to your tuned model using the following format:
// modelName:= "projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}
modelName := "gemini-2.5-flash"
job, err := client.Batches.Create(ctx, modelName, src, config)
if err != nil {
return fmt.Errorf("failed to create batch job: %w", err)
}

fmt.Fprintf(w, "Job name: %s\n", job.Name)
fmt.Fprintf(w, "Job state: %s\n", job.State)
// Example response:
// Job name: projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/9876453210000000000
// Job state: JOB_STATE_PENDING

// See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
completedStates := map[genai.JobState]bool{
genai.JobStateSucceeded: true,
genai.JobStateFailed: true,
genai.JobStateCancelled: true,
genai.JobStatePaused: true,
}

for !completedStates[job.State] {
time.Sleep(30 * time.Second)
job, err = client.Batches.Get(ctx, job.Name, nil)
if err != nil {
return fmt.Errorf("failed to get batch job: %w", err)
}
fmt.Fprintf(w, "Job state: %s\n", job.State)
}

// Example response:
// Job state: JOB_STATE_PENDING
// Job state: JOB_STATE_RUNNING
// Job state: JOB_STATE_RUNNING
// ...
// Job state: JOB_STATE_SUCCEEDED

return nil
}

// [END googlegenaisdk_batchpredict_with_bq]
99 changes: 99 additions & 0 deletions genai/batch_prediction/batchpredict_with_gcs.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
// Copyright 2025 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
//
// https://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 batch_prediction shows how to use the GenAI SDK to batch prediction.

package batch_prediction

// [START googlegenaisdk_batchpredict_with_gcs]
import (
"context"
"fmt"
"io"
"time"

"google.golang.org/genai"
)

// generateBatchPredict runs a batch prediction job using GCS input/output.
func generateBatchPredict(w io.Writer, outputURI string) error {
// outputURI = "gs://your-bucket/your-prefix"
ctx := context.Background()

client, err := genai.NewClient(ctx, &genai.ClientConfig{
HTTPOptions: genai.HTTPOptions{APIVersion: "v1"},
})
if err != nil {
return fmt.Errorf("failed to create genai client: %w", err)
}

// Source file with prompts for prediction
src := &genai.BatchJobSource{
Format: "jsonl",
// Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
GCSURI: []string{"gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl"},
}

// Batch job config with output GCS location
config := &genai.CreateBatchJobConfig{
Dest: &genai.BatchJobDestination{
Format: "jsonl",
GCSURI: outputURI,
},
}
// To use a tuned model, set the model param to your tuned model using the following format:
// modelName:= "projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}
modelName := "gemini-2.5-flash"
// See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.batches.Batches.create
job, err := client.Batches.Create(ctx, modelName, src, config)
if err != nil {
return fmt.Errorf("failed to create batch job: %w", err)
}

fmt.Fprintf(w, "Job name: %s\n", job.Name)
fmt.Fprintf(w, "Job state: %s\n", job.State)
// Example response:
// Job name: projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/9876453210000000000
// Job state: JOB_STATE_PENDING

// See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
completedStates := map[genai.JobState]bool{
genai.JobStateSucceeded: true,
genai.JobStateFailed: true,
genai.JobStateCancelled: true,
genai.JobStatePaused: true,
}

for !completedStates[job.State] {
time.Sleep(30 * time.Second)

job, err = client.Batches.Get(ctx, job.Name, nil)
if err != nil {
return fmt.Errorf("failed to get batch job: %w", err)
}

fmt.Fprintf(w, "Job state: %s\n", job.State)
}

// Example response:
// Job state: JOB_STATE_PENDING
// Job state: JOB_STATE_RUNNING
// Job state: JOB_STATE_RUNNING
// ...
// Job state: JOB_STATE_SUCCEEDED

return nil
}

// [END googlegenaisdk_batchpredict_with_gcs]
80 changes: 80 additions & 0 deletions genai/batch_prediction/batchprediction_examples_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
// Copyright 2025 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 batch_prediction

import (
"bytes"
"fmt"
"testing"
"time"

"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
)

const gcsOutputBucket = "golang-docs-samples-tests"

func TestBatchPrediction(t *testing.T) {
tc := testutil.SystemTest(t)

t.Setenv("GOOGLE_GENAI_USE_VERTEXAI", "1")
t.Setenv("GOOGLE_CLOUD_LOCATION", "us-central1")
t.Setenv("GOOGLE_CLOUD_PROJECT", tc.ProjectID)

prefix := fmt.Sprintf("embeddings_output/%d", time.Now().UnixNano())
outputURI := fmt.Sprintf("gs://%s/%s", gcsOutputBucket, prefix)
buf := new(bytes.Buffer)

t.Run("generate batch embeddings with GCS", func(t *testing.T) {
buf.Reset()
err := generateBatchEmbeddings(buf, outputURI)
if err != nil {
t.Fatalf("generateBatchEmbeddings failed: %v", err)
}

output := buf.String()
if output == "" {
t.Error("expected non-empty output, got empty")
}
})

t.Run("generate batch predict with gcs input/output", func(t *testing.T) {
buf.Reset()
err := generateBatchPredict(buf, outputURI)
if err != nil {
t.Fatalf("generateBatchPredict failed: %v", err)
}

output := buf.String()
if output == "" {
t.Error("expected non-empty output, got empty")
}

})

t.Run("generate batch predict with BigQuery", func(t *testing.T) {
buf.Reset()
outputURIBQ := "bq://your-project.your_dataset.your_table"

err := generateBatchPredictWithBQ(buf, outputURIBQ)
if err != nil {
t.Fatalf("generateBatchPredictWithBQ failed: %v", err)
}

output := buf.String()
if output == "" {
t.Error("expected non-empty output, got empty")
}
})
}
Loading