feat(Update modules interface):

agentic_security/probe_data/__init__.py

@@ -180,6 +180,10 @@ REGISTRY = [
         "selected": False,
         "url": "https://github.com/leondz/garak2",
         "dynamic": True,
+        "opts": {
+            "port": 8718,
+            "modules": ["encoding"],
+        },
     },
     {
         "dataset_name": "InspectAI",

agentic_security/probe_data/data.py

@@ -213,7 +213,7 @@ def load_generic_csv(url, name, column="prompt", predicator=None):
     )
 
 
-def prepare_prompts(dataset_names, budget, tools_inbox=None):
+def prepare_prompts(dataset_names, budget, tools_inbox=None, options=[]):
     # ## Datasets used and cleaned:
     # markush1/LLM-Jailbreak-Classifier
     # 1. Open-Orca/OpenOrca
@@ -255,28 +255,31 @@ def prepare_prompts(dataset_names, budget, tools_inbox=None):
                 logger.error(f"Error loading {dataset_name}: {e}")
 
     dynamic_datasets = {
-        "Steganography": lambda: Stenography(group),
-        "llm-adaptive-attacks": lambda: dataset_from_iterator(
-            "llm-adaptive-attacks", adaptive_attacks.Module(group).apply()
+        "Steganography": lambda opts: Stenography(group),
+        "llm-adaptive-attacks": lambda opts: dataset_from_iterator(
+            "llm-adaptive-attacks",
+            adaptive_attacks.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
         ),
-        "Garak": lambda: dataset_from_iterator(
+        "Garak": lambda opts: dataset_from_iterator(
             "Garak",
-            garak_tool.Module(group, tools_inbox=tools_inbox).apply(),
+            garak_tool.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
             lazy=True,
         ),
-        "InspectAI": lambda: dataset_from_iterator(
+        "InspectAI": lambda opts: dataset_from_iterator(
             "InspectAI",
             inspect_ai_tool.Module(group, tools_inbox=tools_inbox).apply(),
             lazy=True,
         ),
-        "GPT fuzzer": lambda: [],
+        "GPT fuzzer": lambda opts: [],
     }
 
     dynamic_groups = []
-    for dataset_name in dataset_names:
+    options = options or [{} for _ in dataset_names]
+    for dataset_name, opts in zip(dataset_names, options):
         if dataset_name in dynamic_datasets:
             logger.info(f"Loading {dataset_name}")
-            ds = dynamic_datasets[dataset_name]()
+
+            ds = dynamic_datasets[dataset_name](opts)
 
             for g in ds:
                 dynamic_groups.append(g)

agentic_security/probe_data/modules/adaptive_attacks.py

@@ -3,12 +3,13 @@ import io
 import httpx
 import pandas as pd
 from loguru import logger
+import asyncio
 
 url = "https://raw.githubusercontent.com/tml-epfl/llm-adaptive-attacks/main/harmful_behaviors/harmful_behaviors_pair.csv"
 
 
 class Module:
-    def __init__(self, prompt_groups: []):
+    def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
         r = httpx.get(url)
 
         content = r.content

agentic_security/probe_data/modules/garak_tool.py

@@ -11,21 +11,22 @@ from loguru import logger
 GARAK_CONFIG = "garak_rest.json"
 
 
-def write_garak_config_json():
+def write_garak_config_json(port):
     with open(GARAK_CONFIG, "w") as f:
-        f.write(json.dumps(SPEC))
+        f.write(json.dumps(SPEC, indent=4).replace("$PORT", str(port)))
 
 
 # TODO: add config params to data registry
 
 
 class Module:
-    def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
+    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'"
             )
+        self.opts = opts
 
     def is_garak_installed(self) -> bool:
         garak_spec = importlib.util.find_spec("garak")
@@ -46,7 +47,7 @@ class Module:
             "encoding",
         ]
         logger.info("Starting Garak tool. Writing config file.")
-        write_garak_config_json()
+        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(
@@ -63,7 +64,6 @@ class Module:
             ready.set()
         logger.info("Garak tool finished.")
         logger.info(f"stdout: {out}")
-        logger.error(f"exit code: {process.returncode}")
         if process.returncode != 0:
             logger.error(f"Error executing command: {command}")
             logger.error(f"err: {err}")
@@ -73,8 +73,8 @@ class Module:
 SPEC = {
     "rest": {
         "RestGenerator": {
-            "name": "OpenAI GPT-4 Service",
-            "uri": "http://0.0.0.0:8718/proxy/chat/completions",
+            "name": "Agentic Security Proxy Service",
+            "uri": "http://0.0.0.0:$PORT/proxy/chat/completions",
             "method": "POST",
             "headers": {
                 "Authorization": "Bearer $OPENAI_API_KEY",

agentic_security/probe_data/modules/inspect_ai_task.py

@@ -1,13 +0,0 @@
-from inspect_ai import Task, eval, task
-from inspect_ai.dataset import example_dataset
-from inspect_ai.scorer import model_graded_fact
-from inspect_ai.solver import chain_of_thought, generate, self_critique
-
-
-@task
-def theory_of_mind():
-    return Task(
-        dataset=example_dataset("theory_of_mind"),
-        plan=[chain_of_thought(), generate(), self_critique()],
-        scorer=model_graded_fact(),
-    )

agentic_security/probe_data/modules/inspect_ai_tool.py

@@ -14,7 +14,7 @@ inspect_ai_task = (
 class Module:
     name = "Inspect AI"
 
-    def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
+    def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}):
         self.tools_inbox = tools_inbox
         if not self.is_tool_installed():
             logger.error(

agentic_security/probe_data/modules/test_adaptive_attacks.py

@@ -7,7 +7,7 @@ class TestModule:
     # Module can be initialized with a list of prompt groups.
     def test_initialize_with_prompt_groups(self):
         prompt_groups = []
-        module = Module(prompt_groups)
+        module = Module(prompt_groups, None, {})
         assert module is not None
         assert isinstance(module, Module)
         assert len(module.goals) == snapshot(50)