diff --git a/Dockerfile b/Dockerfile index a74d21a..3334c3b 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,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"] diff --git a/ext-proc.yaml b/ext-proc.yaml index 5a0cfd7..d01317c 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,12 @@ 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" + 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 \ No newline at end of file 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. diff --git a/plugins/examples/nemo/nemo_wrapper_plugin.py b/plugins/examples/nemo/nemo_wrapper_plugin.py new file mode 100644 index 0000000..047c46f --- /dev/null +++ b/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, + ] diff --git a/plugins/examples/nemo/pii_detect_config/__init__.py b/plugins/examples/nemo/pii_detect_config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plugins/examples/nemo/pii_detect_config/actions.py b/plugins/examples/nemo/pii_detect_config/actions.py new file mode 100644 index 0000000..aa9dc62 --- /dev/null +++ b/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, + } diff --git a/plugins/examples/nemo/pii_detect_config/config.yml b/plugins/examples/nemo/pii_detect_config/config.yml new file mode 100644 index 0000000..85b76b2 --- /dev/null +++ b/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 diff --git a/plugins/examples/nemo/pii_detect_config/rails.co b/plugins/examples/nemo/pii_detect_config/rails.co new file mode 100644 index 0000000..1ff11c0 --- /dev/null +++ b/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 diff --git a/plugins/examples/nemo/plugin-manifest.yaml b/plugins/examples/nemo/plugin-manifest.yaml new file mode 100644 index 0000000..0292015 --- /dev/null +++ b/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 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 812a5e7..03ff86d 100644 --- a/resources/config/config.yaml +++ b/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: diff --git a/src/server.py b/src/server.py index c78715e..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,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))