fix: pc

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)
 

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 []

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.
         """
-        ...
+        ...