CalvinBackup/agentic_security
1
0
Fork 0
mirror of https://github.com/msoedov/agentic_security.git synced 2026-06-24 06:09:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(csv file generation bug):

Browse Source
This commit is contained in:
Alexander Myasoedov
2025-02-17 20:21:47 +02:00
parent 653e9a7234
commit 4ffca42e48
6 changed files with 52 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agentic_security/models/schemas.py
+2
View File
@@ -35,6 +35,7 @@ class ScanResult(BaseModel):
prompt: str = ""
model: str = ""
refused: bool = False
latency: float = 0.0
@classmethod
def status_msg(cls, msg: str) -> str:
@@ -48,6 +49,7 @@ class ScanResult(BaseModel):
prompt="",
model="",
refused=False,
latency=0,
).model_dump_json()
agentic_security/probe_actor/fuzzer.py
+22 -5
View File
@@ -1,7 +1,7 @@
import asyncio
import random
from collections.abc import AsyncGenerator
import time
import httpx
import pandas as pd
from loguru import logger
@@ -44,7 +44,7 @@ def multi_modality_spec(llm_spec):
async def process_prompt(
request_factory, prompt, tokens, module_name, refusals, errors
request_factory, prompt, tokens, module_name, refusals, errors, outputs
) -> tuple[int, bool]:
"""
Process a single prompt and update the token count and failure status.
@@ -63,10 +63,12 @@ async def process_prompt(
response_text = response.text
tokens += len(response_text.split())
if refusal_heuristic(response.json()):
refused = refusal_heuristic(response.json())
if refused:
refusals.append((module_name, prompt, response.status_code, response_text))
return tokens, True
return tokens, False
outputs.append((module_name, prompt, response_text, refused))
return tokens, refused
except httpx.RequestError as exc:
logger.error(f"Request error: {exc}")
@@ -98,6 +100,7 @@ async def perform_single_shot_scan(
errors = []
refusals = []
outputs = []
total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
processed_prompts = 0
@@ -131,6 +134,7 @@ async def perform_single_shot_scan(
100 * processed_prompts / total_prompts if total_prompts else 0
)
total_tokens -= tokens
start = time.time()
tokens, failed = await process_prompt(
request_factory,
prompt,
@@ -138,7 +142,9 @@ async def perform_single_shot_scan(
module.dataset_name,
refusals,
errors,
outputs,
)
end = time.time()
total_tokens += tokens
# logger.debug(f"Trying prompt: {prompt}, {failed=}")
if failed:
@@ -147,6 +153,13 @@ async def perform_single_shot_scan(
failure_rates.append(failure_rate)
cost = calculate_cost(tokens)
# TODO: improve this cond
last_output = outputs[-1] if outputs else None
if last_output and last_output[1] == prompt:
response_text = last_output[2]
else:
response_text = ""
yield ScanResult(
module=module.dataset_name,
tokens=round(tokens / 1000, 1),
@@ -154,6 +167,8 @@ async def perform_single_shot_scan(
progress=round(progress, 2),
failureRate=round(failure_rate * 100, 2),
prompt=prompt[:MAX_PROMPT_LENGTH],
latency=end - start,
model=response_text,
).model_dump_json()
if optimize and len(failure_rates) >= 5:
@@ -219,6 +234,7 @@ async def perform_many_shot_scan(
errors = []
refusals = []
outputs = []
total_prompts = sum(len(m.prompts) for m in prompt_modules if not m.lazy)
processed_prompts = 0
@@ -270,6 +286,7 @@ async def perform_many_shot_scan(
module.dataset_name,
refusals,
errors,
outputs,
)
if failed:
module_failures += 1
agentic_security/probe_actor/test_fuzzer.py
+5
View File
@@ -209,6 +209,7 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
module_name="module_a",
refusals=[],
errors=[],
outputs=[],
)
self.assertEqual(tokens, 3) # Tokens from "Valid response text"
@@ -226,6 +227,7 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
)
refusals = []
outputs = []
tokens, refusal = await process_prompt(
request_factory=mock_request_factory,
prompt="test prompt",
@@ -233,6 +235,7 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
module_name="module_a",
refusals=refusals,
errors=[],
outputs=outputs,
)
self.assertEqual(tokens, 3) # Tokens from "Response indicating refusal"
@@ -257,6 +260,7 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
module_name="module_a",
refusals=refusals,
errors=[],
outputs=[],
)
async def test_request_error(self):
@@ -273,6 +277,7 @@ class TestProcessPrompt(unittest.IsolatedAsyncioTestCase):
module_name="module_a",
refusals=[],
errors=errors,
outputs=[],
)
self.assertEqual(tokens, 0)
agentic_security/static/base.js
+7 -7
View File
@@ -1,13 +1,13 @@
let URL = window.location.href;
if (URL.endsWith('/')) {
URL = URL.slice(0, -1);
let SELF_URL = window.location.href;
if (SELF_URL.endsWith('/')) {
SELF_URL = SELF_URL.slice(0, -1);
}
URL = URL.replace('/#', '');
SELF_URL = SELF_URL.replace('/#', '');
// Vue application
let LLM_SPECS = [
`POST ${URL}/v1/self-probe
`POST ${SELF_URL}/v1/self-probe
Authorization: Bearer XXXXX
Content-Type: application/json
@@ -79,7 +79,7 @@ Content-Type: application/json
]
}
`,
`POST ${URL}/v1/self-probe-image
`POST ${SELF_URL}/v1/self-probe-image
Authorization: Bearer XXXXX
Content-Type: application/json
@@ -101,7 +101,7 @@ Content-Type: application/json
}
]
`,
`POST ${URL}/v1/self-probe-file
`POST ${SELF_URL}/v1/self-probe-file
Authorization: Bearer $GROQ_API_KEY
Content-Type: multipart/form-data
agentic_security/static/index.html
+4 -1
View File
@@ -95,7 +95,6 @@
<h2 class="text-2xl font-bold">LLM API Spec</h2>
<span :class="statusDotClass"
class="w-3 h-3 rounded-full mr-2"></span>
<span v-if="latency" class="text-sm text-gray-400 ml-2">Latency: {{latency}}s</span>
<svg :class="{'rotate-180': showLLMSpec}"
class="w-6 h-6 transition-transform duration-200"
xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none"
@@ -140,6 +139,8 @@
<strong class="font-bold">></strong>
<span class="block sm:inline">{{okMsg}}</span>
</div>
<span v-if="latency" class="text-sm text-gray-400 ml-2">Latency: {{latency}}s</span>
<!-- Action Buttons -->
<section class="flex justify-center space-x-4 mt-10">
@@ -419,6 +420,8 @@
<strong class="font-bold">></strong>
<span class="block sm:inline">{{okMsg}}</span>
</div>
<span v-if="latency" class="text-sm text-gray-400 ml-2">Latency: {{latency}}s</span>
<!-- Action Buttons -->
<section class="flex justify-center space-x-4">
agentic_security/static/main.js
+12 -7
View File
@@ -193,7 +193,8 @@ var app = new Vue({
let payload = {
spec: this.modelSpec,
};
const response = await fetch(`${URL}/verify`, {
let startTime = performance.now(); // Capture start time
const response = await fetch(`${SELF_URL}/verify`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
@@ -202,7 +203,10 @@ var app = new Vue({
});
console.log(response);
let r = await response.json();
this.latency = r.elapsed;
let endTime = performance.now(); // Capture end time
let latency = endTime - startTime; // Calculate latency in milliseconds
latency = latency.toFixed(3) / 1000; // Round to 2 decimal places
this.latency = latency;
if (!response.ok) {
this.updateStatusDot(false);
this.errorMsg = 'Integration verification failed:' + JSON.stringify(r);
@@ -218,7 +222,7 @@ var app = new Vue({
this.saveStateToLocalStorage();
},
loadConfigs: async function () {
const response = await fetch(`${URL}/v1/data-config`, {
const response = await fetch(`${SELF_URL}/v1/data-config`, {
method: 'GET',
headers: {
'Content-Type': 'application/json',
@@ -290,6 +294,7 @@ var app = new Vue({
this.okMsg = `${event.module}`;
return
}
this.latency = event.latency.toFixed(3);
console.log('New event');
// { "module": "Module 49", "tokens": 480, "cost": 4.800000000000001, "progress": 9.8 }
let progress = event.progress;
@@ -325,14 +330,14 @@ var app = new Vue({
let payload = {
table: this.mainTable,
};
const response = await fetch(`${URL}/plot.jpeg`, {
const response = await fetch(`${SELF_URL}/plot.jpeg`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify(payload),
});
// Convert image response to a data URL for the <img> src
// Convert image response to a data SELF_URL for the <img> src
const blob = await response.blob();
const reader = new FileReader();
reader.readAsDataURL(blob);
@@ -375,7 +380,7 @@ var app = new Vue({
},
stopScan: async function () {
this.scanRunning = false;
const response = await fetch(`${URL}/stop`, {
const response = await fetch(`${SELF_URL}/stop`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
@@ -391,7 +396,7 @@ var app = new Vue({
optimize: this.optimize,
enableMultiStepAttack: this.enableMultiStepAttack,
};
const response = await fetch(`${URL}/scan`, {
const response = await fetch(`${SELF_URL}/scan`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',