feat: ✨ Handle tool response checks with nemo guardrails check

.github/workflows/ci.yaml

@@ -42,6 +42,8 @@ jobs:
             git diff --stat
             exit 1
           }
+
+      # Plugin tests
       - name: Install nemocheck plugin dependencies
         working-directory: ./plugins/examples/nemocheck
         run: |
@@ -50,3 +52,13 @@ jobs:
       - name: Run nemocheck plugin tests
         working-directory: ./plugins/examples/nemocheck
         run: uv run pytest tests
+
+      # Server tests
+      - name: Install server test dependencies
+        run: |
+          echo "Installing pytest and dependencies for server tests..."
+          pip install pytest pytest-asyncio
+      - name: Run server unit tests
+        run: |
+          echo "Running server unit tests..."
+          uv run pytest tests

.pre-commit-config.yaml

@@ -4,6 +4,7 @@ repos:
   hooks:
     # Run the linter.
     - id: ruff
-      args: [ --fix ]
+      args: [ --fix, --line-length=80 ]
     # Run the formatter.
     - id: ruff-format
+      args: [ --line-length=80 ]

plugins/examples/nemo/nemo_wrapper_plugin.py

@@ -62,20 +62,24 @@ async def tool_pre_invoke(
                 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:")
+            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")
+            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)
+        return ToolPreInvokeResult(
+            modified_payload=payload, continue_processing=True
+        )
 
     async def tool_post_invoke(
         self, payload: ToolPostInvokePayload, context: PluginContext

plugins/examples/nemocheck/plugin.py

@@ -62,7 +62,9 @@ def __init__(self, config: PluginConfig):
             )
         else:
             self.check_endpoint = DEFAULT_CHECK_ENDPOINT
-            logger.warning("Plugin config is empty or invalid, using default endpoint")
+            logger.warning(
+                "Plugin config is empty or invalid, using default endpoint"
+            )
         logger.info(f"Nemo Check endpoint: {self.check_endpoint}")
 
     async def prompt_pre_fetch(
@@ -105,9 +107,11 @@ async def tool_pre_invoke(
         Returns:
             The result of the plugin's analysis, including whether the tool can proceed.
         """
-        logger.info("tool_pre_invoke....")
-        logger.info(payload)
-        tool_name = payload.name  # ("tool_name", None)
+        logger.info(
+            f"[NemoCheck] Starting tool pre invoke hook with payload {payload}"
+        )
+
+        tool_name = payload.name
         check_nemo_payload = {
             "model": MODEL_NAME,
             "messages": [
@@ -119,7 +123,9 @@ async def tool_pre_invoke(
                             "type": "function",
                             "function": {
                                 "name": tool_name,
-                                "arguments": payload.args.get("tool_args", None),
+                                "arguments": payload.args.get(
+                                    "tool_args", None
+                                ),
                             },
                         }
                     ],
@@ -135,7 +141,7 @@ async def tool_pre_invoke(
             if response.status_code == 200:
                 data = response.json()
                 status = data.get("status", "blocked")
-                logger.debug(f"rails reply: {data}")
+                logger.debug(f"[NemoCheck] Rails reply: {data}")
 
                 if status == "success":
                     metadata = data.get("rails_status")
@@ -147,7 +153,7 @@ async def tool_pre_invoke(
                     violation = PluginViolation(
                         reason=f"Check tool rails:{status}.",
                         description=json.dumps(data),
-                        code=f"checkserver_http_status_code:{response.status_code}",
+                        code="NEMO_RAILS_BLOCKED",
                         details=metadata,
                     )
                     return ToolPreInvokeResult(
@@ -158,23 +164,25 @@ async def tool_pre_invoke(
             else:
                 violation = PluginViolation(
                     reason="Tool Check Unavailable",
-                    description="Tool arguments check server returned error",
-                    code=f"checkserver_http_status_code:{response.status_code}",
-                    details={},
+                    description=f"Tool arguments check server returned error. Status code: {response.status_code}, Response: {response.text}",
+                    code="NEMO_SERVER_ERROR",
+                    details={"status_code": response.status_code},
                 )
                 return ToolPreInvokeResult(
                     continue_processing=False, violation=violation
                 )
 
         except Exception as e:
-            logger.error(f"Error calling Nemo Check endpoint: {e}")
+            logger.error(f"[NemoCheck] Error checking tool arguments: {e}")
             violation = PluginViolation(
                 reason="Tool Check Error",
                 description=f"Failed to connect to check server: {str(e)}",
-                code="checkserver_connection_error",
-                details={},
+                code="NEMO_CONNECTION_ERROR",
+                details={"error": str(e)},
+            )
+            return ToolPreInvokeResult(
+                continue_processing=False, violation=violation
             )
-            return ToolPreInvokeResult(continue_processing=False, violation=violation)
 
     async def tool_post_invoke(
         self, payload: ToolPostInvokePayload, context: PluginContext
@@ -188,4 +196,90 @@ async def tool_post_invoke(
         Returns:
             The result of the plugin's analysis, including whether the tool result should proceed.
         """
-        return ToolPostInvokeResult(continue_processing=True)
+        logger.info(
+            f"[NemoCheck] Starting tool post invoke hook with payload {payload}"
+        )
+
+        # Extract content from payload.result
+        # payload.result format: {'content': [{'type': 'text', 'text': 'Hello, bob!'}]}
+        result_content = payload.result.get("content", [])
+        tool_name = payload.name
+
+        if not result_content:
+            logger.warning(
+                "[NemoCheck] No content in tool result, skipping check"
+            )
+            return ToolPostInvokeResult(continue_processing=True)
+
+        # Extract text content from the content array
+        # TODO: what to do if there's actually multiple texts?
+        text_content = ""
+        for item in result_content:
+            if item.get("type") == "text":
+                text_content += item.get("text", "")
+
+        # Build NeMo check payload for tool response
+        check_nemo_payload = {
+            "model": MODEL_NAME,  # ideally optional
+            "messages": [
+                {"role": "tool", "content": text_content, "name": tool_name}
+            ],
+        }
+
+        logger.debug(
+            f"[NemoCheck] Payload for guardrail check: {check_nemo_payload}"
+        )
+
+        violation = None
+        try:
+            response = requests.post(
+                self.check_endpoint, headers=HEADERS, json=check_nemo_payload
+            )
+            if response.status_code == 200:
+                data = response.json()
+                status = data.get("status", "blocked")
+                logger.debug(f"[NemoCheck] Rails reply: {data}")
+
+                if status == "success":
+                    metadata = data.get("rails_status")
+                    result = ToolPostInvokeResult(
+                        continue_processing=True, metadata=metadata
+                    )
+                else:  # blocked
+                    metadata = data.get("rails_status")
+                    violation = PluginViolation(
+                        reason=f"Check tool rails:{status}.",
+                        description=json.dumps(data),
+                        code="NEMO_RAILS_BLOCKED",
+                        details=metadata,
+                    )
+                    result = ToolPostInvokeResult(
+                        continue_processing=False,
+                        violation=violation,
+                        metadata=metadata,
+                    )
+            else:
+                violation = PluginViolation(
+                    reason="Tool Check Unavailable",
+                    description=f"Tool response check server returned error. Status code: {response.status_code}, Response: {response.text}",
+                    code="NEMO_SERVER_ERROR",
+                    details={"status_code": response.status_code},
+                )
+                result = ToolPostInvokeResult(
+                    continue_processing=False, violation=violation
+                )
+
+            logger.info(f"[NemoCheck] Tool post invoke result: {result}")
+            return result
+
+        except Exception as e:
+            logger.error(f"[NemoCheck] Error checking tool response: {e}")
+            violation = PluginViolation(
+                reason="Tool Check Error",
+                description=f"Failed to connect to check server: {str(e)}",
+                code="NEMO_CONNECTION_ERROR",
+                details={"error": str(e)},
+            )
+            return ToolPostInvokeResult(
+                continue_processing=False, violation=violation
+            )

plugins/examples/nemocheck/tests/test_all.py

@@ -48,9 +48,13 @@ async def test_prompt_pre_hook(plugin_manager: PluginManager):
 async def test_prompt_post_hook(plugin_manager: PluginManager):
     """Test prompt post hook across all registered plugins."""
     # Customize payload for testing
-    message = Message(content=TextContent(type="text", text="prompt"), role=Role.USER)
+    message = Message(
+        content=TextContent(type="text", text="prompt"), role=Role.USER
+    )
     prompt_result = PromptResult(messages=[message])
-    payload = PromptPosthookPayload(prompt_id="test_prompt", result=prompt_result)
+    payload = PromptPosthookPayload(
+        prompt_id="test_prompt", result=prompt_result
+    )
     global_context = GlobalContext(request_id="1")
     result, _ = await plugin_manager.invoke_hook(
         PromptHookType.PROMPT_POST_FETCH, payload, global_context