fix(csv file generation bug):

Author: msoedov

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 @@ def status_msg(cls, msg: str) -> str:
             prompt="",
             model="",
             refused=False,
+            latency=0,
         ).model_dump_json()
 
 

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,14 +134,17 @@ 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,
                     tokens,
                     module.dataset_name,
                     refusals,
                     errors,
+                    outputs,
                 )
+                end = time.time()
                 total_tokens += tokens
                 # logger.debug(f"Trying prompt: {prompt}, {failed=}")
                 if failed:
@@ -147,13 +153,22 @@ 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),
                     cost=cost,
                     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

@@ -209,6 +209,7 @@ async def test_successful_response_no_refusal(self):
             module_name="module_a",
             refusals=[],
             errors=[],
+            outputs=[],
         )
 
         self.assertEqual(tokens, 3)  # Tokens from "Valid response text"
@@ -226,13 +227,15 @@ async def test_successful_response_with_refusal(self):
         )
 
         refusals = []
+        outputs = []
         tokens, refusal = await process_prompt(
             request_factory=mock_request_factory,
             prompt="test prompt",
             tokens=0,
             module_name="module_a",
             refusals=refusals,
             errors=[],
+            outputs=outputs,
         )
 
         self.assertEqual(tokens, 3)  # Tokens from "Response indicating refusal"
@@ -257,6 +260,7 @@ async def test_http_error_response(self):
             module_name="module_a",
             refusals=refusals,
             errors=[],
+            outputs=[],
         )
 
     async def test_request_error(self):
@@ -273,6 +277,7 @@ async def test_request_error(self):
             module_name="module_a",
             refusals=[],
             errors=errors,
+            outputs=[],
         )
 
         self.assertEqual(tokens, 0)

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
 

agentic_security/static/index.html

@@ -95,7 +95,6 @@ <h2 class="text-2xl font-bold mb-4">Select a Config</h2>
             <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 @@ <h2 class="text-2xl font-bold">LLM API Spec</h2>
               <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 @@ <h3 class="text-lg font-semibold">Enable Many-shot
           <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

@@ -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',