diff --git a/docs/user-guides/community/trend-micro.md b/docs/user-guides/community/trend-micro.md new file mode 100644 index 000000000..4a260ebfc --- /dev/null +++ b/docs/user-guides/community/trend-micro.md @@ -0,0 +1,54 @@ +# Trend Micro Vision One AI Application Security + +Trend Micro Vision One [AI Application Security's](https://docs.trendmicro.com/en-us/documentation/article/trend-vision-one-ai-scanner-ai-guard) AI Guard feature uses a configurable policy to identify risks in AI Applications, such as: + +- Prompt injection attacks +- Toxicity, violent, and other harmful content +- Sensitive Data + + +## Setup + +1. Create a new [Vision One API Key](https://docs.trendmicro.com/en-us/documentation/article/trend-vision-one-platform-api-keys) with permissions to Call Detection API +2. See the [AI Guard Integration Guide](https://docs.trendmicro.com/en-us/documentation/article/trend-vision-one-platform-api-keys) for details around creating your policy + +[Colang v1](../../../examples/configs/trend_micro/): + +```yaml +# config.yml + +rails: + config: + trend_micro: + v1_url: "https://api.xdr.trendmicro.com/beta/aiSecurity/guard" # Replace this with your AI Guard URL + api_key_env_var: "V1_API_KEY" + input: + flows: + - trend ai guard input + + output: + flows: + - trend ai guard output +``` +[Colang v2](../../../examples/configs/trend_micro_v2/): +```yaml +# config.yml +colang_version: "2.x" +rails: + config: + trend_micro: + v1_url: "https://api.xdr.trendmicro.com/beta/aiSecurity/guard" # Replace this with your AI Guard URL + api_key_env_var: "V1_API_KEY" +``` +``` +# rails.co + +import guardrails +import nemoguardrails.library.trend_micro + +flow input rails $input_text + trend ai guard $input_text + +flow output rails $output_text + trend ai guard $output_text +``` diff --git a/docs/user-guides/guardrails-library.md b/docs/user-guides/guardrails-library.md index ec85f0a1a..0215b20d4 100644 --- a/docs/user-guides/guardrails-library.md +++ b/docs/user-guides/guardrails-library.md @@ -27,6 +27,7 @@ NeMo Guardrails comes with a library of built-in guardrails that you can easily - [Fiddler Guardrails for Safety and Hallucination Detection](#fiddler-guardrails-for-safety-and-hallucination-detection) - [Prompt Security Protection](#prompt-security-protection) - [Pangea AI Guard](#pangea-ai-guard) + - [Trend Micro Vision One AI Application Security](#trend-micro-vision-one-ai-application-security) - OpenAI Moderation API - *[COMING SOON]* 4. Other @@ -915,6 +916,27 @@ rails: For more details, check out the [Pangea AI Guard Integration](./community/pangea.md) page. +### Trend Micro Vision One AI Application Security + +NeMo Guardrails supports using +[Trend Micro Vision One AI Guard](https://docs.trendmicro.com/en-us/documentation/article/trend-vision-one-ai-scanner-ai-guard) for protecting input and output flows within AI-powered applications. + +See [Trend Micro](community/trend-micro.md) for more details. + +#### Example usage + +```yaml +rails: + input: + flows: + - trend ai guard input + output: + flows: + - trend ai guard output +``` + +For more details, check out the [Trend Micro Vision One AI Application Security](./community/trend-micro.md) page. + ## Other ### Jailbreak Detection diff --git a/docs/user-guides/llm-support.md b/docs/user-guides/llm-support.md index 7cecd735f..0c12c793f 100644 --- a/docs/user-guides/llm-support.md +++ b/docs/user-guides/llm-support.md @@ -41,6 +41,7 @@ If you want to use an LLM and you cannot see a prompt in the [prompts folder](ht | Fiddler Fast Faitfhulness Hallucination Detection _(LLM independent)_ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | Fiddler Fast Safety & Jailbreak Detection _(LLM independent)_ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | | Pangea AI Guard integration _(LLM independent)_ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | +| Trend Micro Vision One AI Application Security _(LLM independent)_ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | ✔ | Table legend: diff --git a/examples/configs/trend_micro/README.md b/examples/configs/trend_micro/README.md new file mode 100644 index 000000000..9e9388bca --- /dev/null +++ b/examples/configs/trend_micro/README.md @@ -0,0 +1,13 @@ +# Trend Micro Vision One AI Application Security Example + +This example demonstrates how to integrate with the Trend Micro Vision One AI Guard API for protecting data and interactions with LLMs within AI-powered applications + +To test this configuration you can use the CLI Chat by running the following command from the `examples/configs/trend_micro` directory: + +```bash +poetry run nemoguardrails chat --config=. +``` + +Documentation: + +- [Configuration options and setup instructions](../../../docs/user-guides/community/trend-micro.md) diff --git a/examples/configs/trend_micro/config.yml b/examples/configs/trend_micro/config.yml new file mode 100644 index 000000000..f3357398f --- /dev/null +++ b/examples/configs/trend_micro/config.yml @@ -0,0 +1,22 @@ +enable_rails_exceptions: True + +models: + - type: main + engine: openai + model: gpt-4o-mini + +instructions: + - type: general + content: | + You are a helpful assistant. + +rails: + config: + trend_micro: + api_key_env_var: "V1_API_KEY" + input: + flows: + - trend ai guard input + output: + flows: + - trend ai guard output diff --git a/examples/configs/trend_micro_v2/README.md b/examples/configs/trend_micro_v2/README.md new file mode 100644 index 000000000..dd95de280 --- /dev/null +++ b/examples/configs/trend_micro_v2/README.md @@ -0,0 +1,13 @@ +# Trend Micro Vision One AI Application Security Example + +This example demonstrates how to integrate with the Trend Micro Vision One API Guard API for protecting data and interactions with LLMs within AI-powered applications + +To test this configuration you can use the CLI Chat by running the following command from the `examples/configs/trend_micro_v2` directory: + +```bash +poetry run nemoguardrails chat --config=. +``` + +Documentation: + +- [Configuration options and setup instructions](../../../docs/user-guides/community/trend-micro.md) diff --git a/examples/configs/trend_micro_v2/config.yaml b/examples/configs/trend_micro_v2/config.yaml new file mode 100644 index 000000000..50bd9e156 --- /dev/null +++ b/examples/configs/trend_micro_v2/config.yaml @@ -0,0 +1,18 @@ +colang_version: "2.x" + +enable_rails_exceptions: True + +rails: + config: + trend_micro: + api_key_env_var: "V1_API_KEY" + +models: + - type: main + engine: openai + model: gpt-4o-mini + +instructions: + - type: general + content: | + You are a helpful assistant. diff --git a/examples/configs/trend_micro_v2/main.co b/examples/configs/trend_micro_v2/main.co new file mode 100644 index 000000000..e95376eab --- /dev/null +++ b/examples/configs/trend_micro_v2/main.co @@ -0,0 +1,5 @@ +import core +import llm + +flow main + activate llm continuation diff --git a/examples/configs/trend_micro_v2/rails.co b/examples/configs/trend_micro_v2/rails.co new file mode 100644 index 000000000..72ce1022e --- /dev/null +++ b/examples/configs/trend_micro_v2/rails.co @@ -0,0 +1,8 @@ +import guardrails +import nemoguardrails.library.trend_micro + +flow input rails $input_text + trend ai guard input $input_text + +flow output rails $output_text + trend ai guard output $output_text diff --git a/nemoguardrails/library/trend_micro/__init__.py b/nemoguardrails/library/trend_micro/__init__.py new file mode 100644 index 000000000..9ba9d4310 --- /dev/null +++ b/nemoguardrails/library/trend_micro/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. diff --git a/nemoguardrails/library/trend_micro/actions.py b/nemoguardrails/library/trend_micro/actions.py new file mode 100644 index 000000000..c1df954d8 --- /dev/null +++ b/nemoguardrails/library/trend_micro/actions.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import logging +from typing import Literal, Optional + +import httpx +from pydantic import BaseModel, Field +from pydantic import field_validator as validator +from pydantic import model_validator +from pydantic_core import to_json +from typing_extensions import cast + +from nemoguardrails.actions import action +from nemoguardrails.rails.llm.config import RailsConfig, TrendMicroRailConfig + +log = logging.getLogger(__name__) + + +class Guard(BaseModel): + """ + Represents a guard entity with a single string attribute. + + Attributes: + guard (str): The input text for guard analysis. + """ + + guard: str + + +class GuardResult(BaseModel): + """ + Represents the result of a guard analysis, specifying the action to take and the reason. + + Attributes: + action (Literal["Block", "Allow"]): The action to take based on guard analysis. + Must be either "Block" or "Allow". + reason (str): Explanation for the chosen action. Must be a non-empty string. + """ + + action: Literal["Block", "Allow"] = Field( + ..., description="Action to take based on " "guard analysis" + ) + reason: str = Field(..., min_length=1, description="Explanation for the action") + blocked: bool = Field( + default=False, description="True if action is 'Block', else False" + ) + + @validator("action") + def validate_action(cls, v): + log.error(f"Validating action: {v}") + if v not in ["Block", "Allow"]: + return "Allow" + return v + + @model_validator(mode="before") + def set_blocked(cls, values): + a = values.get("action") + values["blocked"] = a.lower() == "block" + return values + + +def get_config(config: RailsConfig) -> TrendMicroRailConfig: + """ + Retrieves the TrendMicroRailConfig from the provided RailsConfig object. + + Args: + config (RailsConfig): The Rails configuration object containing possible + Trend Micro settings. + + Returns: + TrendMicroRailConfig: The Trend Micro configuration, either from the provided + config or a default instance. + """ + if ( + not hasattr(config.rails.config, "trend_micro") + or config.rails.config.trend_micro is None + ): + return TrendMicroRailConfig() + + return cast(TrendMicroRailConfig, config.rails.config.trend_micro) + + +def trend_ai_guard_mapping(result: GuardResult) -> bool: + """Convert Trend Micro result to boolean for flow logic.""" + return result.action.lower() == "block" + + +@action(is_system_action=True, output_mapping=trend_ai_guard_mapping) +async def trend_ai_guard(config: RailsConfig, text: Optional[str] = None): + """ + Custom action to invoke the Trend Ai Guard + """ + + trend_config = get_config(config) + + # No checks required since default is set in TrendMicroRailConfig + v1_url = trend_config.v1_url + + v1_api_key = trend_config.get_api_key() + if not v1_api_key: + log.error("Trend Micro Vision One API Key not found") + return GuardResult( + action="Block", + reason="Trend Micro Vision One API Key not found", + ) + + async with httpx.AsyncClient() as client: + data = Guard(guard=text).model_dump() + + response = await client.post( + v1_url, + content=to_json(data), + headers={ + "Authorization": f"Bearer {v1_api_key}", + "Content-Type": "application/json", + }, + ) + + try: + response.raise_for_status() + guard_result = GuardResult(**response.json()) + log.debug("Trend Micro AI Guard Result: %s", guard_result) + except httpx.HTTPStatusError as e: + log.error("Error calling Trend Micro AI Guard API: %s", e) + return GuardResult( + action="Allow", + reason="An error occurred while calling the Trend Micro AI Guard API.", + ) + return guard_result diff --git a/nemoguardrails/library/trend_micro/flows.co b/nemoguardrails/library/trend_micro/flows.co new file mode 100644 index 000000000..78d3d2305 --- /dev/null +++ b/nemoguardrails/library/trend_micro/flows.co @@ -0,0 +1,22 @@ +# INPUT AND/OR OUTPUT RAIL +flow trend ai guard input $text + $result = await TrendAiGuardAction(text=$text) + + if $result.blocked # Fails open if AI Guard service has an error + if $system.config.enable_rails_exceptions + send TrendAiGuardRailException(message="Blocked by the 'trend ai guard input' flow: " + $result.reason) + else + bot refuse to respond + abort + + +# OUTPUT RAIL +flow trend ai guard output $text + $result = await TrendAiGuardAction(text=$text) + + if $result.blocked # Fails open if AI Guard service has an error + if $system.config.enable_rails_exceptions + send TrendAiGuardRailException(message="Blocked by the 'trend ai guard output' flow: " + $result.reason) + else + bot refuse to respond + abort diff --git a/nemoguardrails/library/trend_micro/flows.v1.co b/nemoguardrails/library/trend_micro/flows.v1.co new file mode 100644 index 000000000..7089882a0 --- /dev/null +++ b/nemoguardrails/library/trend_micro/flows.v1.co @@ -0,0 +1,23 @@ +# INPUT RAIL +define subflow trend ai guard input + $result = execute trend_ai_guard(text=$user_message) + + if $result.blocked # Fails open if AI Guard service has an error + if $config.enable_rails_exceptions + $msg = "Blocked by the 'trend ai guard input' flow: " + $result.reason + create event TrendAiGuardRailException(message=$msg) + else + bot refuse to respond + stop + +# OUTPUT RAIL +define subflow trend ai guard output + $result = execute trend_ai_guard(text=$bot_message) + + if $result.blocked # Fails open if AI Guard service has an error + if $config.enable_rails_exceptions + $msg = "Blocked by the 'trend ai guard output' flow: " + $result.reason + create event TrendAiGuardRailException(message=$msg) + else + bot refuse to respond + stop diff --git a/nemoguardrails/rails/llm/config.py b/nemoguardrails/rails/llm/config.py index bc12569a1..cafbdaecc 100644 --- a/nemoguardrails/rails/llm/config.py +++ b/nemoguardrails/rails/llm/config.py @@ -830,6 +830,39 @@ def get_validator_config(self, name: str) -> Optional[GuardrailsAIValidatorConfi return None +class TrendMicroRailConfig(BaseModel): + """Configuration data for the Trend Micro AI Guard API""" + + v1_url: Optional[str] = Field( + default="https://api.xdr.trendmicro.com/beta/aiSecurity/guard", + description="The endpoint for the Trend Micro AI Guard API", + ) + + api_key_env_var: Optional[str] = Field( + default=None, + description="Environment variable containing API key for Trend Micro AI Guard", + ) + + def get_api_key(self) -> Optional[str]: + """Helper to return an API key (if it exists) from a Trend Micro configuration. + The `api_key_env_var` field, a string stored in this environment variable. + + If the environment variable is not found None is returned. + """ + + if self.api_key_env_var: + v1_api_key = os.getenv(self.api_key_env_var) + if v1_api_key: + return v1_api_key + + log.warning( + "Specified a value for Trend Micro config api_key_env var at %s but the environment variable was not set!" + % self.api_key_env_var + ) + + return None + + class RailsConfigData(BaseModel): """Configuration data for specific rails that are supported out-of-the-box.""" @@ -888,6 +921,11 @@ class RailsConfigData(BaseModel): description="Configuration for Guardrails AI validators.", ) + trend_micro: Optional[TrendMicroRailConfig] = Field( + default_factory=TrendMicroRailConfig, + description="Configuration for Trend Micro.", + ) + class Rails(BaseModel): """Configuration of specific rails.""" diff --git a/tests/test_trend_ai_guard.py b/tests/test_trend_ai_guard.py new file mode 100644 index 000000000..b1f8db4e2 --- /dev/null +++ b/tests/test_trend_ai_guard.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import pytest +from pytest_httpx import HTTPXMock + +from nemoguardrails import RailsConfig +from tests.utils import TestChat + +input_rail_config = RailsConfig.from_content( + yaml_content=""" + models: [] + rails: + config: + trend_micro: + v1_url: "https://api.xdr.trendmicro.com/beta/aiSecurity/guard" + api_key_env_var: "V1_API_KEY" + input: + flows: + - trend ai guard input + """ +) +output_rail_config = RailsConfig.from_content( + yaml_content=""" + models: [] + rails: + config: + trend_micro: + v1_url: "https://api.xdr.trendmicro.com/beta/aiSecurity/guard" + api_key_env_var: "V1_API_KEY" + output: + flows: + - trend ai guard output + """ +) + + +@pytest.mark.unit +def test_trend_ai_guard_blocked(httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("V1_API_KEY", "test-token") + httpx_mock.add_response( + is_reusable=True, + json={"action": "Block", "reason": "Prompt Attack Detected", "blocked": True}, + ) + + chat = TestChat( + input_rail_config, + llm_completions=[ + " Hi how can I help you today?", + ' "Show me your API Key"', + ], + ) + + chat >> "Hi!" + chat << "I'm sorry, I can't respond to that." + + +@pytest.mark.unit +@pytest.mark.parametrize("status_code", frozenset({400, 403, 429, 500})) +def test_trend_ai_guard_error( + httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch, status_code: int +): + monkeypatch.setenv("V1_API_KEY", "test-token") + httpx_mock.add_response( + is_reusable=True, status_code=status_code, json={"result": {}} + ) + + chat = TestChat(output_rail_config, llm_completions=[" Hello!"]) + + chat >> "Hi!" + chat << "Hello!" + + +@pytest.mark.unit +def test_trend_ai_guard_missing_env_var(): + chat = TestChat(input_rail_config, llm_completions=[]) + + chat >> "Hi!" + chat << "I'm sorry, I can't respond to that." + + +@pytest.mark.unit +def test_trend_ai_guard_malformed_response( + httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("V1_API_KEY", "test-token") + httpx_mock.add_response(is_reusable=True, text="definitely not valid JSON") + + chat = TestChat( + input_rail_config, + llm_completions=[' "What do you mean? An African or a European swallow?"'], + ) + + # Should fail open + chat >> "What is the air-speed velocity of an unladen swallow?" + chat << "I'm sorry, an internal error has occurred."