✨ Nemo guardrails example

Dockerfile

@@ -3,7 +3,7 @@
 FROM public.ecr.aws/docker/library/python:3.12.12-slim
 
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends git \
+    && apt-get install -y --no-install-recommends git gcc g++ \
     && apt-get purge -y --auto-remove \
     && rm -rf /var/lib/apt/lists/*
 
@@ -19,10 +19,10 @@ RUN mkdir -p src/resources
 
 COPY src/ ./src/
 COPY resources ./src/resources/
+COPY plugins ./plugins/
 
-WORKDIR /app/src
 # Expose the gRPC port
 EXPOSE 50052
 
 # Run the server
-CMD ["python", "server.py"]
+CMD ["python", "src/server.py"]

ext-proc.yaml

@@ -32,17 +32,20 @@ spec:
           image: plugins-adapter:0.1.0
           # command: ["bash", "-c","--"]
           # args: ["while true; do sleep 3600; done"]
-#          imagePullPolicy: IfNotPresent
+          # imagePullPolicy: IfNotPresent
           env:
             - name: PLUGINS_SERVER_HOST
               value: "0.0.0.0"
             - name: LOGLEVEL
               value: "DEBUG"
             - name: PLUGINS_ENABLED
               value: "true"
+            # Note: The Dockerfile currently moves resources under ./src/resources
             - name: PLUGIN_CONFIG_FILE
-              value: "./resources/config/config.yaml"
+              value: "./src/resources/config/config.yaml"
             - name: PLUGIN_MANAGER_CONFIG
-              value: "./resources/config/config.yaml"
+              value: "./src/resources/config/config.yaml"
+            - name: PYTHONPATH
+              value: "./"
           ports:
-            - containerPort: 50052
+            - containerPort: 50052

plugins/examples/nemo/README.md

@@ -0,0 +1,9 @@
+# Nemo guardrails plugin example
+
+The `NemoWrapperPlugin` in `nemo_wrapper_plugin.py` currently invokes the simple flow in `pii_detect_config` which leverages an ollama model through `host.docker.internal`. The model can be easily replaced in the `config.yml`.
+
+How this works with the adapter:
+- The `NemoWrapperPlugin` is referenced in the plugin manager config (`resources/config/config.yaml` by default).
+- A plugins adapter image can be built with the nemoguardrails library using the `Dockerfile` in the repository.
+- The plugins adapter image can then be replaced in the `ext-proc.yaml` deployment. The Envoy filter `filter.yaml` makes sure the Envoy gateway request will pass through the ext-proc.
+- The MCP gateway can be brought up with `make inspect-gateway` or other methods. Test tool `test2_hello_world` can be used as a simple example to test PII/non-PII. As this is a simple example, there may be false positives.

plugins/examples/nemo/nemo_wrapper_plugin.py

@@ -0,0 +1,98 @@
+import asyncio
+import logging
+import os
+
+from mcpgateway.plugins.framework import (
+    Plugin,
+    PluginConfig,
+    PluginContext,
+    ToolHookType,
+    ToolPreInvokePayload,
+    ToolPreInvokeResult,
+    ToolPostInvokePayload,
+    ToolPostInvokeResult,
+)
+from nemoguardrails import LLMRails, RailsConfig
+
+logger = logging.getLogger(__name__)
+
+# NOTE: Ideally a plugin writer does not need to know what (MCP) primitive that the plugin has to run on
+# This uses the Context Forge plugin interface to inform how much effort it would be to adapt
+# nemo guardrails functionality and leverage this as a plugin server to be leveraged by the ext-proc
+# plugin adapter - currently as an internal plugin. The log levels are also particularly high for development currently
+
+
+class NemoWrapperPlugin(Plugin):
+    def __init__(self, config: PluginConfig) -> None:
+        """Initialize the plugin.
+
+        Args:
+            config: Plugin configuration
+        """
+        super().__init__(config)
+        logger.info(
+            f"[NemoWrapperPlugin] Initializing plugin with config: {config.config}"
+        )
+        # NOTE: very hardcoded
+        nemo_config = RailsConfig.from_path(
+            os.path.join(
+                os.getcwd(), "plugins", "examples", "nemo", "pii_detect_config"
+            )
+        )
+        self._rails = LLMRails(nemo_config)
+        logger.info("[NemoWrapperPlugin] Plugin initialized successfully")
+
+    async def tool_pre_invoke(
+        self, payload: ToolPreInvokePayload, context: PluginContext
+    ) -> ToolPreInvokeResult:
+        """Plugin hook run before a tool is invoked.
+
+        Args:
+            payload: The tool payload to be analyzed.
+            context: Contextual information about the hook call.
+
+        Returns:
+            The result of the plugin's analysis, including whether the tool can proceed.
+        """
+        # Very simple PII detection - attempt to block if any PII and does not alter the payload itself
+        rails_response = None
+        payload_args = payload.args
+        if payload_args:
+            try:
+                rails_response = await self._rails.generate_async(
+                    messages=[{"role": "user", "content": payload_args}]
+                )
+            except (
+                asyncio.CancelledError
+            ):  # asyncio.exceptions.CancelledError is thrown by nemo, need to catch
+                logging.exception("An error occurred in the nemo plugin except block:")
+            finally:
+                logger.warning("[NemoWrapperPlugin] Async rails executed")
+                logger.warning(rails_response)
+        if rails_response and "PII detected" in rails_response["content"]:
+            logger.warning("[NemoWrapperPlugin] PII detected, stopping processing")
+            return ToolPreInvokeResult(
+                modified_payload=payload, continue_processing=False
+            )
+        logger.warning("[NemoWrapperPlugin] No PII detected, continuing")
+        return ToolPreInvokeResult(modified_payload=payload, continue_processing=True)
+
+    async def tool_post_invoke(
+        self, payload: ToolPostInvokePayload, context: PluginContext
+    ) -> ToolPostInvokeResult:
+        """Plugin hook run after a tool is invoked.
+
+        Args:
+            payload: The tool result payload to be analyzed.
+            context: Contextual information about the hook call.
+        Returns:
+            The result of the plugin's analysis, including whether the tool result should proceed.
+        """
+        return ToolPostInvokeResult(modified_payload=payload)
+
+    def get_supported_hooks(self) -> list[str]:
+        """Return list of supported hook types."""
+        return [
+            ToolHookType.TOOL_PRE_INVOKE,
+            ToolHookType.TOOL_POST_INVOKE,
+        ]

plugins/examples/nemo/pii_detect_config/actions.py

@@ -0,0 +1,36 @@
+import logging
+import re
+from nemoguardrails.actions import action
+
+EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
+SSN_RE = re.compile(r"\b\d{3}-\d{2}-\d{4}\b")
+
+logger = logging.getLogger(__name__)
+
+
+@action()
+def debug_log(value, label: str = "DEBUG"):
+    # Purposely high logger level for now to make sure this is visible
+    logger.warning(f"{label}: {value}")
+    return value
+
+
+@action()
+def detect_pii(tool_input):
+    """
+    Simple PII detector to start
+    """
+    text = str(tool_input)
+
+    findings = []
+
+    if EMAIL_RE.search(text):
+        findings.append("email")
+
+    if SSN_RE.search(text):
+        findings.append("ssn")
+
+    return {
+        "found": len(findings) > 0,
+        "types": findings,
+    }

plugins/examples/nemo/pii_detect_config/config.yml

@@ -0,0 +1,13 @@
+id: pii-tool-test
+
+models:
+  - type: main
+    engine: ollama
+    model: llama3.2:3b-instruct-fp16
+    parameters:
+      # Local ollama instance
+      base_url: http://localhost:11434
+rails:
+  input:
+    flows:
+      - detect_pii_on_tool_input

plugins/examples/nemo/pii_detect_config/rails.co

@@ -0,0 +1,15 @@
+# Note: tool_call is w.r.t. LLM for nemo
+# Here, we happen to use the flow on tool input but the LLM itself is not making a tool call
+
+define flow detect_pii_on_tool_input
+    execute debug_log(value=$user_message, label="user_message")
+
+    $pii = execute detect_pii(tool_input=$user_message)
+
+    execute debug_log(value=$pii, label="has_pii")
+
+    if $pii.found
+        bot say "PII detected, refusing to continue"
+    else
+        bot say "No PII!"
+    abort

plugins/examples/nemo/plugin-manifest.yaml

@@ -0,0 +1,11 @@
+name: nemo_wrapper
+kind: nemo.plugin_external.NemoWrapperPlugin
+version: 1.0.0
+description: Attempted Nemo wrapper
+author: Evaline Ju
+priority: 10
+hooks:
+  - TOOL_PRE_INVOKE
+  - TOOL_POST_INVOKE
+config:
+  foo: bar

requirements.txt

@@ -3,4 +3,4 @@ grpcio-tools>=1.76.0
 betterproto2==0.9.1
 #mcp-contextforge-gateway==0.8.0
 git+https://github.com/IBM/mcp-context-forge@v0.9.0
-#nemoguardrails==0.17.0
+nemoguardrails==0.19.0

resources/config/config.yaml

@@ -1,6 +1,6 @@
 # plugins/config.yaml - Main plugin configuration file
 plugins:
-# Self-contained Search Replace Plugin
+  # Self-contained Search Replace Plugin - depends on plugin availability
   - name: "ReplaceBadWordsPlugin"
     kind: "plugins.regex_filter.search_replace.SearchReplacePlugin"
     description: "A plugin for finding and replacing words."
@@ -21,12 +21,25 @@ plugins:
           replace: crud
         - search: crud
           replace: yikes
+  # Nemo example
+  - name: "NemoWrapperPlugin"
+    kind: "plugins.examples.nemo.nemo_wrapper_plugin.NemoWrapperPlugin"
+    description: "A simple Nemo PII detector"
+    version: "0.1.0"
+    author: "Evaline Ju"
+    hooks: ["tool_pre_invoke", "tool_post_invoke"]
+    tags: ["plugin", "pre-post"]
+    mode: "enforce"  # enforce | permissive | disabled
+    priority: 150
+    config:
+      foo: bar
 
 # Plugin directories to scan
 plugin_dirs:
-  - "plugins/native"      # Built-in plugins
-  - "plugins/custom"      # Custom organization plugins
+  - "plugins/native"           # Built-in plugins
+  - "plugins/custom"           # Custom organization plugins
   - "/etc/mcpgateway/plugins"  # System-wide plugins
+  - "plugins/examples/nemo"    # Example Nemo guardrails plugins
 
 # Global plugin settings
 plugin_settings:

src/server.py

@@ -46,12 +46,12 @@ async def getToolPostInvokeResponse(body):
     # for content in body["result"]["content"]:
 
     logger.debug("**** Tool Post Invoke ****")
-    payload = ToolPostInvokePayload(name="replaceme", result=body)
+    payload = ToolPostInvokePayload(name="replaceme", result=body["result"])
     # TODO: hard-coded ids
-    logger.debug("**** Tool Post Invoke result ****")
-    logger.deub(payload)
+    logger.debug("**** Tool Post Invoke payload ****")
+    logger.debug(payload)
     global_context = GlobalContext(request_id="1", server_id="2")
-    result, contexts = await manager.invoke_hook(
+    result, _ = await manager.invoke_hook(
         ToolHookType.TOOL_POST_INVOKE, payload, global_context=global_context
     )
     logger.info(result)
@@ -66,7 +66,7 @@ async def getToolPostInvokeResponse(body):
     else:
         result_payload = result.modified_payload
         if result_payload is not None:
-            body = result_payload.result
+            body["result"] = result_payload.result
         else:
             body = None
         body_resp = ep.ProcessingResponse(
@@ -78,9 +78,11 @@ async def getToolPostInvokeResponse(body):
         )
     return body_resp
 
+
 def set_result_in_body(body, result_args):
     body["params"]["arguments"] = result_args
 
+
 async def getToolPreInvokeResponse(body):
     logger.debug(body)
     payload_args = {
@@ -93,25 +95,44 @@ async def getToolPreInvokeResponse(body):
     global_context = GlobalContext(request_id="1", server_id="2")
     logger.debug("**** Invoking Tool Pre Invoke with payload ****")
     logger.debug(payload)
-    result, contexts = await manager.invoke_hook(
+    result, _ = await manager.invoke_hook(
         ToolHookType.TOOL_PRE_INVOKE, payload, global_context=global_context
     )
     logger.debug("**** Tool Pre Invoke Result ****")
     logger.debug(result)
     if not result.continue_processing:
-        logger.debug("continue_processing false")
+        error_body = {
+            "jsonrpc": body["jsonrpc"],
+            "id": body["id"],
+            "error": {"code": -32000, "message": "No go - Tool args forbidden"},
+        }
         body_resp = ep.ProcessingResponse(
             immediate_response=ep.ImmediateResponse(
-                status=http_status_pb2.HttpStatus(code=http_status_pb2.Forbidden),
-                details="No go",
+                # ok for stream, with error in body
+                status=http_status_pb2.HttpStatus(code=200),
+                headers=ep.HeaderMutation(
+                    set_headers=[
+                        core.HeaderValueOption(
+                            header=core.HeaderValue(
+                                key="content-type",
+                                raw_value="application/json".encode("utf-8"),
+                            )
+                        ),
+                        core.HeaderValueOption(
+                            header=core.HeaderValue(
+                                key="x-mcp-denied", raw_value="True".encode("utf-8")
+                            )
+                        ),
+                    ],
+                ),
+                body=(json.dumps(error_body)).encode("utf-8"),
             )
         )
     else:
         logger.debug("continue_processing true")
         result_payload = result.modified_payload
-        if result_payload is not None and result_payload.get("tool_args", None) is not None:
-            logger.debug("changing tool call args")
-            set_result_in_body(body, result_payload.args)
+        if result_payload is not None and result_payload.args is not None:
+            body["params"]["arguments"] = result_payload.args["tool_args"]
         else:
             logger.debug("No change in tool args")
 
@@ -257,7 +278,7 @@ async def Process(
                         if data:  # List can be empty
                             data = json.loads(data[0].strip("data:"))
                         # TODO: check for tool call
-                        if "result" in data:
+                        if "result" in data and "content" in data["result"]:
                             body_resp = await getToolPostInvokeResponse(data)
                         else:
                             body_resp = ep.ProcessingResponse(
@@ -276,6 +297,7 @@ async def Process(
 async def serve(host: str = "0.0.0.0", port: int = 50052):
     await manager.initialize()
     logger.info(manager.config)
+    logger.debug(f"Loaded {manager.plugin_count} plugins")
 
     server = grpc.aio.server()
     # server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))