feat(Test optimizer):

agentic_security/probe_actor/fuzzer.py

@@ -1,11 +1,16 @@
 import os
+import asyncio
+from typing import List, Dict, AsyncGenerator
 
 import httpx
-from loguru import logger
-from pydantic import BaseModel
-
+import numpy as np
+import pandas as pd
 from agentic_security.probe_actor.refusal import refusal_heuristic
 from agentic_security.probe_data.data import prepare_prompts
+from loguru import logger
+from pydantic import BaseModel
+from skopt import Optimizer
+from skopt.space import Real
 
 IS_VERCEL = os.getenv("IS_VERCEL", "f") == "t"
 
@@ -19,7 +24,7 @@ class ScanResult(BaseModel):
     status: bool = False
 
     @classmethod
-    def status_msg(cls, msg: str):
+    def status_msg(cls, msg: str) -> str:
         return cls(
             module=msg,
             tokens=0,
@@ -30,24 +35,29 @@ class ScanResult(BaseModel):
         ).model_dump_json()
 
 
-async def prompt_iter(prompts):
+async def prompt_iter(prompts: List[str] | AsyncGenerator) -> AsyncGenerator[str, None]:
     if isinstance(prompts, list):
         for p in prompts:
             yield p
-        return
-    async for p in prompts:
-        yield p
+    else:
+        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...")
+    request_factory,
+    max_budget: int,
+    datasets: List[Dict[str, str]] = [],
+    tools_inbox=None,
+    optimize=False,
+) -> AsyncGenerator[str, None]:
     if IS_VERCEL:
         yield ScanResult.status_msg(
-            "Vercel deployment detected. Streaming messages are not supported by serverless, plz run it locally."
+            "Vercel deployment detected. Streaming messages are not supported by serverless, please run it locally."
         )
         return
+
+    yield ScanResult.status_msg("Loading datasets...")
     prompt_modules = prepare_prompts(
         dataset_names=[m["dataset_name"] for m in datasets if m["selected"]],
         budget=max_budget,
@@ -57,63 +67,83 @@ async def perform_scan(
 
     errors = []
     refusals = []
-    size = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
-    step = 0
-    for mi, module in enumerate(prompt_modules):
+    total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
+    processed_prompts = 0
+
+    failure_rates = []
+
+    for module in prompt_modules:
         tokens = 0
         module_failures = 0
-        size = 0 if module.lazy else len(module.prompts)
-        logger.info(f"Scanning {module.dataset_name} {size}")
-        i = 0
+        module_size = 0 if module.lazy else len(module.prompts)
+        logger.info(f"Scanning {module.dataset_name} {module_size}")
+        optimizer = Optimizer(
+            [Real(0, 1)], base_estimator="GP", n_initial_points=25, acq_func="EI"
+        )
+        should_stop_early = False
         async for prompt in prompt_iter(module.prompts):
-            i += 1
-            step += 1
-            progress = 100 * (step) / size if size else 0
+            processed_prompts += 1
+            progress = 100 * processed_prompts / total_prompts if total_prompts else 0
 
-            # Naive token count
             tokens += len(prompt.split())
             try:
                 r = await request_factory.fn(prompt=prompt)
-            except httpx.RequestError as e:
+                if r.status_code >= 400:
+                    raise httpx.HTTPStatusError(
+                        f"HTTP {r.status_code}", request=r.request, response=r
+                    )
+
+                response_text = r.text
+                tokens += len(response_text.split())
+
+                if not refusal_heuristic(r.json()):
+                    refusals.append(
+                        (module.dataset_name, prompt, r.status_code, response_text)
+                    )
+                    module_failures += 1
+            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                 logger.error(f"Request error: {e}")
-                errors.append((module.dataset_name, prompt.replace("\n", ";"), e))
+                errors.append((module.dataset_name, prompt, str(e)))
                 module_failures += 1
                 continue
-            if r.status_code >= 400:
-                module_failures += 1
-                errors.append(
-                    (
-                        module.dataset_name,
-                        prompt.replace("\n", ";"),
-                        r.status_code,
-                        r.text,
-                    )
-                )
-            elif not refusal_heuristic(r.json()):
-                refusals.append(
-                    (
-                        module.dataset_name,
-                        prompt.replace("\n", ";"),
-                        r.status_code,
-                        r.text,
-                    )
-                )
-                module_failures += 1
-            # Naive token count for llm response
-            tokens += len(r.text.split())
-            total = size if size else i
+
+            failure_rate = module_failures / max(processed_prompts, 1)
+            failure_rates.append(failure_rate)
+
             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(total, 1),
+                failureRate=round(failure_rate * 100, 2),
             ).model_dump_json()
-    yield ScanResult.status_msg("Done.")
-    import pandas as pd
+
+            if not optimize:
+                continue
+            # Use the optimizer to decide whether to stop early
+            if len(failure_rates) >= 5:  # Wait for at least 5 data points
+                next_point = optimizer.ask()
+                optimizer.tell(
+                    next_point, -failure_rate
+                )  # We want to minimize failure rate
+
+                # Get the best point found so far
+                best_failure_rate = -optimizer.get_result().fun
+
+                # If the best failure rate is high, consider stopping
+                if best_failure_rate > 0.5:  # Threshold can be adjusted
+                    yield ScanResult.status_msg(
+                        f"High failure rate detected ({best_failure_rate:.2%}). Stopping this module..."
+                    )
+                    should_stop_early = True
+                    break  # Break out of the prompt loop
+
+        if should_stop_early:
+            continue  # Move to the next module
+
+    yield ScanResult.status_msg("Scan completed.")
 
     df = pd.DataFrame(
         errors + refusals, columns=["module", "prompt", "status_code", "content"]
     )
     df.to_csv("failures.csv", index=False)
-    # TODO: save all results