diff --git a/agentic_security/models/schemas.py b/agentic_security/models/schemas.py index 3f9620e..ae60778 100644 --- a/agentic_security/models/schemas.py +++ b/agentic_security/models/schemas.py @@ -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() diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index 68c069f..eacce62 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -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 diff --git a/agentic_security/probe_actor/test_fuzzer.py b/agentic_security/probe_actor/test_fuzzer.py index ec6c8c0..830d606 100644 --- a/agentic_security/probe_actor/test_fuzzer.py +++ b/agentic_security/probe_actor/test_fuzzer.py @@ -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) diff --git a/agentic_security/static/base.js b/agentic_security/static/base.js index 911ebca..0b7be6c 100644 --- a/agentic_security/static/base.js +++ b/agentic_security/static/base.js @@ -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 diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html index 6eded55..ff00c50 100644 --- a/agentic_security/static/index.html +++ b/agentic_security/static/index.html @@ -95,7 +95,6 @@

LLM API Spec

- Latency: {{latency}}s > {{okMsg}} + Latency: {{latency}}s +
@@ -419,6 +420,8 @@ > {{okMsg}} + Latency: {{latency}}s +
diff --git a/agentic_security/static/main.js b/agentic_security/static/main.js index fa96844..c988ef7 100644 --- a/agentic_security/static/main.js +++ b/agentic_security/static/main.js @@ -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 src + // Convert image response to a data SELF_URL for the 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',