mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-08-09 03:06:05 +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
|
# TODO: full log file
|
||||||
|
|
||||||
MAX_PROMPT_LENGTH = 2048
|
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(
|
async def generate_prompts(
|
||||||
@@ -82,6 +86,27 @@ async def process_prompt(
|
|||||||
return tokens, True
|
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(
|
async def perform_single_shot_scan(
|
||||||
request_factory,
|
request_factory,
|
||||||
max_budget: int,
|
max_budget: int,
|
||||||
|
|||||||
Reference in New Issue
Block a user