Executing Input and Output Rails in Parallel

1
import os
2
from nemo_platform import NeMoPlatform, ConflictError
3
4
client = NeMoPlatform(
5
    base_url=os.environ.get("NMP_BASE_URL", "http://localhost:8080"),
6
    workspace="default",
7
)

1
guardrails_config = {
2
    "models": [
3
        {
4
            "type": "content_safety",
5
            "engine": "nim",
6
            "model": "default/nvidia-llama-3-1-nemotron-safety-guard-8b-v3",
7
        },
8
        {
9
            "type": "topic_control",
10
            "engine": "nim",
11
            "model": "default/nvidia-llama-3-1-nemoguard-8b-topic-control",
12
        },
13
    ],
14
    "rails": {
15
        "input": {
16
            "parallel": True,
17
            "flows": [
18
                "content safety check input $model=content_safety",
19
                "topic safety check input $model=topic_control",
20
            ],
21
        },
22
        "output": {
23
            "flows": [
24
                "content safety check output $model=content_safety",
25
            ],
26
        },
27
    },
28
    "prompts": [
29
        {
30
            "task": "content_safety_check_input $model=content_safety",
31
            "content": 'Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a \'safe\' or \'unsafe\' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.\n\n<BEGIN UNSAFE CONTENT CATEGORIES>\nS1: Violence.\nS2: Sexual.\nS3: Criminal Planning/Confessions.\nS4: Guns and Illegal Weapons.\nS5: Controlled/Regulated Substances.\nS6: Suicide and Self Harm.\nS7: Sexual (minor).\nS8: Hate/Identity Hate.\nS9: PII/Privacy.\nS10: Harassment.\nS11: Threat.\nS12: Profanity.\nS13: Needs Caution.\nS14: Other.\nS15: Manipulation.\nS16: Fraud/Deception.\nS17: Malware.\nS18: High Risk Gov Decision Making.\nS19: Political/Misinformation/Conspiracy.\nS20: Copyright/Trademark/Plagiarism.\nS21: Unauthorized Advice.\nS22: Illegal Activity.\nS23: Immoral/Unethical.\n<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\nuser: {{ user_input }}\n\n<END CONVERSATION>\n\nProvide your safety assessment for the above conversation in the below output JSON format:\n{"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}\n\nDo not include anything other than the output JSON in your response.\nOutput JSON:',
32
            "output_parser": "nemoguard_parse_prompt_safety",
33
            "max_tokens": 50,
34
        },
35
        {
36
            "task": "content_safety_check_output $model=content_safety",
37
            "content": 'Task: Check if there is unsafe content in the user or agent messages in the below conversation in accordance with our safety policy with the below categories. Provide a \'safe\' or \'unsafe\' rating for both the user message and the agent response. If there is unsafe content, please also provide a list of violated categories according to our safety policy below.\n\n<BEGIN UNSAFE CONTENT CATEGORIES>\nS1: Violence.\nS2: Sexual.\nS3: Criminal Planning/Confessions.\nS4: Guns and Illegal Weapons.\nS5: Controlled/Regulated Substances.\nS6: Suicide and Self Harm.\nS7: Sexual (minor).\nS8: Hate/Identity Hate.\nS9: PII/Privacy.\nS10: Harassment.\nS11: Threat.\nS12: Profanity.\nS13: Needs Caution.\nS14: Other.\nS15: Manipulation.\nS16: Fraud/Deception.\nS17: Malware.\nS18: High Risk Gov Decision Making.\nS19: Political/Misinformation/Conspiracy.\nS20: Copyright/Trademark/Plagiarism.\nS21: Unauthorized Advice.\nS22: Illegal Activity.\nS23: Immoral/Unethical.\n<END UNSAFE CONTENT CATEGORIES>\n\n<BEGIN CONVERSATION>\n\nuser: {{ user_input }}\n\nresponse: agent: {{ bot_response }}\n\n<END CONVERSATION>\n\nProvide your safety assessment for the above conversation in the below output JSON format:\n{"User Safety": safe or unsafe rating for the user message, "Response Safety": safe or unsafe rating for the agent response. Omit if no agent response present. "Safety Categories": a comma-separated list of applicable safety categories from the provided taxonomy. Omit if all safe.}\n\nDo not include anything other than the output JSON in your response.\nOutput JSON:',
38
            "output_parser": "nemoguard_parse_response_safety",
39
            "max_tokens": 50,
40
        },
41
        {
42
            "task": "topic_safety_check_input $model=topic_control",
43
            "content": "You are to act as a customer service agent, providing users with factual information in accordance to the knowledge base. Your role is to ensure that you respond only to relevant queries and adhere to the following guidelines\n\nGuidelines for the user messages:\n- Do not answer questions related to personal opinions or advice on user's order, future recommendations\n- Do not provide any information on non-company products or services.\n- Do not answer enquiries unrelated to the company policies.\n- Do not answer questions asking for personal details about the agent or its creators.\n- Do not answer questions about sensitive topics related to politics, religion, or other sensitive subjects.\n- If a user asks topics irrelevant to the company's customer service relations, politely redirect the conversation or end the interaction.\n- Your responses should be professional, accurate, and compliant with customer relations guidelines, focusing solely on providing transparent, up-to-date information about the company that is already publicly available.\n- allow user comments that are related to small talk and chit-chat.\n\nUser message: \"{{ user_input }}\"",
44
            "max_tokens": 50,
45
        },
46
    ],
47
}
48
49
config_name = "parallel-rails-config"
50
try:
51
    config = client.guardrail.configs.create(
52
        name=config_name,
53
        description="Parallel rails guardrail configuration",
54
        data=guardrails_config,
55
    )
56
except ConflictError:
57
    print(f"Config {config_name} already exists, continuing...")

1
oai_client = client.models.get_openai_client()
2
3
response = oai_client.chat.completions.create(
4
    model="default/guarded-parallel-rails",
5
    messages=[{"role": "user", "content": "What is your return policy?"}],
6
    max_tokens=200,
7
)
8
9
print(response.model_dump_json(indent=2))

1
response = oai_client.chat.completions.create(
2
    model="default/guarded-parallel-rails",
3
    messages=[{"role": "user", "content": "Tell me a joke about quantum gravity."}],
4
    max_tokens=200,
5
)
6
7
print(response.model_dump_json(indent=2))

1
response = oai_client.chat.completions.create(
2
    model="default/guarded-parallel-rails",
3
    messages=[{"role": "user", "content": "What is your return policy?"}],
4
    max_tokens=200,
5
    extra_body={"guardrails": {"options": {"log": {"activated_rails": True}}}},
6
)
7
8
print(response.model_dump_json(indent=2))

1
client.inference.virtual_models.delete(name="guarded-parallel-rails")
2
client.guardrail.configs.delete(name=config_name)
3
print("Cleanup complete")