feat(Integrated Garak):

agentic_security/probe_actor/fuzzer.py

@@ -30,7 +30,18 @@ class ScanResult(BaseModel):
         ).model_dump_json()
 
 
-async def perform_scan(request_factory, max_budget: int, datasets: list[dict] = []):
+async def prompt_iter(prompts):
+    if isinstance(prompts, list):
+        for p in prompts:
+            yield p
+        return
+    async for p in prompts:
+        yield p
+
+
+async def perform_scan(
+    request_factory, max_budget: int, datasets: list[dict] = [], tools_inbox=None
+):
     yield ScanResult.status_msg("Loading datasets...")
     if IS_VERCEL:
         yield ScanResult.status_msg(
@@ -40,20 +51,24 @@ async def perform_scan(request_factory, max_budget: int, datasets: list[dict] =
     prompt_modules = prepare_prompts(
         dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
         budget=max_budget,
+        tools_inbox=tools_inbox,
     )
     yield ScanResult.status_msg("Datasets loaded. Starting scan...")
 
     errors = []
     refusals = []
-    size = sum(len(m.prompts) for m in prompt_modules)
+    size = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
     step = 0
     for mi, module in enumerate(prompt_modules):
         tokens = 0
         module_failures = 0
-        logger.info(f"Scanning {module.dataset_name} {len(module.prompts)}")
-        for i, prompt in enumerate(module.prompts):
+        size = 0 if module.lazy else len(module.prompts)
+        logger.info(f"Scanning {module.dataset_name} {size}")
+        i = 0
+        async for prompt in prompt_iter(module.prompts):
+            i += 1
             step += 1
-            progress = 100 * (step) / size
+            progress = 100 * (step) / size if size else 0
 
             # Naive token count
             tokens += len(prompt.split())
@@ -86,12 +101,13 @@ async def perform_scan(request_factory, max_budget: int, datasets: list[dict] =
                 module_failures += 1
             # Naive token count for llm response
             tokens += len(r.text.split())
+            total = size if size else i
             yield ScanResult(
                 module=module.dataset_name,
                 tokens=round(tokens / 1000, 1),
                 cost=round(tokens * 1.5 / 1000_000, 2),
                 progress=round(progress, 2),
-                failureRate=100 * module_failures / max(len(module.prompts), 1),
+                failureRate=100 * module_failures / max(total, 1),
             ).model_dump_json()
     yield ScanResult.status_msg("Done.")
     import pandas as pd