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", } } -} +}