diff --git a/agentic_security/probe_data/modules/adaptive_attacks.py b/agentic_security/probe_data/modules/adaptive_attacks.py index 33a5dd9..6263d17 100644 --- a/agentic_security/probe_data/modules/adaptive_attacks.py +++ b/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 `_ + 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) diff --git a/agentic_security/probe_data/modules/fine_tuned.py b/agentic_security/probe_data/modules/fine_tuned.py index eb70920..49eef4a 100644 --- a/agentic_security/probe_data/modules/fine_tuned.py +++ b/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( @@ -84,4 +94,4 @@ class Module: return data.get("prompts", []) except httpx.RequestError as e: logger.error(f"Failed to fetch prompts: {e}") - return [] + return [] diff --git a/agentic_security/probe_data/modules/garak_tool.py b/agentic_security/probe_data/modules/garak_tool.py index a68b6e7..0268ead 100644 --- a/agentic_security/probe_data/modules/garak_tool.py +++ b/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 `_ 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 ) @@ -90,4 +105,4 @@ SPEC = { "response_json_field": "$.choices[0].message.content", } } -} +} diff --git a/agentic_security/probe_data/modules/inspect_ai_tool.py b/agentic_security/probe_data/modules/inspect_ai_tool.py index d2058e7..78afea8 100644 --- a/agentic_security/probe_data/modules/inspect_ai_tool.py +++ b/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 `_ + 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}") @@ -66,4 +80,4 @@ class Module: yield message ready.set() logger.info(f"{self.name} tool finished.") - await proc + await proc diff --git a/agentic_security/probe_data/modules/module_protocol.py b/agentic_security/probe_data/modules/module_protocol.py new file mode 100644 index 0000000..30355ba --- /dev/null +++ b/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. + """ + ... \ No newline at end of file diff --git a/docs/external_module.md b/docs/external_module.md index 3721b5c..95851a8 100644 --- a/docs/external_module.md +++ b/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 -from typing import List, Dict, Any, AsyncGenerator -import asyncio -from .module_protocol import ModuleProtocol +* :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` -class ModuleProtocol(ModuleProtocol): - 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 +### Interface Summary - async def apply(self) -> AsyncGenerator[str, None]: - for group in self.prompt_groups: - await asyncio.sleep(1) - result = f"Processed {group}" - await self.tools_inbox.put(result) - yield result -``` +Every module class accepts three constructor arguments: -#### Usage Example +``def __init__(self, prompt_groups: list[Any], tools_inbox: asyncio.Queue, opts: dict = {}): ...`` -```python -import asyncio -import ModuleProtocol +The ``apply`` method is an async generator that yields result strings: -tools_inbox = asyncio.Queue() -prompt_groups = ["group1", "group2"] -opts = {"max_prompts": 1000, "batch_size": 100} +``async def apply(self) -> AsyncGenerator[str, None]: yield "result message"`` -module = ModuleProtocol(prompt_groups, tools_inbox, opts) +### Usage Example -async def main(): - async for result in module.apply(): - print(result) +``import asyncio`` +``from agentic_security.probe_data.modules.garak_tool import Module as GarakModule`` +``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`` \ No newline at end of file