mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-09-30 03:41:45 +02:00
docs: add module docstrings
This commit is contained in:
@@ -1,4 +1,4 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import importlib.util
|
import importlib.util
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
@@ -20,21 +20,37 @@ def write_garak_config_json(port):
|
|||||||
|
|
||||||
|
|
||||||
class Module:
|
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 = {}):
|
def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
|
||||||
self.tools_inbox = tools_inbox
|
self.tools_inbox = tools_inbox
|
||||||
if not self.is_garak_installed():
|
if not self.is_garak_installed():
|
||||||
logger.error(
|
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
|
self.opts = opts
|
||||||
|
|
||||||
def is_garak_installed(self) -> bool:
|
def is_garak_installed(self) -> bool:
|
||||||
|
"""Return ``True`` if the ``garak`` package is importable."""
|
||||||
garak_spec = importlib.util.find_spec("garak")
|
garak_spec = importlib.util.find_spec("garak")
|
||||||
return garak_spec is not None
|
return garak_spec is not None
|
||||||
|
|
||||||
async def apply(self) -> []:
|
async def apply(self) -> []:
|
||||||
env = os.environ.copy()
|
env = os.environ.copy()
|
||||||
# Command to be executed
|
|
||||||
command = [
|
command = [
|
||||||
"python",
|
"python",
|
||||||
"-m",
|
"-m",
|
||||||
@@ -49,7 +65,6 @@ class Module:
|
|||||||
logger.info("Starting Garak tool. Writing config file.")
|
logger.info("Starting Garak tool. Writing config file.")
|
||||||
write_garak_config_json(port=self.opts.get("port", 8718))
|
write_garak_config_json(port=self.opts.get("port", 8718))
|
||||||
logger.info(f"Executing command: {command}")
|
logger.info(f"Executing command: {command}")
|
||||||
# Execute the command with the specific environment
|
|
||||||
process = subprocess.Popen(
|
process = subprocess.Popen(
|
||||||
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
|
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user