Skip to content

feat(modelarmor): Added snippets for create template with labels and metadata #5273

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 23 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
d0b6026
feat(modelarmor): Added samples for creating, listing, updating and …
tirthrajsinh-zala-crest Apr 8, 2025
78110c9
feat(modelarmor): Added samples for creating, listing, updating and d…
tirthrajsinh-zala-crest Apr 9, 2025
a7b4cff
Merge branch 'main' into modelarmor-curd-snippets
tirthrajsinh-zala-crest Apr 9, 2025
4095b9a
feat(modelarmor): resolve review comments
tirthrajsinh-zala-crest Apr 9, 2025
caf8c7e
Merge branch 'main' into modelarmor-curd-snippets
tirthrajsinh-zala-crest Apr 9, 2025
1e5a58d
Merge branch 'main' into modelarmor-curd-snippets
tirthrajsinh-zala-crest Apr 10, 2025
cb7c703
feat(modelarmor): added samples for to create ma template with labels…
tirthrajsinh-zala-crest Apr 11, 2025
2f95b55
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 11, 2025
ca0ca05
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 14, 2025
cfdbe57
feat(modelarmor): resolve review comments and code refactored
tirthrajsinh-zala-crest Apr 14, 2025
a008169
feat(modelarmor): refactor code
tirthrajsinh-zala-crest Apr 14, 2025
cc25938
feat(modelarmor): code refactor
tirthrajsinh-zala-crest Apr 14, 2025
176cd99
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 17, 2025
0cb1b88
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 18, 2025
cec29ec
feat(modelarmor): resolve conflicts
tirthrajsinh-zala-crest Apr 18, 2025
2e82ed0
feat(modelarmor): resolve conflicts
tirthrajsinh-zala-crest Apr 18, 2025
5413312
feat(modelarmor): resolve conflicts
tirthrajsinh-zala-crest Apr 18, 2025
eb76c8a
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 23, 2025
6d95a8a
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
telpirion Apr 23, 2025
0239164
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
tirthrajsinh-zala-crest Apr 24, 2025
dcfb164
feat(modelarmor): resolve comments
tirthrajsinh-zala-crest Apr 24, 2025
abdee68
feat(modelarmor): resolve comments
tirthrajsinh-zala-crest Apr 25, 2025
ddafc31
Merge branch 'main' into modelarmor-create-template-with-labels-metadata
harshnasitcrest Apr 25, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 94 additions & 0 deletions modelarmor/create_template_with_labels.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// 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.

// Sample code for creating a new model armor template with labels.

package modelarmor

// [START modelarmor_create_template_with_labels]

import (
"context"
"fmt"
"io"

modelarmor "cloud.google.com/go/modelarmor/apiv1"
modelarmorpb "cloud.google.com/go/modelarmor/apiv1/modelarmorpb"
"google.golang.org/api/option"
)

// createModelArmorTemplateWithLabels method creates a
// new Model Armor template with custom labels.
//
// w io.Writer: The writer to use for logging.
// projectID string: The ID of the Google Cloud project.
// locationID string: The ID of the Google Cloud location.
// templateID string: The ID of the template to create.
// labels map[string]string: A map of custom labels to apply to the template.
func createModelArmorTemplateWithLabels(w io.Writer, projectID, locationID, templateID string, labels map[string]string) error {
ctx := context.Background()

// Create options for Model Armor client.
endpoint := fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID)
opts := option.WithEndpoint(endpoint)
// Create the Model Armor client.
client, err := modelarmor.NewClient(ctx,opts)
if err != nil {
return fmt.Errorf("failed to create client for project %s, location %s: %w", projectID, locationID, err)
}
defer client.Close()

parent := fmt.Sprintf("projects/%s/locations/%s", projectID, locationID)

// Build the Model Armor template with your preferred filters.
// For more details on filters, please refer to the following doc:
// [https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters](https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters)
template := &modelarmorpb.Template{
FilterConfig: &modelarmorpb.FilterConfig{
RaiSettings: &modelarmorpb.RaiFilterSettings{
RaiFilters: []*modelarmorpb.RaiFilterSettings_RaiFilter{
{
FilterType: modelarmorpb.RaiFilterType_HATE_SPEECH,
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_HIGH,
},
{
FilterType: modelarmorpb.RaiFilterType_SEXUALLY_EXPLICIT,
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_MEDIUM_AND_ABOVE,
},
},
},
},
Labels: labels,
}

// Prepare the request for creating the template.
req := &modelarmorpb.CreateTemplateRequest{
Parent: parent,
TemplateId: templateID,
Template: template,
}

// Create the template.
response, err := client.CreateTemplate(ctx, req)
if err != nil {
return fmt.Errorf("failed to create template: %w", err)
}

// Print the new template name using fmt.Fprintf with the io.Writer.
fmt.Fprintf(w, "Created Template with labels: %s\n", response.Name)

return nil
}

// [END modelarmor_create_template_with_labels]
100 changes: 100 additions & 0 deletions modelarmor/create_template_with_metadata.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
// 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.

// Sample code for creating a new model armor template with template metadata.

package modelarmor

// [START modelarmor_create_template_with_metadata]

import (
"context"
"fmt"
"io"

modelarmor "cloud.google.com/go/modelarmor/apiv1"
modelarmorpb "cloud.google.com/go/modelarmor/apiv1/modelarmorpb"
"google.golang.org/api/option"
)

// createModelArmorTemplateWithMetadata method creates a
// new Model Armor template with template metadata.
//
// w io.Writer: The writer to use for logging.
// projectID string: The ID of the Google Cloud project.
// locationID string: The ID of the Google Cloud location.
// templateID string: The ID of the template to create.
// metadata *modelarmorpb.TemplateMetadata: The template metadata to apply.
func createModelArmorTemplateWithMetadata(w io.Writer, projectID, locationID, templateID string) error {
ctx := context.Background()

// Create options for Model Armor client.
endpoint := fmt.Sprintf("modelarmor.%s.rep.googleapis.com:443", locationID)
opts := option.WithEndpoint(endpoint)
// Create the Model Armor client.
client, err := modelarmor.NewClient(ctx,opts)
if err != nil {
return fmt.Errorf("failed to create client for project %s, location %s: %w", projectID, locationID, err)
}
defer client.Close()

parent := fmt.Sprintf("projects/%s/locations/%s", projectID, locationID)

// Build the Model Armor template with your preferred filters.
// For more details on filters, please refer to the following doc:
// [https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters](https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters)
template := &modelarmorpb.Template{
FilterConfig: &modelarmorpb.FilterConfig{
RaiSettings: &modelarmorpb.RaiFilterSettings{
RaiFilters: []*modelarmorpb.RaiFilterSettings_RaiFilter{
{
FilterType: modelarmorpb.RaiFilterType_HATE_SPEECH,
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_HIGH,
},
{
FilterType: modelarmorpb.RaiFilterType_SEXUALLY_EXPLICIT,
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_MEDIUM_AND_ABOVE,
},
},
},
},
// Add template metadata to the template.
// For more details on template metadata, please refer to the following doc:
// [https://cloud.google.com/security-command-center/docs/reference/model-armor/rest/v1/projects.locations.templates#templatemetadata](https://cloud.google.com/security-command-center/docs/reference/model-armor/rest/v1/projects.locations.templates#templatemetadata)
TemplateMetadata: &modelarmorpb.Template_TemplateMetadata{
IgnorePartialInvocationFailures: true,
LogSanitizeOperations: true,
},
}

// Prepare the request for creating the template.
req := &modelarmorpb.CreateTemplateRequest{
Parent: parent,
TemplateId: templateID,
Template: template,
}

// Create the template.
response, err := client.CreateTemplate(ctx, req)
if err != nil {
return fmt.Errorf("failed to create template: %w", err)
}

// Print the new template name using fmt.Fprintf with the io.Writer.
fmt.Fprintf(w, "Created Model Armor Template: %s\n", response.Name)

return nil
}

// [END modelarmor_create_template_with_metadata]
83 changes: 77 additions & 6 deletions modelarmor/modelarmor_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,42 @@ func testClient(t *testing.T) (*modelarmor.Client, context.Context) {
return client, ctx
}

// testCleanupTemplate deletes the specified Model Armor template if it exists,
// ignoring the error if the template is already deleted.
// testModelArmorTemplate creates a new ModelArmor template for use in tests.
// It returns the created template or an error.
func testModelArmorTemplate(t *testing.T, templateID string) (*modelarmorpb.Template, error) {
t.Helper()
tc := testutil.SystemTest(t)
locationID := testLocation(t)
client, ctx := testClient(t)

template := &modelarmorpb.Template{
FilterConfig: &modelarmorpb.FilterConfig{
PiAndJailbreakFilterSettings: &modelarmorpb.PiAndJailbreakFilterSettings{
FilterEnforcement: modelarmorpb.PiAndJailbreakFilterSettings_ENABLED,
ConfidenceLevel: modelarmorpb.DetectionConfidenceLevel_MEDIUM_AND_ABOVE,
},
MaliciousUriFilterSettings: &modelarmorpb.MaliciousUriFilterSettings{
FilterEnforcement: modelarmorpb.MaliciousUriFilterSettings_ENABLED,
},
},
}

req := &modelarmorpb.CreateTemplateRequest{
Parent: fmt.Sprintf("projects/%s/locations/%s", tc.ProjectID, locationID),
TemplateId: templateID,
Template: template,
}

response, err := client.CreateTemplate(ctx, req)
if err != nil {
return nil, fmt.Errorf("failed to create template: %v", err)
}

return response, nil
}

// testCleanupTemplate deletes a given template by name if it exists.
// Ignores NotFound errors, which means the resource has already been deleted.
func testCleanupTemplate(t *testing.T, templateName string) {
t.Helper()

Expand Down Expand Up @@ -185,20 +219,57 @@ func TestCreateModelArmorTemplateWithAdvancedSDP(t *testing.T) {
// It ensures the output contains a confirmation message after creation.
func TestCreateModelArmorTemplate(t *testing.T) {
tc := testutil.SystemTest(t)

templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, testLocation(t), templateID)
var b bytes.Buffer
if err := createModelArmorTemplate(&b, tc.ProjectID, testLocation(t), templateID); err != nil {
var buf bytes.Buffer
if err := createModelArmorTemplate(&buf, tc.ProjectID, testLocation(t), templateID); err != nil {
t.Fatal(err)
}
defer testCleanupTemplate(t, templateName)

if got, want := b.String(), "Created template:"; !strings.Contains(got, want) {
if got, want := buf.String(), "Created template:"; !strings.Contains(got, want) {
t.Errorf("createModelArmorTemplate: expected %q to contain %q", got, want)
}
}

// TestCreateModelArmorTemplateWithMetadata tests the creation of a template with metadata.
// Verifies the success message is printed after template creation.
func TestCreateModelArmorTemplateWithMetadata(t *testing.T) {
tc := testutil.SystemTest(t)
locationID := testLocation(t)
templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)

var buf bytes.Buffer
if err := createModelArmorTemplateWithMetadata(&buf, tc.ProjectID, locationID, templateID); err != nil {
t.Fatal(err)
}
defer testCleanupTemplate(t, templateName)

if got, want := buf.String(), "Created Model Armor Template:"; !strings.Contains(got, want) {
t.Errorf("createModelArmorTemplateWithMetadata: expected %q to contain %q", got, want)
}
}

// TestCreateModelArmorTemplateWithLabels tests the creation of a template with labels.
// Verifies the output contains confirmation of successful template creation.
func TestCreateModelArmorTemplateWithLabels(t *testing.T) {
tc := testutil.SystemTest(t)
locationID := testLocation(t)
templateID := fmt.Sprintf("test-model-armor-%s", uuid.New().String())
templateName := fmt.Sprintf("projects/%s/locations/%s/templates/%s", tc.ProjectID, locationID, templateID)

var buf bytes.Buffer
if err := createModelArmorTemplateWithLabels(&buf, tc.ProjectID, locationID, templateID, map[string]string{"testkey": "testvalue"}); err != nil {
t.Fatal(err)
}
defer testCleanupTemplate(t, templateName)

if got, want := buf.String(), "Created Template with labels: "; !strings.Contains(got, want) {
t.Errorf("createModelArmorTemplateWithLabels: expected %q to contain %q", got, want)
}
}

// TestDeleteModelArmorTemplate verifies the deletion of a Model Armor template.
// It ensures the output contains a confirmation message after deletion.
func TestDeleteModelArmorTemplate(t *testing.T) {
Expand Down