Merge pull request #343 from TianHengZhuang/main

docs: document module interface and add ModuleProtocol

agentic_security/probe_data/modules/adaptive_attacks.py

@@ -10,6 +10,23 @@ url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harm
 
 
 class Module:
+    """:class:`Module` that generates adversarial prompt templates for guard testing.
+
+    Loads the `LLM-Adaptive-Attacks <https://github.com/tml-epfl/llm-adaptive-attacks>`_
+    dataset and yields structured prompt templates targeting specific harmful
+    content categories. The prompts are intended for guard evaluation, not for
+    use in production systems.
+
+    Each prompt is wrapped in a template that instructs the target model to
+    comply with potentially harmful requests, and the module produces multiple
+    template variants (refined, one-shot, simplified) for thorough guard testing.
+
+    Attributes:
+        targets: List of harmful target topics loaded from the dataset.
+        goals: List of corresponding goals for each target topic.
+    """
+
+
     def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
         r = httpx.get(url)
 

agentic_security/probe_data/modules/fine_tuned.py

@@ -1,4 +1,4 @@
-import asyncio
+import asyncio
 import os
 import uuid as U
 
@@ -10,35 +10,47 @@ AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115
 
 
 class Module:
+    """:class:`Module` for fetching prompts from a remote API and submitting LLM outputs.
+
+    Retrieves batches of prompts from a configurable HTTP endpoint (default:
+    ``https://mcp.metaheuristic.co/infer``), optionally posts each prompt to an
+    LLM proxy for evaluation, and yields the raw prompts for downstream guard
+    processing.
+
+    Configuration is passed via the ``opts`` dict:
+
+    * ``port`` (int): LLM proxy port. Defaults to ``8718``.
+    * ``max_prompts`` (int): Total prompts to process. Defaults to ``2000``.
+    * ``batch_size`` (int): Prompts fetched per API call. Defaults to ``500``.
+
+    Attributes:
+        tools_inbox: Async queue that receives guard evaluation results.
+        opts: Module configuration dictionary.
+        prompt_groups: Passed through but not used directly by this module.
+    """
+
     def __init__(
         self, prompt_groups: list[str], tools_inbox: asyncio.Queue, opts: dict = {}
     ):
         self.tools_inbox = tools_inbox
         self.opts = opts
         self.prompt_groups = prompt_groups
-        self.max_prompts = self.opts.get("max_prompts", 2000)  # Default max M prompts
+        self.max_prompts = self.opts.get("max_prompts", 2000)
         self.run_id = U.uuid4().hex
         self.batch_size = self.opts.get("batch_size", 500)
 
     async def apply(self):
         for _ in range(max(self.max_prompts // self.batch_size, 1)):
-            # Fetch prompts from the API
             prompts = await self.fetch_prompts()
-
             if not prompts:
                 logger.error("No prompts retrieved from the API.")
                 return
-
             logger.info(f"Retrieved {len(prompts)} prompts.")
-
             for i, prompt in enumerate(
                 prompts[: self.max_prompts]
-            ):  # Limit to max_prompts
+            ):
                 logger.info(f"Processing prompt {i+1}/{len(prompts)}: {prompt}")
-                # response = await self.post_prompt(prompt)
-                # logger.info(f"Response: {response}")
                 yield prompt
-
                 while not self.tools_inbox.empty():
                     ref = await self.tools_inbox.get()
                     message, _, ready = ref["message"], ref["reply"], ref["ready"]
@@ -55,7 +67,6 @@ class Module:
             "max_tokens": 1050,
             "temperature": 0.7,
         }
-
         async with httpx.AsyncClient() as client:
             try:
                 response = await client.post(uri, headers=headers, json=data)
@@ -71,7 +82,6 @@ class Module:
             "Authorization": f"Bearer {AUTH_TOKEN}",
             "Content-Type": "application/json",
         }
-
         async with httpx.AsyncClient() as client:
             try:
                 response = await client.post(

agentic_security/probe_data/modules/garak_tool.py

@@ -1,4 +1,4 @@
-import asyncio
+import asyncio
 import importlib.util
 import json
 import os
@@ -20,21 +20,37 @@ def write_garak_config_json(port):
 
 
 class Module:
+    """:class:`Module` for running Garak vulnerability probes against LLM endpoints.
+
+    Executes the `Garak <https://github.com/NVIDIA/garak>`_ adversarial probe suite
+    against a local or remote LLM proxy. The module writes a Garak REST config,
+    launches the ``python -m garak`` subprocess, and streams results into the
+    shared :data:`tools_inbox`` queue.
+
+    The module requires ``garak`` to be installed (``pip install garak``) and
+    an LLM proxy running at the address configured via the ``port`` option
+    (default: ``8718``).
+
+    Attributes:
+        tools_inbox: Async queue that receives Garak probe results.
+        opts: Configuration dictionary. Recognized key: ``port`` (int).
+    """
+
     def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
         self.tools_inbox = tools_inbox
         if not self.is_garak_installed():
             logger.error(
-                "Garak module is not installed. Please install it using 'pip install garak'"
+                "Garak module is not installed. Please install it using '"'"'pip install garak'"'"'"
             )
         self.opts = opts
 
     def is_garak_installed(self) -> bool:
+        """Return ``True`` if the ``garak`` package is importable."""
         garak_spec = importlib.util.find_spec("garak")
         return garak_spec is not None
 
     async def apply(self) -> []:
         env = os.environ.copy()
-        # Command to be executed
         command = [
             "python",
             "-m",
@@ -49,7 +65,6 @@ class Module:
         logger.info("Starting Garak tool. Writing config file.")
         write_garak_config_json(port=self.opts.get("port", 8718))
         logger.info(f"Executing command: {command}")
-        # Execute the command with the specific environment
         process = subprocess.Popen(
             command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
         )

agentic_security/probe_data/modules/inspect_ai_tool.py

@@ -1,4 +1,4 @@
-import asyncio
+import asyncio
 import importlib.util
 import os
 
@@ -12,17 +12,34 @@ inspect_ai_task = (
 
 
 class Module:
+    """:class:`Module` that runs Inspect AI evaluations against an LLM proxy.
+
+    Launches the `Inspect AI <https://github.com/UKGovernmentBEIS/inspect_ai>`_
+    evaluation framework via ``inspect eval``, targeting a local LLM proxy.
+    The module requires ``inspect_ai`` to be installed
+    (``pip install inspect_ai``).
+
+    Configuration is passed via ``opts``:
+
+    * ``port`` (int): LLM proxy base URL port. Defaults to ``8718``.
+
+    Attributes:
+        tools_inbox: Async queue that receives evaluation results.
+        opts: Configuration dictionary.
+    """
+
     name = "Inspect AI"
 
     def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
         self.tools_inbox = tools_inbox
         if not self.is_tool_installed():
             logger.error(
-                "inspect_ai module is not installed. Please install it using 'pip install inspect_ai'"
+                "inspect_ai module is not installed. Please install it using '"'"'pip install inspect_ai'"'"'"
             )
         self.opts = opts
 
     def is_tool_installed(self) -> bool:
+        """Return ``True`` if the ``inspect_ai`` package is importable."""
         inspect_ai = importlib.util.find_spec("inspect_ai")
         return inspect_ai is not None
 
@@ -38,11 +55,9 @@ class Module:
 
         logger.info(f"Started {command}")
 
-        # Read output as it becomes available
         async for line in process.stdout:
             logger.info(line.decode().strip())
 
-        # Check for errors
         err = await process.stderr.read()
         if err:
             logger.error(err.decode().strip())
@@ -52,7 +67,6 @@ class Module:
 
     async def apply(self) -> []:
         port = self.opts.get("port", 8718)
-        # Command to be executed
         command = f"inspect eval {inspect_ai_task} --model openai/gpt-4  --model-base-url=http://0.0.0.0:{port}/proxy"
         logger.info(f"Executing command: {command}")
 

agentic_security/probe_data/modules/module_protocol.py

@@ -0,0 +1,46 @@
+""":mod:`module_protocol` -- Base protocol for probe data modules.
+
+Defines the abstract Protocol that all probe data modules must implement,
+providing a standardized interface for module initialization and execution.
+
+See Also:
+    :mod:`agentic_security.probe_data.modules.garak_tool`
+    :mod:`agentic_security.probe_data.modules.fine_tuned`
+    :mod:`agentic_security.probe_data.modules.inspect_ai_tool`
+    :mod:`agentic_security.probe_data.modules.rl_model`
+    :doc:`/external_module`
+"""""
+
+from typing import Protocol, Any, AsyncGenerator, runtime_checkable
+
+
+@runtime_checkable
+class ModuleProtocol(Protocol):
+    """:class:`Protocol` defining the interface for probe data modules.
+
+    All modules in :mod:`agentic_security.probe_data.modules` share the same
+    constructor signature and the same async ``apply`` generator method.
+    The protocol captures these shared elements to support type checking.
+
+    Attributes:
+        prompt_groups: List of prompt groups to be processed.
+        tools_inbox: Async queue that receives tool execution results.
+        opts: Module-specific configuration dictionary.
+
+    Note:
+        Use the concrete :class:`Module` base class rather than this
+        protocol directly. This protocol exists to document the shared
+        interface and to enable runtime type checking.
+    """
+
+    prompt_groups: list[Any]
+    tools_inbox: asyncio.Queue
+    opts: dict
+
+    async def apply(self) -> AsyncGenerator[str, None]:
+        """Execute the module and yield result messages.
+
+        Yields:
+            str: Result messages generated during module execution.
+        """
+        ...

docs/external_module.md

@@ -1,43 +1,49 @@
-## Module Interface Documentation
+## Module Interface Documentation
 
-The `Module` class interface provides a standardized way to create and use modules in the `agentic_security` project.
+The ``Module`` class provides a standardized way to create and use probe data
+modules in the ``agentic_security`` project.
 
-Here is an example of a module that implements the `ModuleProtocol` interface. This example shows how to create a module that processes prompts and sends results to a queue.
+All modules in :mod:`agentic_security.probe_data.modules` share the same
+constructor signature and the same ``apply`` async generator method. See the
+concrete implementations for real-world usage:
 
-```python
+* :mod:`agentic_security.probe_data.modules.garak_tool`
-from typing import List, Dict, Any, AsyncGenerator
+* :mod:`agentic_security.probe_data.modules.fine_tuned`
-import asyncio
+* :mod:`agentic_security.probe_data.modules.inspect_ai_tool`
-from .module_protocol import ModuleProtocol
+* :mod:`agentic_security.probe_data.modules.rl_model`
 
-class ModuleProtocol(ModuleProtocol):
+### Interface Summary
-    def __init__(self, prompt_groups: List[Any], tools_inbox: asyncio.Queue, opts: Dict[str, Any]):
-        self.prompt_groups = prompt_groups
-        self.tools_inbox = tools_inbox
-        self.opts = opts
 
-    async def apply(self) -> AsyncGenerator[str, None]:
+Every module class accepts three constructor arguments:
-        for group in self.prompt_groups:
-            await asyncio.sleep(1) 
-            result = f"Processed {group}"
-            await self.tools_inbox.put(result)
-            yield result
-```
 
-#### Usage Example
+``def __init__(self, prompt_groups: list[Any], tools_inbox: asyncio.Queue, opts: dict = {}): ...``
 
-```python
+The ``apply`` method is an async generator that yields result strings:
-import asyncio
-import ModuleProtocol
 
-tools_inbox = asyncio.Queue()
+``async def apply(self) -> AsyncGenerator[str, None]: yield "result message"``
-prompt_groups = ["group1", "group2"]
-opts = {"max_prompts": 1000, "batch_size": 100}
 
-module = ModuleProtocol(prompt_groups, tools_inbox, opts)
+### Usage Example
 
-async def main():
+``import asyncio``
-    async for result in module.apply():
+``from agentic_security.probe_data.modules.garak_tool import Module as GarakModule``
-        print(result)
+``tools_inbox = asyncio.Queue()``
+``prompt_groups = ["group_a", "group_b"]``
+``opts = {"port": 8718}``
+``module = GarakModule(prompt_groups, tools_inbox, opts)``
+``async def main(): async for result in module.apply(): print(result)``
+``asyncio.run(main())``
 
-asyncio.run(main())
+### Defining a Custom Module
-```
+
+``import asyncio``
+``from typing import Any``
+``class MyModule:``
+``    def __init__(self, prompt_groups, tools_inbox, opts={}):``
+``        self.prompt_groups = prompt_groups``
+``        self.tools_inbox = tools_inbox``
+``        self.opts = opts``
+``    async def apply(self):``
+``        for group in self.prompt_groups:``
+``            result = "processed {0}".format(group)``
+``            await self.tools_inbox.put({"message": result})``
+``            yield result``