model armor call

model_armor.py

from google.cloud import modelarmor_v1
import google.auth
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()

def is_prompt_clean(prompt):
    """
        Take raw input and run it through the model armor API to check if it is clean
    """
    is_clean=False
    credentials,project_id=google.auth.default()
    # Create a client
    client = modelarmor_v1.ModelArmorClient(credentials=credentials,transport="rest", client_options = {"api_endpoint" : "modelarmor.us-central1.rep.googleapis.com"})    

    # Initialize request argument(s)
    user_prompt_data = modelarmor_v1.DataItem()
    user_prompt_data.text = prompt

    request = modelarmor_v1.SanitizeUserPromptRequest(
        name=f"projects/{project_id}/locations/us-central1/templates/model-armor-template",
        user_prompt_data=user_prompt_data,
    )

    # Make the request
    response = client.sanitize_user_prompt(request=request)
    if response.sanitization_result.filter_match_state == modelarmor_v1.FilterMatchState.MATCH_FOUND:
        logger.error(f"Prompt is NOT clean: {prompt}")
        is_clean = False
    else:
        logger.info("Prompt is clean")
        is_clean = True
    

    return is_clean