From 31af5255ee3552a876e0bb5c8f3c2cf7e04814df Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Tue, 30 Dec 2025 09:45:40 -0700 Subject: [PATCH 1/9] :alembic: Simple experimental Nemo config for PII detection Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- nemo/pii_detect_config/__init__.py | 0 nemo/pii_detect_config/actions.py | 35 ++++++++++++++++++++++++++++++ nemo/pii_detect_config/config.yml | 13 +++++++++++ nemo/pii_detect_config/rails.co | 17 +++++++++++++++ 4 files changed, 65 insertions(+) create mode 100644 nemo/pii_detect_config/__init__.py create mode 100644 nemo/pii_detect_config/actions.py create mode 100644 nemo/pii_detect_config/config.yml create mode 100644 nemo/pii_detect_config/rails.co diff --git a/nemo/pii_detect_config/__init__.py b/nemo/pii_detect_config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nemo/pii_detect_config/actions.py b/nemo/pii_detect_config/actions.py new file mode 100644 index 0000000..5db5e8c --- /dev/null +++ b/nemo/pii_detect_config/actions.py @@ -0,0 +1,35 @@ +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, + } diff --git a/nemo/pii_detect_config/config.yml b/nemo/pii_detect_config/config.yml new file mode 100644 index 0000000..85b76b2 --- /dev/null +++ b/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 diff --git a/nemo/pii_detect_config/rails.co b/nemo/pii_detect_config/rails.co new file mode 100644 index 0000000..5b12ee8 --- /dev/null +++ b/nemo/pii_detect_config/rails.co @@ -0,0 +1,17 @@ +# 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" + # bot refuse just marks the response as refused instead of stopping + stop + else + # bot say "No PII!" + bot continue From ca30abb7fe262b66a65647d8f8cb99ebf1057708 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Wed, 31 Dec 2025 14:23:07 -0700 Subject: [PATCH 2/9] :alembic::construction: Initial CF nemo plugin experimentation Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- Dockerfile | 12 +++++++- nemo/pii_detect_config/__init__.py | 0 nemo/pii_detect_config/actions.py | 35 ---------------------- nemo/pii_detect_config/config.yml | 4 +-- nemo/pii_detect_config/rails.co | 5 +++- requirements.txt | 2 +- resources/config/config.yaml | 47 +++++++++++++++++++++--------- src/server.py | 47 ++++++++++++++++++++++++++++-- 8 files changed, 97 insertions(+), 55 deletions(-) delete mode 100644 nemo/pii_detect_config/__init__.py delete mode 100644 nemo/pii_detect_config/actions.py diff --git a/Dockerfile b/Dockerfile index a74d21a..473366a 100644 --- a/Dockerfile +++ b/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,6 +19,16 @@ RUN mkdir -p src/resources COPY src/ ./src/ COPY resources ./src/resources/ +RUN mkdir resources +RUN mkdir src +RUN mkdir contextforge-plugins-python +RUN mkdir nemo + +COPY src/ ./src/ +COPY resources ./resources/ +COPY plugins ./plugins/ +COPY contextforge-plugins-python ./contextforge-plugins-python/ +COPY nemo ./nemo/ WORKDIR /app/src # Expose the gRPC port diff --git a/nemo/pii_detect_config/__init__.py b/nemo/pii_detect_config/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/nemo/pii_detect_config/actions.py b/nemo/pii_detect_config/actions.py deleted file mode 100644 index 5db5e8c..0000000 --- a/nemo/pii_detect_config/actions.py +++ /dev/null @@ -1,35 +0,0 @@ -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, - } diff --git a/nemo/pii_detect_config/config.yml b/nemo/pii_detect_config/config.yml index 85b76b2..233e856 100644 --- a/nemo/pii_detect_config/config.yml +++ b/nemo/pii_detect_config/config.yml @@ -5,8 +5,8 @@ models: engine: ollama model: llama3.2:3b-instruct-fp16 parameters: - # Local ollama instance - base_url: http://localhost:11434 + # Docker ollama instance + base_url: http://host.docker.internal:11434 rails: input: flows: diff --git a/nemo/pii_detect_config/rails.co b/nemo/pii_detect_config/rails.co index 5b12ee8..654078c 100644 --- a/nemo/pii_detect_config/rails.co +++ b/nemo/pii_detect_config/rails.co @@ -14,4 +14,7 @@ define flow detect_pii_on_tool_input stop else # bot say "No PII!" - bot continue + bot finish + + execute debug_log(value=$pii.found, label="has_ended_flow_with_pii") + diff --git a/requirements.txt b/requirements.txt index fb992e9..7821bba 100644 --- a/requirements.txt +++ b/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 diff --git a/resources/config/config.yaml b/resources/config/config.yaml index 23c3ab8..ba8b6e4 100644 --- a/resources/config/config.yaml +++ b/resources/config/config.yaml @@ -1,24 +1,45 @@ # plugins/config.yaml - Main plugin configuration file plugins: - - name: CfDLGuard - kind: external - hooks: ["tool_pre_invoke", "tool_post_invoke"] - mode: enforce - mcp: - proto: STREAMABLEHTTP - url: http://dlguard-plugin-service:8000/mcp -# tls: -# certfile: /path/to/client-cert.pem -# keyfile: /path/to/client-key.pem -# ca_bundle: /path/to/ca-bundle.pem -# verify: false -# Self-contained Search Replace Plugin + - name: "ReplaceBadWordsPlugin" + kind: "plugins.examples.context_forge.internal.regex_filter.search_replace.SearchReplacePlugin" + # kind: "contextforge-plugins-python.regex_filter.search_replace.SearchReplacePlugin" + description: "A plugin for finding and replacing words." + version: "0.1.0" + author: "Teryl Taylor" + #hooks: ["prompt_pre_fetch", "prompt_post_fetch", "tool_pre_invoke", "tool_post_invoke"] + hooks: ["tool_pre_invoke"] + tags: ["plugin", "transformer", "regex", "search-and-replace", "pre-post"] + mode: "enforce" # enforce | permissive | disabled + priority: 150 + # conditions: + # # Apply to specific tools/servers + # - prompts: ["test_prompt"] + # server_ids: [] # Apply to all servers + # tenant_ids: [] # Apply to all tenants + config: + words: + - search: crap + replace: crud + - search: crud + replace: yikes + - name: "NemoWrapperPlugin" + kind: "plugins.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 - "/etc/mcpgateway/plugins" # System-wide plugins + - "plugins/nemo" # Example Nemo Guardrails plugins # Global plugin settings plugin_settings: diff --git a/src/server.py b/src/server.py index c78715e..c6e5e9f 100644 --- a/src/server.py +++ b/src/server.py @@ -100,12 +100,55 @@ async def getToolPreInvokeResponse(body): logger.debug(result) if not result.continue_processing: logger.debug("continue_processing false") + # Note: ImmediateResponse closes the stream, which may be rejected by MCP inspector? + error_body = { + "error": { + "code": http_status_pb2.Forbidden, + "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 but body has error + status=http_status_pb2.HttpStatus(code=http_status_pb2.OK), + headers=ep.HeaderMutation( + set_headers=[ + # content-type does not seem to make a difference + # core.HeaderValueOption( + # header=core.HeaderValue(key="content-type", value="application/x-ndjson") + # ), + core.HeaderValueOption( + header=core.HeaderValue(key="x-mcp-denied", value="true") + ), + ], + ), + # This will lead to invalid MCP request + # body=(json.dumps(error_body) + "\n").encode("utf-8") + details="No go - Tool args forbidden" + ) + ) + # Just status and header_mutation will work - but this necessitates that the server will have + # to process the header + body_resp = ep.ProcessingResponse( + request_body=ep.BodyResponse( + response=ep.CommonResponse( + status=ep.CommonResponse.CONTINUE_AND_REPLACE, + header_mutation=ep.HeaderMutation( + set_headers=[ + core.HeaderValueOption( + header=core.HeaderValue( + key="x-mcp-denied", + raw_value="true".encode( + "utf-8" + ), + ), + ) + ] + ), + ) ) ) + else: logger.debug("continue_processing true") result_payload = result.modified_payload From 69e1ada7e85db75a49d46ed97099655c3ce756dc Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 10:55:51 -0700 Subject: [PATCH 3/9] :truck: Move nemo plugin to plugins dir Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- plugins/nemo/nemo_wrapper_plugin.py | 76 ++++++++++++++++++++++ plugins/nemo/pii_detect_config/__init__.py | 0 plugins/nemo/pii_detect_config/actions.py | 35 ++++++++++ plugins/nemo/pii_detect_config/config.yml | 13 ++++ plugins/nemo/pii_detect_config/rails.co | 17 +++++ plugins/nemo/plugin-manifest.yaml | 11 ++++ 6 files changed, 152 insertions(+) create mode 100644 plugins/nemo/nemo_wrapper_plugin.py create mode 100644 plugins/nemo/pii_detect_config/__init__.py create mode 100644 plugins/nemo/pii_detect_config/actions.py create mode 100644 plugins/nemo/pii_detect_config/config.yml create mode 100644 plugins/nemo/pii_detect_config/rails.co create mode 100644 plugins/nemo/plugin-manifest.yaml diff --git a/plugins/nemo/nemo_wrapper_plugin.py b/plugins/nemo/nemo_wrapper_plugin.py new file mode 100644 index 0000000..09cf372 --- /dev/null +++ b/plugins/nemo/nemo_wrapper_plugin.py @@ -0,0 +1,76 @@ +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", "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 + if payload.args: + rails_response = await self._rails.generate_async(messages=[{"role": "user", "content": payload.args}]) + 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, + ] diff --git a/plugins/nemo/pii_detect_config/__init__.py b/plugins/nemo/pii_detect_config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plugins/nemo/pii_detect_config/actions.py b/plugins/nemo/pii_detect_config/actions.py new file mode 100644 index 0000000..5db5e8c --- /dev/null +++ b/plugins/nemo/pii_detect_config/actions.py @@ -0,0 +1,35 @@ +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, + } diff --git a/plugins/nemo/pii_detect_config/config.yml b/plugins/nemo/pii_detect_config/config.yml new file mode 100644 index 0000000..85b76b2 --- /dev/null +++ b/plugins/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 diff --git a/plugins/nemo/pii_detect_config/rails.co b/plugins/nemo/pii_detect_config/rails.co new file mode 100644 index 0000000..5b12ee8 --- /dev/null +++ b/plugins/nemo/pii_detect_config/rails.co @@ -0,0 +1,17 @@ +# 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" + # bot refuse just marks the response as refused instead of stopping + stop + else + # bot say "No PII!" + bot continue diff --git a/plugins/nemo/plugin-manifest.yaml b/plugins/nemo/plugin-manifest.yaml new file mode 100644 index 0000000..0292015 --- /dev/null +++ b/plugins/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 From c5030c008685da10167cb52503ec7f6f8c3ce390 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 10:56:21 -0700 Subject: [PATCH 4/9] :bug: Fix nemo rails to abort Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- plugins/nemo/pii_detect_config/rails.co | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/plugins/nemo/pii_detect_config/rails.co b/plugins/nemo/pii_detect_config/rails.co index 5b12ee8..1ff11c0 100644 --- a/plugins/nemo/pii_detect_config/rails.co +++ b/plugins/nemo/pii_detect_config/rails.co @@ -10,8 +10,6 @@ define flow detect_pii_on_tool_input if $pii.found bot say "PII detected, refusing to continue" - # bot refuse just marks the response as refused instead of stopping - stop else - # bot say "No PII!" - bot continue + bot say "No PII!" + abort From 091fdcec9f5b3faedf3d6cd19702ab1e90a96e49 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 12:40:54 -0700 Subject: [PATCH 5/9] :wrench::goal_net: Improve error handling for deny case Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Co-authored-by: mvle --- Dockerfile | 12 +--- ext-proc.yaml | 5 +- plugins/nemo/nemo_wrapper_plugin.py | 40 +++++++++--- plugins/nemo/pii_detect_config/actions.py | 1 + resources/config/config.yaml | 36 +++++------ src/server.py | 75 ++++++++--------------- 6 files changed, 75 insertions(+), 94 deletions(-) diff --git a/Dockerfile b/Dockerfile index 473366a..3334c3b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,20 +19,10 @@ RUN mkdir -p src/resources COPY src/ ./src/ COPY resources ./src/resources/ -RUN mkdir resources -RUN mkdir src -RUN mkdir contextforge-plugins-python -RUN mkdir nemo - -COPY src/ ./src/ -COPY resources ./resources/ COPY plugins ./plugins/ -COPY contextforge-plugins-python ./contextforge-plugins-python/ -COPY nemo ./nemo/ -WORKDIR /app/src # Expose the gRPC port EXPOSE 50052 # Run the server -CMD ["python", "server.py"] +CMD ["python", "src/server.py"] diff --git a/ext-proc.yaml b/ext-proc.yaml index 5a0cfd7..166b690 100644 --- a/ext-proc.yaml +++ b/ext-proc.yaml @@ -32,7 +32,7 @@ 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" @@ -40,9 +40,8 @@ spec: 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" - name: PLUGIN_MANAGER_CONFIG value: "./resources/config/config.yaml" - ports: - - containerPort: 50052 diff --git a/plugins/nemo/nemo_wrapper_plugin.py b/plugins/nemo/nemo_wrapper_plugin.py index 09cf372..7d6c897 100644 --- a/plugins/nemo/nemo_wrapper_plugin.py +++ b/plugins/nemo/nemo_wrapper_plugin.py @@ -1,3 +1,4 @@ +import asyncio import logging import os @@ -20,8 +21,8 @@ # 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): +class NemoWrapperPlugin(Plugin): def __init__(self, config: PluginConfig) -> None: """Initialize the plugin. @@ -29,13 +30,19 @@ def __init__(self, config: PluginConfig) -> None: config: Plugin configuration """ super().__init__(config) - logger.info(f"[NemoWrapperPlugin] Initializing plugin with config: {config.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", "nemo", "pii_detect_config")) + nemo_config = RailsConfig.from_path( + os.path.join(os.getcwd(), "plugins", "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: + async def tool_pre_invoke( + self, payload: ToolPreInvokePayload, context: PluginContext + ) -> ToolPreInvokeResult: """Plugin hook run before a tool is invoked. Args: @@ -47,16 +54,30 @@ async def tool_pre_invoke(self, payload: ToolPreInvokePayload, context: PluginCo """ # Very simple PII detection - attempt to block if any PII and does not alter the payload itself rails_response = None - if payload.args: - rails_response = await self._rails.generate_async(messages=[{"role": "user", "content": payload.args}]) + 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) + 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: + async def tool_post_invoke( + self, payload: ToolPostInvokePayload, context: PluginContext + ) -> ToolPostInvokeResult: """Plugin hook run after a tool is invoked. Args: @@ -67,7 +88,6 @@ async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: Plugin """ return ToolPostInvokeResult(modified_payload=payload) - def get_supported_hooks(self) -> list[str]: """Return list of supported hook types.""" return [ diff --git a/plugins/nemo/pii_detect_config/actions.py b/plugins/nemo/pii_detect_config/actions.py index 5db5e8c..aa9dc62 100644 --- a/plugins/nemo/pii_detect_config/actions.py +++ b/plugins/nemo/pii_detect_config/actions.py @@ -14,6 +14,7 @@ def debug_log(value, label: str = "DEBUG"): logger.warning(f"{label}: {value}") return value + @action() def detect_pii(tool_input): """ diff --git a/resources/config/config.yaml b/resources/config/config.yaml index ba8b6e4..aeec7ba 100644 --- a/resources/config/config.yaml +++ b/resources/config/config.yaml @@ -1,27 +1,19 @@ # plugins/config.yaml - Main plugin configuration file plugins: - - name: "ReplaceBadWordsPlugin" - kind: "plugins.examples.context_forge.internal.regex_filter.search_replace.SearchReplacePlugin" - # kind: "contextforge-plugins-python.regex_filter.search_replace.SearchReplacePlugin" - description: "A plugin for finding and replacing words." - version: "0.1.0" - author: "Teryl Taylor" - #hooks: ["prompt_pre_fetch", "prompt_post_fetch", "tool_pre_invoke", "tool_post_invoke"] - hooks: ["tool_pre_invoke"] - tags: ["plugin", "transformer", "regex", "search-and-replace", "pre-post"] - mode: "enforce" # enforce | permissive | disabled - priority: 150 - # conditions: - # # Apply to specific tools/servers - # - prompts: ["test_prompt"] - # server_ids: [] # Apply to all servers - # tenant_ids: [] # Apply to all tenants - config: - words: - - search: crap - replace: crud - - search: crud - replace: yikes + # Currently unavailable plugins will bring down the adapter +# - name: CfDLGuard +# kind: external +# hooks: ["tool_pre_invoke", "tool_post_invoke"] +# mode: enforce +# mcp: +# proto: STREAMABLEHTTP +# url: http://dlguard-plugin-service:8000/mcp +# # tls: +# # certfile: /path/to/client-cert.pem +# # keyfile: /path/to/client-key.pem +# # ca_bundle: /path/to/ca-bundle.pem +# # verify: false + # Nemo example - name: "NemoWrapperPlugin" kind: "plugins.nemo.nemo_wrapper_plugin.NemoWrapperPlugin" description: "A simple Nemo PII detector" diff --git a/src/server.py b/src/server.py index c6e5e9f..576d225 100644 --- a/src/server.py +++ b/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,68 +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") - # Note: ImmediateResponse closes the stream, which may be rejected by MCP inspector? error_body = { - "error": { - "code": http_status_pb2.Forbidden, - "message": "No go - Tool args forbidden" - } + "jsonrpc": body["jsonrpc"], + "id": body["id"], + "error": {"code": -32000, "message": "No go - Tool args forbidden"}, } body_resp = ep.ProcessingResponse( immediate_response=ep.ImmediateResponse( - # ok for stream but body has error - status=http_status_pb2.HttpStatus(code=http_status_pb2.OK), + # ok for stream, with error in body + status=http_status_pb2.HttpStatus(code=200), headers=ep.HeaderMutation( set_headers=[ - # content-type does not seem to make a difference - # core.HeaderValueOption( - # header=core.HeaderValue(key="content-type", value="application/x-ndjson") - # ), core.HeaderValueOption( - header=core.HeaderValue(key="x-mcp-denied", value="true") + 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") + ) ), ], ), - # This will lead to invalid MCP request - # body=(json.dumps(error_body) + "\n").encode("utf-8") - details="No go - Tool args forbidden" + body=(json.dumps(error_body)).encode("utf-8"), ) ) - # Just status and header_mutation will work - but this necessitates that the server will have - # to process the header - body_resp = ep.ProcessingResponse( - request_body=ep.BodyResponse( - response=ep.CommonResponse( - status=ep.CommonResponse.CONTINUE_AND_REPLACE, - header_mutation=ep.HeaderMutation( - set_headers=[ - core.HeaderValueOption( - header=core.HeaderValue( - key="x-mcp-denied", - raw_value="true".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") @@ -300,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( @@ -319,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)) From c119184dbad7882b8274ebccedccb6fdcb737ed4 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 12:54:15 -0700 Subject: [PATCH 6/9] :wrench: Update config Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- ext-proc.yaml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ext-proc.yaml b/ext-proc.yaml index 166b690..d01317c 100644 --- a/ext-proc.yaml +++ b/ext-proc.yaml @@ -42,6 +42,10 @@ spec: 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 \ No newline at end of file From 9827916619ff26f615a5a6a38f301df0df73eb9e Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 12:59:25 -0700 Subject: [PATCH 7/9] :truck: Move to examples dir Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- nemo/pii_detect_config/config.yml | 13 ------------ nemo/pii_detect_config/rails.co | 20 ------------------- .../nemo/nemo_wrapper_plugin.py | 2 +- .../nemo/pii_detect_config/__init__.py | 0 .../nemo/pii_detect_config/actions.py | 0 .../nemo/pii_detect_config/config.yml | 0 .../nemo/pii_detect_config/rails.co | 0 .../{ => examples}/nemo/plugin-manifest.yaml | 0 resources/config/config.yaml | 8 ++++---- 9 files changed, 5 insertions(+), 38 deletions(-) delete mode 100644 nemo/pii_detect_config/config.yml delete mode 100644 nemo/pii_detect_config/rails.co rename plugins/{ => examples}/nemo/nemo_wrapper_plugin.py (97%) rename plugins/{ => examples}/nemo/pii_detect_config/__init__.py (100%) rename plugins/{ => examples}/nemo/pii_detect_config/actions.py (100%) rename plugins/{ => examples}/nemo/pii_detect_config/config.yml (100%) rename plugins/{ => examples}/nemo/pii_detect_config/rails.co (100%) rename plugins/{ => examples}/nemo/plugin-manifest.yaml (100%) diff --git a/nemo/pii_detect_config/config.yml b/nemo/pii_detect_config/config.yml deleted file mode 100644 index 233e856..0000000 --- a/nemo/pii_detect_config/config.yml +++ /dev/null @@ -1,13 +0,0 @@ -id: pii-tool-test - -models: - - type: main - engine: ollama - model: llama3.2:3b-instruct-fp16 - parameters: - # Docker ollama instance - base_url: http://host.docker.internal:11434 -rails: - input: - flows: - - detect_pii_on_tool_input diff --git a/nemo/pii_detect_config/rails.co b/nemo/pii_detect_config/rails.co deleted file mode 100644 index 654078c..0000000 --- a/nemo/pii_detect_config/rails.co +++ /dev/null @@ -1,20 +0,0 @@ -# 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" - # bot refuse just marks the response as refused instead of stopping - stop - else - # bot say "No PII!" - bot finish - - execute debug_log(value=$pii.found, label="has_ended_flow_with_pii") - diff --git a/plugins/nemo/nemo_wrapper_plugin.py b/plugins/examples/nemo/nemo_wrapper_plugin.py similarity index 97% rename from plugins/nemo/nemo_wrapper_plugin.py rename to plugins/examples/nemo/nemo_wrapper_plugin.py index 7d6c897..8f2083c 100644 --- a/plugins/nemo/nemo_wrapper_plugin.py +++ b/plugins/examples/nemo/nemo_wrapper_plugin.py @@ -35,7 +35,7 @@ def __init__(self, config: PluginConfig) -> None: ) # NOTE: very hardcoded nemo_config = RailsConfig.from_path( - os.path.join(os.getcwd(), "plugins", "nemo", "pii_detect_config") + os.path.join(os.getcwd(), "plugins", "examples", "nemo", "pii_detect_config") ) self._rails = LLMRails(nemo_config) logger.info("[NemoWrapperPlugin] Plugin initialized successfully") diff --git a/plugins/nemo/pii_detect_config/__init__.py b/plugins/examples/nemo/pii_detect_config/__init__.py similarity index 100% rename from plugins/nemo/pii_detect_config/__init__.py rename to plugins/examples/nemo/pii_detect_config/__init__.py diff --git a/plugins/nemo/pii_detect_config/actions.py b/plugins/examples/nemo/pii_detect_config/actions.py similarity index 100% rename from plugins/nemo/pii_detect_config/actions.py rename to plugins/examples/nemo/pii_detect_config/actions.py diff --git a/plugins/nemo/pii_detect_config/config.yml b/plugins/examples/nemo/pii_detect_config/config.yml similarity index 100% rename from plugins/nemo/pii_detect_config/config.yml rename to plugins/examples/nemo/pii_detect_config/config.yml diff --git a/plugins/nemo/pii_detect_config/rails.co b/plugins/examples/nemo/pii_detect_config/rails.co similarity index 100% rename from plugins/nemo/pii_detect_config/rails.co rename to plugins/examples/nemo/pii_detect_config/rails.co diff --git a/plugins/nemo/plugin-manifest.yaml b/plugins/examples/nemo/plugin-manifest.yaml similarity index 100% rename from plugins/nemo/plugin-manifest.yaml rename to plugins/examples/nemo/plugin-manifest.yaml diff --git a/resources/config/config.yaml b/resources/config/config.yaml index 1adc9fd..03ff86d 100644 --- a/resources/config/config.yaml +++ b/resources/config/config.yaml @@ -23,7 +23,7 @@ plugins: replace: yikes # Nemo example - name: "NemoWrapperPlugin" - kind: "plugins.nemo.nemo_wrapper_plugin.NemoWrapperPlugin" + kind: "plugins.examples.nemo.nemo_wrapper_plugin.NemoWrapperPlugin" description: "A simple Nemo PII detector" version: "0.1.0" author: "Evaline Ju" @@ -36,10 +36,10 @@ plugins: # 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/nemo" # Example Nemo Guardrails plugins + - "plugins/examples/nemo" # Example Nemo guardrails plugins # Global plugin settings plugin_settings: From afa4fffa5b0648c878c62eaf58156209c1d13864 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:10:23 -0700 Subject: [PATCH 8/9] :memo: Document small description README Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- plugins/examples/nemo/README.md | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 plugins/examples/nemo/README.md diff --git a/plugins/examples/nemo/README.md b/plugins/examples/nemo/README.md new file mode 100644 index 0000000..deceffb --- /dev/null +++ b/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. From ee4419bad4b036220cfe528da319cf9abe4a4cc4 Mon Sep 17 00:00:00 2001 From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:18:22 -0700 Subject: [PATCH 9/9] :art: Lint plugin Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com> --- plugins/examples/nemo/nemo_wrapper_plugin.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/plugins/examples/nemo/nemo_wrapper_plugin.py b/plugins/examples/nemo/nemo_wrapper_plugin.py index 8f2083c..047c46f 100644 --- a/plugins/examples/nemo/nemo_wrapper_plugin.py +++ b/plugins/examples/nemo/nemo_wrapper_plugin.py @@ -35,7 +35,9 @@ def __init__(self, config: PluginConfig) -> None: ) # NOTE: very hardcoded nemo_config = RailsConfig.from_path( - os.path.join(os.getcwd(), "plugins", "examples", "nemo", "pii_detect_config") + os.path.join( + os.getcwd(), "plugins", "examples", "nemo", "pii_detect_config" + ) ) self._rails = LLMRails(nemo_config) logger.info("[NemoWrapperPlugin] Plugin initialized successfully")