Skip to content

Commit e383a54

Browse files
tirthrajsinh-zala-crestharshnasitcrest
tirthrajsinh-zala-crest
and
harshnasitcrest
authored
feat(modelarmor): Added code samples for screening user prompt, model response and a PDF file (#10065)
* feat(modelarmor): Added sanitization code snippets * feat(modelarmor): Added the sanitization and PDF screening snippets * feat(modelarmor): refactor code * address-review-comments * Update sanitization methods to return response objects and update tests * increased-test-coverage * add-basic-sdp-test * add-advanced-sdp-test * revert-commented-code * added-sdp-template-cleanup-in-tests * updated-rai-tests-to-match-all-filters --------- Co-authored-by: Harsh Nasit <harsh.nasit@crestdata.ai>
1 parent 26355bd commit e383a54

File tree

5 files changed

+757
-17
lines changed

5 files changed

+757
-17
lines changed

‎modelarmor/src/main/java/modelarmor/SanitizeModelResponse.java

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package modelarmor;
18+
19+
// [START modelarmor_sanitize_model_response]
20+
21+
import com.google.cloud.modelarmor.v1.DataItem;
22+
import com.google.cloud.modelarmor.v1.ModelArmorClient;
23+
import com.google.cloud.modelarmor.v1.ModelArmorSettings;
24+
import com.google.cloud.modelarmor.v1.SanitizeModelResponseRequest;
25+
import com.google.cloud.modelarmor.v1.SanitizeModelResponseResponse;
26+
import com.google.cloud.modelarmor.v1.TemplateName;
27+
import com.google.protobuf.util.JsonFormat;
28+
import java.io.IOException;
29+
30+
public class SanitizeModelResponse {
31+
32+
public static void main(String[] args) throws IOException {
33+
// TODO(developer): Replace these variables before running the sample.
34+
35+
// Specify the Google Project ID.
36+
String projectId = "your-project-id";
37+
// Specify the location ID. For example, us-central1.
38+
String locationId = "your-location-id";
39+
// Specify the template ID.
40+
String templateId = "your-template-id";
41+
// Specify the model response.
42+
String modelResponse = "Unsanitized model output";
43+
44+
sanitizeModelResponse(projectId, locationId, templateId, modelResponse);
45+
}
46+
47+
public static SanitizeModelResponseResponse sanitizeModelResponse(String projectId,
48+
String locationId, String templateId, String modelResponse) throws IOException {
49+
50+
// Endpoint to call the Model Armor server.
51+
String apiEndpoint = String.format("modelarmor.%s.rep.googleapis.com:443", locationId);
52+
ModelArmorSettings modelArmorSettings = ModelArmorSettings.newBuilder().setEndpoint(apiEndpoint)
53+
.build();
54+
55+
try (ModelArmorClient client = ModelArmorClient.create(modelArmorSettings)) {
56+
// Build the resource name of the template.
57+
String name = TemplateName.of(projectId, locationId, templateId).toString();
58+
59+
// Prepare the request.
60+
SanitizeModelResponseRequest request =
61+
SanitizeModelResponseRequest.newBuilder()
62+
.setName(name)
63+
.setModelResponseData(
64+
DataItem.newBuilder().setText(modelResponse)
65+
.build())
66+
.build();
67+
68+
SanitizeModelResponseResponse response = client.sanitizeModelResponse(request);
69+
System.out.println("Result for the provided model response: "
70+
+ JsonFormat.printer().print(response.getSanitizationResult()));
71+
72+
return response;
73+
}
74+
}
75+
}
76+
// [END modelarmor_sanitize_model_response]

‎modelarmor/src/main/java/modelarmor/SanitizeUserPrompt.java

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package modelarmor;
18+
19+
// [START modelarmor_sanitize_user_prompt]
20+
21+
import com.google.cloud.modelarmor.v1.DataItem;
22+
import com.google.cloud.modelarmor.v1.ModelArmorClient;
23+
import com.google.cloud.modelarmor.v1.ModelArmorSettings;
24+
import com.google.cloud.modelarmor.v1.SanitizeUserPromptRequest;
25+
import com.google.cloud.modelarmor.v1.SanitizeUserPromptResponse;
26+
import com.google.cloud.modelarmor.v1.TemplateName;
27+
import com.google.protobuf.util.JsonFormat;
28+
import java.io.IOException;
29+
30+
public class SanitizeUserPrompt {
31+
32+
public static void main(String[] args) throws IOException {
33+
// TODO(developer): Replace these variables before running the sample.
34+
35+
// Specify the Google Project ID.
36+
String projectId = "your-project-id";
37+
// Specify the location ID. For example, us-central1.
38+
String locationId = "your-location-id";
39+
// Specify the template ID.
40+
String templateId = "your-template-id";
41+
// Specify the user prompt.
42+
String userPrompt = "Unsafe user prompt";
43+
44+
sanitizeUserPrompt(projectId, locationId, templateId, userPrompt);
45+
}
46+
47+
public static SanitizeUserPromptResponse sanitizeUserPrompt(String projectId, String locationId,
48+
String templateId, String userPrompt) throws IOException {
49+
50+
// Endpoint to call the Model Armor server.
51+
String apiEndpoint = String.format("modelarmor.%s.rep.googleapis.com:443", locationId);
52+
ModelArmorSettings modelArmorSettings = ModelArmorSettings.newBuilder()
53+
.setEndpoint(apiEndpoint)
54+
.build();
55+
56+
try (ModelArmorClient client = ModelArmorClient.create(modelArmorSettings)) {
57+
// Build the resource name of the template.
58+
String templateName = TemplateName.of(projectId, locationId, templateId).toString();
59+
60+
// Prepare the request.
61+
SanitizeUserPromptRequest request = SanitizeUserPromptRequest.newBuilder()
62+
.setName(templateName)
63+
.setUserPromptData(DataItem.newBuilder().setText(userPrompt).build())
64+
.build();
65+
66+
SanitizeUserPromptResponse response = client.sanitizeUserPrompt(request);
67+
System.out.println("Result for the provided user prompt: "
68+
+ JsonFormat.printer().print(response.getSanitizationResult()));
69+
70+
return response;
71+
}
72+
}
73+
}
74+
// [END modelarmor_sanitize_user_prompt]

‎modelarmor/src/main/java/modelarmor/ScreenPdfFile.java

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package modelarmor;
18+
19+
// [START modelarmor_screen_pdf_file]
20+
21+
import com.google.cloud.modelarmor.v1.ByteDataItem;
22+
import com.google.cloud.modelarmor.v1.ByteDataItem.ByteItemType;
23+
import com.google.cloud.modelarmor.v1.DataItem;
24+
import com.google.cloud.modelarmor.v1.ModelArmorClient;
25+
import com.google.cloud.modelarmor.v1.ModelArmorSettings;
26+
import com.google.cloud.modelarmor.v1.SanitizeUserPromptRequest;
27+
import com.google.cloud.modelarmor.v1.SanitizeUserPromptResponse;
28+
import com.google.cloud.modelarmor.v1.TemplateName;
29+
import com.google.protobuf.ByteString;
30+
import com.google.protobuf.util.JsonFormat;
31+
import java.io.IOException;
32+
import java.nio.file.Files;
33+
import java.nio.file.Paths;
34+
35+
public class ScreenPdfFile {
36+
37+
public static void main(String[] args) throws IOException {
38+
// TODO(developer): Replace these variables before running the sample.
39+
40+
// Specify the Google Project ID.
41+
String projectId = "your-project-id";
42+
// Specify the location ID. For example, us-central1.
43+
String locationId = "your-location-id";
44+
// Specify the template ID.
45+
String templateId = "your-template-id";
46+
// Specify the PDF file path. Replace with your PDF file path.
47+
String pdfFilePath = "src/main/resources/test_sample.pdf";
48+
49+
screenPdfFile(projectId, locationId, templateId, pdfFilePath);
50+
}
51+
52+
public static SanitizeUserPromptResponse screenPdfFile(String projectId, String locationId,
53+
String templateId, String pdfFilePath) throws IOException {
54+
55+
// Endpoint to call the Model Armor server.
56+
String apiEndpoint = String.format("modelarmor.%s.rep.googleapis.com:443", locationId);
57+
ModelArmorSettings modelArmorSettings = ModelArmorSettings.newBuilder().setEndpoint(apiEndpoint)
58+
.build();
59+
60+
try (ModelArmorClient client = ModelArmorClient.create(modelArmorSettings)) {
61+
// Build the resource name of the template.
62+
String name = TemplateName.of(projectId, locationId, templateId).toString();
63+
64+
// Read the PDF file content and encode it to Base64.
65+
byte[] fileContent = Files.readAllBytes(Paths.get(pdfFilePath));
66+
67+
// Prepare the request.
68+
DataItem userPromptData = DataItem.newBuilder()
69+
.setByteItem(
70+
ByteDataItem.newBuilder()
71+
.setByteDataType(ByteItemType.PDF)
72+
.setByteData(ByteString.copyFrom(fileContent))
73+
.build())
74+
.build();
75+
76+
SanitizeUserPromptRequest request =
77+
SanitizeUserPromptRequest.newBuilder()
78+
.setName(name)
79+
.setUserPromptData(userPromptData)
80+
.build();
81+
82+
// Send the request and get the response.
83+
SanitizeUserPromptResponse response = client.sanitizeUserPrompt(request);
84+
85+
// Print the sanitization result.
86+
System.out.println("Result for the provided PDF file: "
87+
+ JsonFormat.printer().print(response.getSanitizationResult()));
88+
89+
return response;
90+
}
91+
}
92+
}
93+
// [END modelarmor_screen_pdf_file]

‎modelarmor/src/main/resources/test_sample.pdf

26.4 KB
Binary file not shown.

0 commit comments

Comments
 (0)