kailash 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

kailash/nodes/mixins/mcp.py

@@ -5,7 +5,7 @@ allowing them to discover and use MCP tools without being an LLM agent.
 """
 
 import asyncio
-from typing import Any, Dict, List, Union
+from typing import Any
 
 from kailash.mcp import MCPClient
 
@@ -50,8 +50,8 @@ class MCPCapabilityMixin:
         return self._mcp_client
 
     async def discover_mcp_tools(
-        self, mcp_servers: List[Union[str, Dict[str, Any]]]
-    ) -> List[Dict[str, Any]]:
+        self, mcp_servers: list[str | dict[str, Any]]
+    ) -> list[dict[str, Any]]:
         """Discover tools from MCP servers asynchronously.
 
         Args:
@@ -75,9 +75,9 @@ class MCPCapabilityMixin:
 
     async def call_mcp_tool(
         self,
-        server_config: Union[str, Dict[str, Any]],
+        server_config: str | dict[str, Any],
         tool_name: str,
-        arguments: Dict[str, Any],
+        arguments: dict[str, Any],
     ) -> Any:
         """Call an MCP tool asynchronously.
 
@@ -92,8 +92,8 @@ class MCPCapabilityMixin:
         return await self.mcp_client.call_tool(server_config, tool_name, arguments)
 
     async def list_mcp_resources(
-        self, server_config: Union[str, Dict[str, Any]]
-    ) -> List[Dict[str, Any]]:
+        self, server_config: str | dict[str, Any]
+    ) -> list[dict[str, Any]]:
         """List available resources from an MCP server.
 
         Args:
@@ -105,7 +105,7 @@ class MCPCapabilityMixin:
         return await self.mcp_client.list_resources(server_config)
 
     async def read_mcp_resource(
-        self, server_config: Union[str, Dict[str, Any]], uri: str
+        self, server_config: str | dict[str, Any], uri: str
     ) -> Any:
         """Read a resource from an MCP server.
 
@@ -121,8 +121,8 @@ class MCPCapabilityMixin:
     # Synchronous wrappers for non-async nodes
 
     def discover_mcp_tools_sync(
-        self, mcp_servers: List[Union[str, Dict[str, Any]]]
-    ) -> List[Dict[str, Any]]:
+        self, mcp_servers: list[str | dict[str, Any]]
+    ) -> list[dict[str, Any]]:
         """Synchronous wrapper for discovering MCP tools.
 
         Args:
@@ -139,9 +139,9 @@ class MCPCapabilityMixin:
 
     def call_mcp_tool_sync(
         self,
-        server_config: Union[str, Dict[str, Any]],
+        server_config: str | dict[str, Any],
         tool_name: str,
-        arguments: Dict[str, Any],
+        arguments: dict[str, Any],
     ) -> Any:
         """Synchronous wrapper for calling MCP tools.
 
@@ -162,8 +162,8 @@ class MCPCapabilityMixin:
             loop.close()
 
     def list_mcp_resources_sync(
-        self, server_config: Union[str, Dict[str, Any]]
-    ) -> List[Dict[str, Any]]:
+        self, server_config: str | dict[str, Any]
+    ) -> list[dict[str, Any]]:
         """Synchronous wrapper for listing MCP resources.
 
         Args:
@@ -179,7 +179,7 @@ class MCPCapabilityMixin:
             loop.close()
 
     def read_mcp_resource_sync(
-        self, server_config: Union[str, Dict[str, Any]], uri: str
+        self, server_config: str | dict[str, Any], uri: str
     ) -> Any:
         """Synchronous wrapper for reading MCP resources.
 
@@ -198,7 +198,7 @@ class MCPCapabilityMixin:
 
     # Helper methods for common patterns
 
-    def get_mcp_parameter_defaults(self) -> Dict[str, Any]:
+    def get_mcp_parameter_defaults(self) -> dict[str, Any]:
         """Get default MCP-related parameters for nodes.
 
         Returns:
@@ -206,7 +206,7 @@ class MCPCapabilityMixin:
         """
         return {"mcp_servers": [], "mcp_context": [], "auto_discover_tools": False}
 
-    def format_mcp_tools_for_display(self, tools: List[Dict[str, Any]]) -> str:
+    def format_mcp_tools_for_display(self, tools: list[dict[str, Any]]) -> str:
         """Format MCP tools for human-readable display.
 
         Args:

kailash/nodes/mixins.py

@@ -12,7 +12,7 @@ Design Philosophy:
 """
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any
 
 from kailash.security import (
     SecurityConfig,
@@ -43,9 +43,7 @@ class SecurityMixin:
                 return self.process_safely(safe_params)
     """
 
-    def __init__(
-        self, *args, security_config: Optional[SecurityConfig] = None, **kwargs
-    ):
+    def __init__(self, *args, security_config: SecurityConfig | None = None, **kwargs):
         """
         Initialize security mixin.
 
@@ -60,7 +58,7 @@ class SecurityMixin:
         if self.security_config.enable_audit_logging:
             logger.info(f"Security mixin initialized for {self.__class__.__name__}")
 
-    def validate_and_sanitize_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+    def validate_and_sanitize_inputs(self, inputs: dict[str, Any]) -> dict[str, Any]:
         """
         Validate and sanitize input parameters.
 
@@ -151,7 +149,7 @@ class ValidationMixin:
     """
 
     def validate_required_params(
-        self, inputs: Dict[str, Any], required_params: list
+        self, inputs: dict[str, Any], required_params: list
     ) -> None:
         """
         Validate that all required parameters are present.
@@ -168,8 +166,8 @@ class ValidationMixin:
             raise ValueError(f"Missing required parameters: {missing_params}")
 
     def validate_param_types(
-        self, inputs: Dict[str, Any], type_mapping: Dict[str, type]
-    ) -> Dict[str, Any]:
+        self, inputs: dict[str, Any], type_mapping: dict[str, type]
+    ) -> dict[str, Any]:
         """
         Validate and convert parameter types.
 
@@ -203,7 +201,7 @@ class ValidationMixin:
         return converted
 
     def validate_param_ranges(
-        self, inputs: Dict[str, Any], range_mapping: Dict[str, tuple]
+        self, inputs: dict[str, Any], range_mapping: dict[str, tuple]
     ) -> None:
         """
         Validate that numeric parameters are within acceptable ranges.
@@ -284,7 +282,7 @@ class PerformanceMixin:
 
         return wrapper
 
-    def get_performance_stats(self) -> Dict[str, Any]:
+    def get_performance_stats(self) -> dict[str, Any]:
         """
         Get performance statistics for this node.
 

kailash/nodes/transform/chunkers.py

@@ -1,6 +1,6 @@
 """Document chunking nodes for splitting text into manageable pieces."""
 
-from typing import Any, Dict
+from typing import Any
 
 from kailash.nodes.base import Node, NodeParameter, register_node
 
@@ -9,7 +9,7 @@ from kailash.nodes.base import Node, NodeParameter, register_node
 class HierarchicalChunkerNode(Node):
     """Splits documents into hierarchical chunks for better retrieval."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "documents": NodeParameter(
                 name="documents",
@@ -33,7 +33,7 @@ class HierarchicalChunkerNode(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         documents = kwargs.get("documents", [])
         chunk_size = kwargs.get("chunk_size", 200)
         # overlap = kwargs.get("overlap", 50)  # Currently not used in chunking logic

kailash/nodes/transform/formatters.py

@@ -1,6 +1,6 @@
 """Text formatting nodes for transforming and preparing text data."""
 
-from typing import Any, Dict
+from typing import Any
 
 from kailash.nodes.base import Node, NodeParameter, register_node
 
@@ -9,7 +9,7 @@ from kailash.nodes.base import Node, NodeParameter, register_node
 class ChunkTextExtractorNode(Node):
     """Extracts text content from chunks for embedding generation."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "chunks": NodeParameter(
                 name="chunks",
@@ -19,7 +19,7 @@ class ChunkTextExtractorNode(Node):
             )
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         chunks = kwargs.get("chunks", [])
         # Extract just the content text from chunks
         texts = [chunk["content"] for chunk in chunks]
@@ -30,7 +30,7 @@ class ChunkTextExtractorNode(Node):
 class QueryTextWrapperNode(Node):
     """Wraps query string in list for embedding generation."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "query": NodeParameter(
                 name="query",
@@ -40,7 +40,7 @@ class QueryTextWrapperNode(Node):
             )
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         query = kwargs.get("query", "")
         print(f"Debug QueryTextWrapper: received query='{query}'")
         # Use input_texts for batch embedding (single item list)
@@ -53,7 +53,7 @@ class QueryTextWrapperNode(Node):
 class ContextFormatterNode(Node):
     """Formats relevant chunks into context for LLM."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "relevant_chunks": NodeParameter(
                 name="relevant_chunks",
@@ -69,7 +69,7 @@ class ContextFormatterNode(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         relevant_chunks = kwargs.get("relevant_chunks", [])
         query = kwargs.get("query", "")
         # Format context from relevant chunks

kailash/nodes/transform/processors.py

@@ -1,7 +1,7 @@
 """Transform nodes for data processing."""
 
 import traceback
-from typing import Any, Dict
+from typing import Any
 
 from kailash.nodes.base import Node, NodeParameter, register_node
 
@@ -119,7 +119,7 @@ class FilterNode(Node):
         >>> assert len(result["filtered_data"]) == 2
     """
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "data": NodeParameter(
                 name="data",
@@ -148,7 +148,7 @@ class FilterNode(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         data = kwargs["data"]
         field = kwargs.get("field")
         operator = kwargs.get("operator", "==")
@@ -220,7 +220,7 @@ class FilterNode(Node):
 class Map(Node):
     """Maps data using a transformation."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "data": NodeParameter(
                 name="data",
@@ -255,7 +255,7 @@ class Map(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         data = kwargs["data"]
         field = kwargs.get("field")
         new_field = kwargs.get("new_field")
@@ -307,11 +307,11 @@ class DataTransformer(Node):
     or other Python code as strings. These are compiled and executed against the input data.
     """
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "data": NodeParameter(
                 name="data",
-                type=list,
+                type=Any,
                 required=False,
                 description="Primary input data to transform",
             ),
@@ -332,7 +332,7 @@ class DataTransformer(Node):
             },  # Support for up to 5 additional arguments
         }
 
-    def validate_inputs(self, **kwargs) -> Dict[str, Any]:
+    def validate_inputs(self, **kwargs) -> dict[str, Any]:
         """Override validate_inputs to accept arbitrary parameters for transformations.
 
         DataTransformer needs to accept any input parameters that might be mapped
@@ -351,7 +351,7 @@ class DataTransformer(Node):
 
         return validated
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         # Extract the transformation functions
         transformations = kwargs.get("transformations", [])
         if not transformations:
@@ -484,7 +484,7 @@ class DataTransformer(Node):
 class Sort(Node):
     """Sorts data."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "data": NodeParameter(
                 name="data",
@@ -507,7 +507,7 @@ class Sort(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         data = kwargs["data"]
         field = kwargs.get("field")
         reverse = kwargs.get("reverse", False)

kailash/runtime/access_controlled.py

@@ -27,7 +27,7 @@ Example with access control (opt-in):
 """
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 from kailash.access_control import (
     AccessControlManager,
@@ -108,7 +108,7 @@ class AccessControlledRuntime:
     """
 
     def __init__(
-        self, user_context: UserContext, base_runtime: Optional[LocalRuntime] = None
+        self, user_context: UserContext, base_runtime: LocalRuntime | None = None
     ):
         """
         Initialize access-controlled runtime.
@@ -123,11 +123,11 @@ class AccessControlledRuntime:
 
         # Track skipped nodes for alternative routing
         self._skipped_nodes: set[str] = set()
-        self._node_outputs: Dict[str, Any] = {}
+        self._node_outputs: dict[str, Any] = {}
 
     def execute(
-        self, workflow: Workflow, parameters: Optional[Dict[str, Any]] = None
-    ) -> Tuple[Any, str]:
+        self, workflow: Workflow, parameters: dict[str, Any] | None = None
+    ) -> tuple[Any, str]:
         """
         Execute workflow with access control.
 
@@ -272,7 +272,7 @@ class AccessControlledRuntime:
         return wrapper
 
     @staticmethod
-    def _mask_fields(data: Dict[str, Any], fields: List[str]) -> Dict[str, Any]:
+    def _mask_fields(data: dict[str, Any], fields: list[str]) -> dict[str, Any]:
         """Mask sensitive fields in data"""
         masked = data.copy()
         for field in fields:
@@ -281,8 +281,8 @@ class AccessControlledRuntime:
         return masked
 
     def _handle_conditional_routing(
-        self, node_id: str, true_path: List[str], false_path: List[str]
-    ) -> List[str]:
+        self, node_id: str, true_path: list[str], false_path: list[str]
+    ) -> list[str]:
         """
         Determine which path to take based on permissions.
 
@@ -345,14 +345,14 @@ class AccessControlConfig:
     """
 
     def __init__(self):
-        self.rules: List[PermissionRule] = []
+        self.rules: list[PermissionRule] = []
 
     def add_workflow_permission(
         self,
         workflow_id: str,
         permission: WorkflowPermission,
-        user_id: Optional[str] = None,
-        role: Optional[str] = None,
+        user_id: str | None = None,
+        role: str | None = None,
         effect: PermissionEffect = PermissionEffect.ALLOW,
     ):
         """Add a workflow-level permission rule"""
@@ -372,11 +372,11 @@ class AccessControlConfig:
         workflow_id: str,
         node_id: str,
         permission: NodePermission,
-        user_id: Optional[str] = None,
-        role: Optional[str] = None,
+        user_id: str | None = None,
+        role: str | None = None,
         effect: PermissionEffect = PermissionEffect.ALLOW,
-        masked_fields: Optional[List[str]] = None,
-        redirect_node: Optional[str] = None,
+        masked_fields: list[str] | None = None,
+        redirect_node: str | None = None,
     ):
         """Add a node-level permission rule"""
         rule = PermissionRule(
@@ -406,9 +406,9 @@ class AccessControlConfig:
 def execute_with_access_control(
     workflow: Workflow,
     user_context: UserContext,
-    parameters: Optional[Dict[str, Any]] = None,
-    access_config: Optional[AccessControlConfig] = None,
-) -> Tuple[Any, str]:
+    parameters: dict[str, Any] | None = None,
+    access_config: AccessControlConfig | None = None,
+) -> tuple[Any, str]:
     """
     Convenience function to execute a workflow with access control.
 

kailash/runtime/async_local.py

@@ -6,8 +6,8 @@ database queries, or LLM interactions.
 """
 
 import logging
-from datetime import datetime, timezone
-from typing import Any, Dict, Optional, Tuple
+from datetime import UTC, datetime
+from typing import Any
 
 import networkx as nx
 
@@ -60,9 +60,9 @@ class AsyncLocalRuntime:
     async def execute(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager] = None,
-        parameters: Optional[Dict[str, Dict[str, Any]]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None = None,
+        parameters: dict[str, dict[str, Any]] | None = None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a workflow asynchronously.
 
         Args:
@@ -144,10 +144,10 @@ class AsyncLocalRuntime:
     async def _execute_workflow(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-        parameters: Dict[str, Dict[str, Any]],
-    ) -> Dict[str, Any]:
+        task_manager: TaskManager | None,
+        run_id: str | None,
+        parameters: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
         """Execute the workflow nodes asynchronously.
 
         Args:
@@ -195,7 +195,7 @@ class AsyncLocalRuntime:
                         run_id=run_id,
                         node_id=node_id,
                         node_type=node_instance.__class__.__name__,
-                        started_at=datetime.now(timezone.utc),
+                        started_at=datetime.now(UTC),
                     )
                 except Exception as e:
                     self.logger.warning(
@@ -220,18 +220,16 @@ class AsyncLocalRuntime:
                     task.update_status(TaskStatus.RUNNING)
 
                 # Execute node - check if it supports async execution
-                start_time = datetime.now(timezone.utc)
+                start_time = datetime.now(UTC)
 
                 if isinstance(node_instance, AsyncNode):
                     # Use async execution
                     outputs = await node_instance.execute_async(**inputs)
                 else:
                     # Fall back to synchronous execution
-                    outputs = node_instance.execute(**inputs)
+                    outputs = node_instance.run(**inputs)
 
-                execution_time = (
-                    datetime.now(timezone.utc) - start_time
-                ).total_seconds()
+                execution_time = (datetime.now(UTC) - start_time).total_seconds()
 
                 # Store outputs
                 node_outputs[node_id] = outputs
@@ -245,7 +243,7 @@ class AsyncLocalRuntime:
                     task.update_status(
                         TaskStatus.COMPLETED,
                         result=outputs,
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                         metadata={"execution_time": execution_time},
                     )
 
@@ -262,7 +260,7 @@ class AsyncLocalRuntime:
                     task.update_status(
                         TaskStatus.FAILED,
                         error=str(e),
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                     )
 
                 # Determine if we should continue or stop
@@ -287,9 +285,9 @@ class AsyncLocalRuntime:
         workflow: Workflow,
         node_id: str,
         node_instance: Any,
-        node_outputs: Dict[str, Dict[str, Any]],
-        parameters: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        node_outputs: dict[str, dict[str, Any]],
+        parameters: dict[str, Any],
+    ) -> dict[str, Any]:
         """Prepare inputs for a node execution.
 
         Args: