fixup: Address PR Comments

Switched to use colang config for endpoint and api key env var

Added onboarding steps from Trend's side
Expanded on documentation and examples

Author: trend-willem-gooderham

docs/user-guides/community/trend-micro.md

@@ -7,23 +7,21 @@ Trend Micro Vision One [AI Application Security's](https://docs.trendmicro.com/e
 - Sensitive Data
 
 
-The following environment variable is required to use the integration:
-
-- `V1_API_KEY`: A Vision One API Token with AI Guard Permissions
-
-You can optionally set:
-
-- `V1_URL`: The URL for which instances of AI Guard should be invoked
-  Defaults to `https://api.xdr.trendmicro.com/beta/aiSecurity/guard` for Vision One's hosted US SaaS deployment
-
 ## 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
@@ -36,6 +34,11 @@ rails:
 ```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

docs/user-guides/guardrails-library.md

@@ -918,7 +918,10 @@ For more details, check out the [Pangea AI Guard Integration](./community/pangea
 
 ### Trend Micro Vision One AI Application Security
 
-NeMo Guardrails supports using Trend Micro Vision One AI Guard for protecting input and output flows within AI-powered applications.
+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
 

examples/configs/trend_micro/config.yml

@@ -11,6 +11,9 @@ instructions:
       You are a helpful assistant.
 
 rails:
+  config:
+    trend_micro:
+      api_key_env_var: "V1_API_KEY"
   input:
     flows:
       - trend ai guard input

examples/configs/trend_micro_v2/config.yaml

@@ -2,6 +2,11 @@ colang_version: "2.x"
 
 enable_rails_exceptions: True
 
+rails:
+  config:
+    trend_micro:
+      api_key_env_var: "V1_API_KEY"
+
 models:
   - type: main
     engine: openai

examples/configs/trend_micro_v2/rails.co

@@ -2,7 +2,11 @@ import guardrails
 import nemoguardrails.library.trend_micro
 
 flow input rails $input_text
-    trend ai guard $input_text
+    $result = await TrendAiGuardAction(text=$input_text)
+
+    if $result.action == "Block"
+       send AiGuardException(message="AI Guard detection: " + $result.reason)
+       abort
 
 flow output rails $output_text
     trend ai guard $output_text

nemoguardrails/library/trend_micro/actions.py

@@ -14,14 +14,15 @@
 # limitations under the License.
 
 import logging
-import os
 from typing import Optional
 
 import httpx
 from pydantic import BaseModel
 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__)
 
@@ -35,17 +36,30 @@ class GuardResult(BaseModel):
     reason: str
 
 
+def get_config(config: RailsConfig) -> TrendMicroRailConfig:
+    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)
+
+
 @action(is_system_action=True)
-async def trend_ai_guard(text: Optional[str] = None):
+async def trend_ai_guard(config: RailsConfig, text: Optional[str] = None):
     """
     Custom action to invoke the Trend Ai Guard
     """
-    v1_url = os.environ.get(
-        "V1_URL", "https://api.xdr.trendmicro.com/beta/aiSecurity/guard"
-    )
-    v1_api_key = os.environ.get("V1_API_KEY")
+
+    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:
-        raise ValueError("V1_API_KEY environment variable is not set.")
+        raise ValueError("Trend Micro Vision One API Key not found")
 
     if text is None:
         raise ValueError("No prompt/response found in the last event.")

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."""

tests/test_trend_ai_guard.py

@@ -23,6 +23,10 @@
     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
@@ -32,6 +36,10 @@
     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
@@ -60,7 +68,7 @@ def test_trend_ai_guard_blocked(httpx_mock: HTTPXMock, monkeypatch: pytest.Monke
 
 
 @pytest.mark.unit
-@pytest.mark.parametrize("status_code", frozenset({429, 500, 502, 503, 504}))
+@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
 ):