diff --git a/agentic_security/probe_data/modules/adaptive_attacks.py b/agentic_security/probe_data/modules/adaptive_attacks.py index 6263d17..208d6a5 100644 --- a/agentic_security/probe_data/modules/adaptive_attacks.py +++ b/agentic_security/probe_data/modules/adaptive_attacks.py @@ -26,7 +26,6 @@ class Module: goals: List of corresponding goals for each target topic. """ - def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}): r = httpx.get(url) diff --git a/agentic_security/probe_data/modules/fine_tuned.py b/agentic_security/probe_data/modules/fine_tuned.py index 49eef4a..38732b5 100644 --- a/agentic_security/probe_data/modules/fine_tuned.py +++ b/agentic_security/probe_data/modules/fine_tuned.py @@ -46,9 +46,7 @@ class Module: logger.error("No prompts retrieved from the API.") return logger.info(f"Retrieved {len(prompts)} prompts.") - for i, prompt in enumerate( - prompts[: self.max_prompts] - ): + for i, prompt in enumerate(prompts[: self.max_prompts]): logger.info(f"Processing prompt {i+1}/{len(prompts)}: {prompt}") yield prompt while not self.tools_inbox.empty(): @@ -94,4 +92,4 @@ class Module: return data.get("prompts", []) except httpx.RequestError as e: logger.error(f"Failed to fetch prompts: {e}") - return [] + return [] diff --git a/agentic_security/probe_data/modules/module_protocol.py b/agentic_security/probe_data/modules/module_protocol.py index 30355ba..1f39351 100644 --- a/agentic_security/probe_data/modules/module_protocol.py +++ b/agentic_security/probe_data/modules/module_protocol.py @@ -1,4 +1,4 @@ -""":mod:`module_protocol` -- Base protocol for probe data modules. +""":mod:`module_protocol` -- Base protocol for probe data modules. Defines the abstract Protocol that all probe data modules must implement, providing a standardized interface for module initialization and execution. @@ -9,9 +9,11 @@ See Also: :mod:`agentic_security.probe_data.modules.inspect_ai_tool` :mod:`agentic_security.probe_data.modules.rl_model` :doc:`/external_module` -""""" +""" -from typing import Protocol, Any, AsyncGenerator, runtime_checkable +import asyncio +from collections.abc import AsyncGenerator +from typing import Any, Protocol, runtime_checkable @runtime_checkable @@ -37,10 +39,10 @@ class ModuleProtocol(Protocol): tools_inbox: asyncio.Queue opts: dict - async def apply(self) -> AsyncGenerator[str, None]: + async def apply(self) -> AsyncGenerator[str]: """Execute the module and yield result messages. Yields: str: Result messages generated during module execution. """ - ... \ No newline at end of file + ...