fixing the linter errors from ruff

Signed-off-by: Ambrish Rawat <[email protected]>

Author: ambrishrawat

docs/examples/security/taint_example.py

@@ -3,38 +3,44 @@
 from mellea.security import SecLevel, privileged, SecurityError
 
 # Create tainted content
-tainted_desc = CBlock("Process this sensitive user data", sec_level=SecLevel.tainted_by(None))
+tainted_desc = CBlock(
+    "Process this sensitive user data", sec_level=SecLevel.tainted_by(None)
+)
 
-print(f"Original CBlock is tainted: {tainted_desc.sec_level.is_tainted() if tainted_desc.sec_level else False}")
+print(
+    f"Original CBlock is tainted: {tainted_desc.sec_level.is_tainted() if tainted_desc.sec_level else False}"
+)
 
 # Create session
 session = start_session()
 
 # Use tainted CBlock in session.instruct
 print("Testing session.instruct with tainted CBlock...")
-result = session.instruct(
-    description=tainted_desc, 
-)
+result = session.instruct(description=tainted_desc)
 
 # The result should be tainted
-print(f"Result is tainted: {result.sec_level.is_tainted() if result.sec_level else False}")
+print(
+    f"Result is tainted: {result.sec_level.is_tainted() if result.sec_level else False}"
+)
 if result.sec_level and result.sec_level.is_tainted():
     taint_source = result.sec_level.get_taint_source()
     print(f"Taint source: {taint_source}")
     print("✅ SUCCESS: Taint preserved!")
 else:
     print("❌ FAIL: Result should be tainted but isn't!")
 
+
 # Mock privileged function that requires un-tainted input
 @privileged
 def process_un_tainted_data(data: CBlock) -> str:
     """A function that requires un-tainted input."""
     return f"Processed: {data.value}"
 
+
 print("\nTesting privileged function with tainted result...")
 try:
     # This should raise a SecurityError
     processed = process_un_tainted_data(result)
     print("❌ FAIL: Should have raised SecurityError!")
 except SecurityError as e:
-    print(f"✅ SUCCESS: SecurityError raised - {e}")
+    print(f"✅ SUCCESS: SecurityError raised - {e}")

mellea/backends/litellm.py

@@ -28,6 +28,7 @@
     chat_completion_delta_merge,
     extract_model_tool_requests,
 )
+from mellea.security import taint_sources
 from mellea.stdlib.base import (
     CBlock,
     Component,
@@ -37,7 +38,6 @@
     ModelOutputThunk,
     ModelToolCall,
 )
-from mellea.security import taint_sources
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement
 
@@ -312,11 +312,9 @@ async def _generate_from_chat_context_standard(
 
         # Compute taint sources from action and context
         sources = taint_sources(action, ctx)
-        
+
         output = ModelOutputThunk.from_generation(
-            value=None,
-            taint_sources=sources,
-            meta={}
+            value=None, taint_sources=sources, meta={}
         )
         output._context = linearized_context
         output._action = action

mellea/backends/ollama.py

@@ -25,6 +25,7 @@
 )
 from mellea.helpers.event_loop_helper import _run_async_in_thread
 from mellea.helpers.fancy_logger import FancyLogger
+from mellea.security import taint_sources
 from mellea.stdlib.base import (
     CBlock,
     Component,
@@ -34,7 +35,6 @@
     ModelOutputThunk,
     ModelToolCall,
 )
-from mellea.security import taint_sources
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement
 
@@ -357,11 +357,9 @@ async def generate_from_chat_context(
 
         # Compute taint sources from action and context
         sources = taint_sources(action, ctx)
-        
+
         output = ModelOutputThunk.from_generation(
-            value=None,
-            taint_sources=sources,
-            meta={}
+            value=None, taint_sources=sources, meta={}
         )
         output._context = linearized_context
         output._action = action
@@ -442,9 +440,7 @@ async def generate_from_raw(
             error = None
             if isinstance(response, BaseException):
                 result = ModelOutputThunk.from_generation(
-                    value="",
-                    taint_sources=taint_sources(actions[i], None),
-                    meta={}
+                    value="", taint_sources=taint_sources(actions[i], None), meta={}
                 )
                 error = response
             else:

mellea/backends/openai.py

@@ -46,6 +46,7 @@
     chat_completion_delta_merge,
     extract_model_tool_requests,
 )
+from mellea.security import taint_sources
 from mellea.stdlib.base import (
     CBlock,
     Component,
@@ -55,7 +56,6 @@
     GenerateType,
     ModelOutputThunk,
 )
-from mellea.security import taint_sources
 from mellea.stdlib.chat import Message
 from mellea.stdlib.intrinsics.intrinsic import Intrinsic
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
@@ -642,11 +642,9 @@ async def _generate_from_chat_context_standard(
 
         # Compute taint sources from action and context
         sources = taint_sources(action, ctx)
-        
+
         output = ModelOutputThunk.from_generation(
-            value=None,
-            taint_sources=sources,
-            meta={}
+            value=None, taint_sources=sources, meta={}
         )
         output._context = linearized_context
         output._action = action

mellea/security/__init__.py

@@ -7,19 +7,19 @@
 from .core import (
     AccessType,
     SecLevel,
-    SecurityMetadata,
     SecurityError,
-    privileged,
+    SecurityMetadata,
     declassify,
+    privileged,
     taint_sources,
 )
 
 __all__ = [
     "AccessType",
     "SecLevel",
-    "SecurityMetadata", 
     "SecurityError",
-    "privileged",
+    "SecurityMetadata",
     "declassify",
+    "privileged",
     "taint_sources",
-]
+]