diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index af26241..8b0a2d7 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/agentic_security/__init__.py b/agentic_security/__init__.py index b37466a..a996613 100644 --- a/agentic_security/__init__.py +++ b/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"] diff --git a/agentic_security/executor/concurrent.py b/agentic_security/executor/concurrent.py index 821b418..d5bde15 100644 --- a/agentic_security/executor/concurrent.py +++ b/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 diff --git a/agentic_security/probe_actor/refusal.py b/agentic_security/probe_actor/refusal.py index 2d5b077..c02f28a 100644 --- a/agentic_security/probe_actor/refusal.py +++ b/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): diff --git a/agentic_security/probe_data/audio_generator.py b/agentic_security/probe_data/audio_generator.py index 27f00a1..0b7a488 100644 --- a/agentic_security/probe_data/audio_generator.py +++ b/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") diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py index 027e56c..b2f5132 100644 --- a/agentic_security/probe_data/modules/rl_model.py +++ b/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): diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py index 9cb848b..7d37a75 100644 --- a/agentic_security/probe_data/modules/test_rl_model.py +++ b/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]] diff --git a/agentic_security/probe_data/unified_loader.py b/agentic_security/probe_data/unified_loader.py index 1f077d3..834ef4a 100644 --- a/agentic_security/probe_data/unified_loader.py +++ b/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( - default=None, description="File path for CSV sources" - ) - prompt_column: Optional[str] = Field( + path: str | None = Field(default=None, description="File path for CSV sources") + prompt_column: str | None = 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( - default=None, description="URL for remote CSV files" - ) + url: str | None = Field(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( diff --git a/tests/executor/test_circuit_breaker.py b/tests/executor/test_circuit_breaker.py index 4cf4d5b..1c8213a 100644 --- a/tests/executor/test_circuit_breaker.py +++ b/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 diff --git a/tests/executor/test_concurrent.py b/tests/executor/test_concurrent.py index 993e26a..c076028 100644 --- a/tests/executor/test_concurrent.py +++ b/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 diff --git a/tests/executor/test_rate_limiter.py b/tests/executor/test_rate_limiter.py index 0a36607..6e6efb6 100644 --- a/tests/executor/test_rate_limiter.py +++ b/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 diff --git a/tests/probe_actor/test_fuzzer.py b/tests/probe_actor/test_fuzzer.py index 24e6018..f81035d 100644 --- a/tests/probe_actor/test_fuzzer.py +++ b/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 diff --git a/tests/probe_data/test_unified_loader.py b/tests/probe_data/test_unified_loader.py index 8d19ce9..e1d63c8 100644 --- a/tests/probe_data/test_unified_loader.py +++ b/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()