mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 06:09:55 +02:00
feat(add process_prompt_batch):
This commit is contained in:
@@ -20,6 +20,10 @@ from agentic_security.probe_data.data import prepare_prompts
|
||||
# TODO: full log file
|
||||
|
||||
MAX_PROMPT_LENGTH = 2048
|
||||
BUDGET_MULTIPLIER = 100_000_000
|
||||
INITIAL_OPTIMIZER_POINTS = 25
|
||||
MIN_FAILURE_SAMPLES = 5
|
||||
FAILURE_RATE_THRESHOLD = 0.5
|
||||
|
||||
|
||||
async def generate_prompts(
|
||||
@@ -82,6 +86,27 @@ async def process_prompt(
|
||||
return tokens, True
|
||||
|
||||
|
||||
async def process_prompt_batch(
|
||||
request_factory,
|
||||
prompts: list[str],
|
||||
tokens: int,
|
||||
module_name: str,
|
||||
refusals,
|
||||
errors,
|
||||
outputs,
|
||||
) -> tuple[int, int]:
|
||||
tasks = [
|
||||
process_prompt(
|
||||
request_factory, p, tokens, module_name, refusals, errors, outputs
|
||||
)
|
||||
for p in prompts
|
||||
]
|
||||
results = await asyncio.gather(*tasks)
|
||||
total_tokens = sum(r[0] for r in results)
|
||||
failures = sum(1 for r in results if r[1])
|
||||
return total_tokens, failures
|
||||
|
||||
|
||||
async def perform_single_shot_scan(
|
||||
request_factory,
|
||||
max_budget: int,
|
||||
|
||||
Reference in New Issue
Block a user