diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py
index 6293b6a..e8a29cc 100644
--- a/agentic_security/mcp/client.py
+++ b/agentic_security/mcp/client.py
@@ -1,10 +1,10 @@
-
 import asyncio
-from agentic_security.logutils import logger
 
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 
+from agentic_security.logutils import logger
+
 # Create server parameters for stdio connection
 server_params = StdioServerParameters(
     command="python",  # Executable
@@ -15,13 +15,15 @@ server_params = StdioServerParameters(
 
 async def run() -> None:
     try:
-        logger.info("Starting stdio client session with server parameters: %s", server_params)
+        logger.info(
+            "Starting stdio client session with server parameters: %s", server_params
+        )
         async with stdio_client(server_params) as (read, write):
             async with ClientSession(read, write) as session:
                 # Initialize the connection --> connection does not work
                 logger.info("Initializing client session...")
                 await session.initialize()
-                
+
                 # List available prompts, resources, and tools --> no avalialbe tools
                 logger.info("Listing available prompts...")
                 prompts = await session.list_prompts()
diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py
index 66ba25e..5d7351e 100644
--- a/agentic_security/probe_data/data.py
+++ b/agentic_security/probe_data/data.py
@@ -8,7 +8,6 @@ from typing import Any, TypeVar
 import httpx
 import pandas as pd
 from cache_to_disk import cache_to_disk
-from datasets import load_dataset
 
 from agentic_security.logutils import logger
 from agentic_security.probe_data import stenography_fn
@@ -20,6 +19,7 @@ from agentic_security.probe_data.modules import (
     inspect_ai_tool,
     rl_model,
 )
+from datasets import load_dataset
 
 # Type aliases for clarity
 T = TypeVar("T")
diff --git a/tests/test_registry.py b/tests/test_registry.py
index 2cf0cbc..6338e2b 100644
--- a/tests/test_registry.py
+++ b/tests/test_registry.py
@@ -1,7 +1,7 @@
 import pytest
-from datasets import load_dataset
 
 from agentic_security.probe_data import REGISTRY
+from datasets import load_dataset
 
 
 @pytest.mark.slow