fix(pc):

.pre-commit-config.yaml

@@ -20,6 +20,7 @@ repos:
       - id: flake8
         language_version: python3.11
         additional_dependencies: [flake8-docstrings]
+        exclude: '^(tests)/'
 
   # - repo: https://github.com/PyCQA/isort
   #   rev: 7.0.0

agentic_security/__init__.py

@@ -2,6 +2,6 @@ from agentic_security.cache_config import ensure_cache_dir
 
 ensure_cache_dir()
 
-from .lib import SecurityScanner
+from .lib import SecurityScanner  # noqa: E402
 
 __all__ = ["SecurityScanner", "ensure_cache_dir"]

agentic_security/executor/concurrent.py

@@ -61,7 +61,11 @@ class ExecutorMetrics:
         if self.latencies:
             sorted_latencies = sorted(self.latencies)
             p95_index = int(len(sorted_latencies) * 0.95)
-            p95_latency_ms = sorted_latencies[p95_index] * 1000 if p95_index < len(sorted_latencies) else 0.0
+            p95_latency_ms = (
+                sorted_latencies[p95_index] * 1000
+                if p95_index < len(sorted_latencies)
+                else 0.0
+            )
         else:
             p95_latency_ms = 0.0
 

agentic_security/probe_actor/refusal.py

@@ -50,7 +50,6 @@ class RefusalClassifierPlugin(ABC):
         Returns:
             bool: True if the response contains a refusal, False otherwise.
         """
-        pass
 
 
 class DefaultRefusalClassifier(RefusalClassifierPlugin):

agentic_security/probe_data/audio_generator.py

@@ -16,8 +16,6 @@ logger = logging.getLogger(__name__)
 class AudioGenerationError(Exception):
     """Custom exception for errors during audio generation."""
 
-    pass
-
 
 def encode(content: bytes) -> str:
     encoded_content = base64.b64encode(content).decode("utf-8")

agentic_security/probe_data/modules/rl_model.py

@@ -20,12 +20,10 @@ class PromptSelectionInterface(ABC):
     @abstractmethod
     def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         """Selects the next prompt based on current state and guard result."""
-        pass
 
     @abstractmethod
     def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
         """Selects the next prompts based on current state and guard result."""
-        pass
 
     @abstractmethod
     def update_rewards(
@@ -36,7 +34,6 @@ class PromptSelectionInterface(ABC):
         passed_guard: bool,
     ) -> None:
         """Updates internal rewards based on the outcome of the last selected prompt."""
-        pass
 
 
 class RandomPromptSelector(PromptSelectionInterface):

agentic_security/probe_data/modules/test_rl_model.py

@@ -39,7 +39,9 @@ def mock_rl_selector(dataset_prompts) -> Mock:
             self.prompts = prompts
             self.idx = 0
 
-        def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        def select_next_prompts(
+            self, current_prompt: str, passed_guard: bool
+        ) -> list[str]:
             self.idx = (self.idx + 1) % len(self.prompts)
             return [self.prompts[self.idx]]
 

agentic_security/probe_data/unified_loader.py

@@ -1,6 +1,6 @@
 """Unified dataset loader for CSV, HuggingFace, and proxy sources."""
 
-from typing import Any, Literal, Optional
+from typing import Literal
 from pydantic import BaseModel, Field
 
 from agentic_security.logutils import logger
@@ -20,28 +20,26 @@ class InputSourceConfig(BaseModel):
     )
     enabled: bool = Field(default=True, description="Whether this source is enabled")
     dataset_name: str = Field(description="Name/identifier of the dataset")
-    weight: float = Field(default=1.0, ge=0.0, description="Sampling weight for merging")
+    weight: float = Field(
+        default=1.0, ge=0.0, description="Sampling weight for merging"
+    )
 
     # CSV-specific fields
-    path: Optional[str] = Field(
+    path: str | None = Field(default=None, description="File path for CSV sources")
-        default=None, description="File path for CSV sources"
+    prompt_column: str | None = Field(
-    )
-    prompt_column: Optional[str] = Field(
         default="prompt", description="Column name containing prompts"
     )
 
     # HuggingFace-specific fields
-    split: Optional[str] = Field(
+    split: str | None = Field(
         default="train", description="Dataset split to load (train/test/validation)"
     )
-    max_samples: Optional[int] = Field(
+    max_samples: int | None = Field(
         default=None, ge=1, description="Maximum number of samples to load"
     )
 
     # URL for custom sources
-    url: Optional[str] = Field(
+    url: str | None = Field(default=None, description="URL for remote CSV files")
-        default=None, description="URL for remote CSV files"
-    )
 
 
 class UnifiedDatasetLoader:
@@ -122,15 +120,19 @@ class UnifiedDatasetLoader:
         elif config.url:
             # Remote CSV file
             logger.info(f"Loading CSV from URL: {config.url}")
-            mappings = {config.prompt_column: "prompt"} if config.prompt_column else None
+            mappings = (
+                {config.prompt_column: "prompt"} if config.prompt_column else None
+            )
             dataset = load_dataset_generic(
                 name=config.dataset_name,
                 url=config.url,
                 mappings=mappings,
-                metadata={"source_type": "csv", "url": config.url}
+                metadata={"source_type": "csv", "url": config.url},
             )
         else:
-            raise ValueError(f"CSV source {config.dataset_name} requires either path or url")
+            raise ValueError(
+                f"CSV source {config.dataset_name} requires either path or url"
+            )
 
         # Apply max_samples limit if specified
         if config.max_samples and len(dataset.prompts) > config.max_samples:
@@ -138,7 +140,7 @@ class UnifiedDatasetLoader:
                 f"Limiting {config.dataset_name} from {len(dataset.prompts)} "
                 f"to {config.max_samples} samples"
             )
-            dataset.prompts = dataset.prompts[:config.max_samples]
+            dataset.prompts = dataset.prompts[: config.max_samples]
 
         return dataset
 
@@ -167,7 +169,7 @@ class UnifiedDatasetLoader:
             metadata={
                 "source_type": "huggingface",
                 "split": config.split,
-            }
+            },
         )
 
         # Apply max_samples limit if specified
@@ -176,7 +178,7 @@ class UnifiedDatasetLoader:
                 f"Limiting {config.dataset_name} from {len(dataset.prompts)} "
                 f"to {config.max_samples} samples"
             )
-            dataset.prompts = dataset.prompts[:config.max_samples]
+            dataset.prompts = dataset.prompts[: config.max_samples]
 
         return dataset
 
@@ -195,7 +197,7 @@ class UnifiedDatasetLoader:
         return create_probe_dataset(
             config.dataset_name,
             [],
-            {"source_type": "proxy", "status": "not_implemented"}
+            {"source_type": "proxy", "status": "not_implemented"},
         )
 
     def _merge_weighted(

tests/executor/test_circuit_breaker.py

@@ -1,8 +1,6 @@
 """Tests for CircuitBreaker."""
 
-import pytest
 import time
-from unittest.mock import patch
 from agentic_security.executor.circuit_breaker import CircuitBreaker
 
 

tests/executor/test_concurrent.py

@@ -2,7 +2,7 @@
 
 import pytest
 import asyncio
-from unittest.mock import Mock, patch, AsyncMock
+from unittest.mock import Mock, patch
 from agentic_security.executor.concurrent import ConcurrentExecutor, ExecutorMetrics
 from agentic_security.probe_actor.state import FuzzerState
 
@@ -268,7 +268,10 @@ class TestConcurrentExecutor:
             # 5 requests with rate=5/s and burst=2
             # First 2 immediate, next 3 should take ~0.6s total
             await executor.execute_batch(
-                request_factory, ["p1", "p2", "p3", "p4", "p5"], "test_module", fuzzer_state
+                request_factory,
+                ["p1", "p2", "p3", "p4", "p5"],
+                "test_module",
+                fuzzer_state,
             )
             elapsed = time.monotonic() - start
 

tests/executor/test_rate_limiter.py

@@ -3,7 +3,6 @@
 import asyncio
 import pytest
 import time
-from unittest.mock import patch
 from agentic_security.executor.rate_limiter import TokenBucketRateLimiter
 
 

tests/probe_actor/test_fuzzer.py

@@ -88,7 +88,9 @@ async def test_perform_many_shot_scan_probe_injection(
     ]
 
     msj_prepare_prompts_mock.return_value = [
-        MagicMock(dataset_name="msj_probe_module", prompts=["msj_probe_prompt"], lazy=False)
+        MagicMock(
+            dataset_name="msj_probe_module", prompts=["msj_probe_prompt"], lazy=False
+        )
     ]
 
     # Mock request_factory

tests/probe_data/test_unified_loader.py

@@ -1,7 +1,7 @@
 """Tests for unified dataset loader."""
 
 import pytest
-from unittest.mock import Mock, patch, AsyncMock
+from unittest.mock import patch
 from agentic_security.probe_data.unified_loader import (
     InputSourceConfig,
     UnifiedDatasetLoader,
@@ -85,12 +85,12 @@ class TestUnifiedDatasetLoader:
             prompts=["prompt1", "prompt2", "prompt3"],
             tokens=10,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            return_value=mock_dataset
+            return_value=mock_dataset,
         ):
             result = await loader.load_all()
 
@@ -114,12 +114,12 @@ class TestUnifiedDatasetLoader:
             prompts=["hf_prompt1", "hf_prompt2"],
             tokens=8,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_dataset_generic",
-            return_value=mock_dataset
+            return_value=mock_dataset,
         ):
             result = await loader.load_all()
 
@@ -151,19 +151,19 @@ class TestUnifiedDatasetLoader:
             prompts=["prompt1"],
             tokens=5,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
         mock_dataset2 = ProbeDataset(
             dataset_name="csv2",
             prompts=["prompt2", "prompt3"],
             tokens=10,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            side_effect=[mock_dataset1, mock_dataset2]
+            side_effect=[mock_dataset1, mock_dataset2],
         ):
             result = await loader.load_all()
 
@@ -199,12 +199,12 @@ class TestUnifiedDatasetLoader:
             prompts=["prompt1"],
             tokens=5,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            return_value=mock_dataset
+            return_value=mock_dataset,
         ) as mock_load:
             result = await loader.load_all()
 
@@ -229,12 +229,12 @@ class TestUnifiedDatasetLoader:
             prompts=["prompt1", "prompt2", "prompt3", "prompt4", "prompt5"],
             tokens=20,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            return_value=mock_dataset
+            return_value=mock_dataset,
         ):
             result = await loader.load_all()
 
@@ -253,7 +253,7 @@ class TestUnifiedDatasetLoader:
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            side_effect=Exception("File not found")
+            side_effect=Exception("File not found"),
         ):
             result = await loader.load_all()
 
@@ -299,19 +299,19 @@ class TestUnifiedDatasetLoader:
             prompts=["a"],
             tokens=1,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
         mock_dataset2 = ProbeDataset(
             dataset_name="high_weight",
             prompts=["b"],
             tokens=1,
             approx_cost=0.0,
-            metadata={}
+            metadata={},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_csv",
-            side_effect=[mock_dataset1, mock_dataset2]
+            side_effect=[mock_dataset1, mock_dataset2],
         ):
             result = await loader.load_all()
 
@@ -347,12 +347,12 @@ class TestUnifiedDatasetLoader:
             prompts=["remote_prompt"],
             tokens=5,
             approx_cost=0.0,
-            metadata={"source_type": "csv", "url": "https://example.com/data.csv"}
+            metadata={"source_type": "csv", "url": "https://example.com/data.csv"},
         )
 
         with patch(
             "agentic_security.probe_data.unified_loader.load_dataset_generic",
-            return_value=mock_dataset
+            return_value=mock_dataset,
         ):
             result = await loader.load_all()