kailash 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

kailash/workflow/builder.py

@@ -2,11 +2,17 @@
 
 import logging
 import uuid
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional, Union
 
-from kailash.nodes.base import Node
+from kailash.nodes.base import Node, NodeRegistry
 from kailash.sdk_exceptions import ConnectionError, WorkflowValidationError
+from kailash.workflow.contracts import ConnectionContract, get_contract_registry
 from kailash.workflow.graph import Workflow
+from kailash.workflow.validation import (
+    IssueSeverity,
+    ParameterDeclarationValidator,
+    ValidationIssue,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -14,8 +20,12 @@ logger = logging.getLogger(__name__)
 class WorkflowBuilder:
     """Builder pattern for creating Workflow instances."""
 
-    def __init__(self):
-        """Initialize an empty workflow builder."""
+    def __init__(self, edge_config: dict[str, Any] | None = None):
+        """Initialize an empty workflow builder.
+
+        Args:
+            edge_config: Optional edge infrastructure configuration
+        """
         self.nodes: dict[str, dict[str, Any]] = {}
         self.connections: list[dict[str, str]] = []
         self._metadata: dict[str, Any] = {}
@@ -23,6 +33,169 @@ class WorkflowBuilder:
         self.workflow_parameters: dict[str, Any] = {}
         self.parameter_mappings: dict[str, dict[str, str]] = {}
 
+        # Edge infrastructure support
+        self.edge_config = edge_config
+        self._has_edge_nodes = False
+        self._edge_infrastructure = None
+
+        # Connection contracts support
+        self.connection_contracts: dict[str, ConnectionContract] = {}
+        self._contract_registry = get_contract_registry()
+
+        # Parameter validation support
+        self._param_validator = ParameterDeclarationValidator()
+
+    def _is_sdk_node(self, node_class: type) -> bool:
+        """Detect if node is SDK-provided vs custom implementation.
+
+        SDK nodes are registered in the NodeRegistry via @register_node decorator.
+        Custom nodes are not registered and require class reference usage.
+
+        Args:
+            node_class: The node class to check
+
+        Returns:
+            True if node is registered in SDK (can use string reference),
+            False if custom node (must use class reference)
+        """
+        if not hasattr(node_class, "__name__"):
+            return False
+
+        # Check if the node class is registered in the NodeRegistry
+        try:
+            registered_class = NodeRegistry.get(node_class.__name__)
+            # Check if it's the same class (identity check)
+            return registered_class is node_class
+        except Exception:
+            # Node not found in registry = custom node
+            return False
+
+    def _generate_intelligent_node_warning(self, node_class: type, node_id: str) -> str:
+        """Generate context-aware warnings based on node type.
+
+        Args:
+            node_class: The node class being added
+            node_id: The node ID
+
+        Returns:
+            Appropriate warning message for the node type
+        """
+        if self._is_sdk_node(node_class):
+            # SDK node using class reference - suggest string pattern
+            return (
+                f"SDK node detected. Consider using string reference for better compatibility:\n"
+                f"  CURRENT: add_node({node_class.__name__}, '{node_id}', {{...}})\n"
+                f"  PREFERRED: add_node('{node_class.__name__}', '{node_id}', {{...}})\n"
+                f"String references work for all @register_node() decorated SDK nodes."
+            )
+        else:
+            # Custom node using class reference - this is CORRECT
+            return (
+                f"✅ CUSTOM NODE USAGE CORRECT\n"
+                f"\n"
+                f"Pattern: add_node({node_class.__name__}, '{node_id}', {{...}})\n"
+                f"Status: This is the CORRECT pattern for custom nodes\n"
+                f"\n"
+                f'⚠️  IGNORE "preferred pattern" suggestions for custom nodes\n'
+                f"String references only work for @register_node() decorated SDK nodes.\n"
+                f"Custom nodes MUST use class references as shown above.\n"
+                f"\n"
+                f"📚 Guide: sdk-users/7-gold-standards/GOLD-STANDARD-custom-node-development-guide.md"
+            )
+
+    def validate_parameter_declarations(
+        self, warn_on_issues: bool = True
+    ) -> list[ValidationIssue]:
+        """Validate parameter declarations for all nodes in the workflow.
+
+        This method detects common parameter declaration issues that lead to
+        silent parameter dropping and debugging difficulties.
+
+        Args:
+            warn_on_issues: Whether to log warnings for detected issues
+
+        Returns:
+            List of ValidationIssue objects for any problems found
+        """
+        all_issues = []
+
+        for node_id, node_info in self.nodes.items():
+            try:
+                # Create a temporary instance to validate parameter declarations
+                if "instance" in node_info:
+                    # Use existing instance
+                    node_instance = node_info["instance"]
+                    workflow_params = {}  # Instance already has config
+                elif "class" in node_info:
+                    # Create temporary instance of custom node
+                    node_class = node_info["class"]
+                    node_config = node_info.get("config", {})
+                    # Create minimal instance just for parameter validation
+                    try:
+                        node_instance = node_class(**node_config)
+                        workflow_params = node_config
+                    except Exception as e:
+                        # If we can't create instance, skip detailed validation
+                        all_issues.append(
+                            ValidationIssue(
+                                severity=IssueSeverity.WARNING,
+                                category="parameter_declaration",
+                                code="PAR005",
+                                message=f"Could not validate parameters for custom node '{node_id}': {e}",
+                                suggestion="Ensure node constructor accepts provided configuration parameters",
+                                node_id=node_id,
+                            )
+                        )
+                        continue
+                else:
+                    # SDK node - validate if we can create it
+                    node_type = node_info["type"]
+                    node_config = node_info.get("config", {})
+                    try:
+                        # Try to get the class from registry
+                        node_class = NodeRegistry.get(node_type)
+                        node_instance = node_class(**node_config)
+                        workflow_params = node_config
+                    except Exception:
+                        # Skip validation for nodes we can't instantiate
+                        continue
+
+                # Validate parameter declarations
+                issues = self._param_validator.validate_node_parameters(
+                    node_instance, workflow_params
+                )
+
+                # Add node_id to issues
+                for issue in issues:
+                    issue.node_id = node_id
+                    all_issues.append(issue)
+
+                    # Log warnings if requested
+                    if warn_on_issues:
+                        if issue.severity == IssueSeverity.ERROR:
+                            logger.error(
+                                f"Parameter validation error in node '{node_id}': {issue.message}"
+                            )
+                        elif issue.severity == IssueSeverity.WARNING:
+                            logger.warning(
+                                f"Parameter validation warning in node '{node_id}': {issue.message}"
+                            )
+
+            except Exception as e:
+                # General validation error
+                all_issues.append(
+                    ValidationIssue(
+                        severity=IssueSeverity.WARNING,
+                        category="parameter_declaration",
+                        code="PAR006",
+                        message=f"Parameter validation failed for node '{node_id}': {e}",
+                        suggestion="Check node configuration and parameter declarations",
+                        node_id=node_id,
+                    )
+                )
+
+        return all_issues
+
     def add_node(self, *args, **kwargs) -> str:
         """
         Unified add_node method supporting multiple API patterns.
@@ -91,6 +264,9 @@ class WorkflowBuilder:
                 # Two strings - assume current API: add_node("NodeType", "node_id")
                 config = kwargs if kwargs else (args[2] if len(args) > 2 else {})
                 return self._add_node_current(args[0], args[1], config)
+            elif isinstance(args[1], dict):
+                # Pattern: add_node("NodeType", {config}) - treat as add_node("NodeType", None, {config})
+                return self._add_node_current(args[0], None, args[1])
             else:
                 # Invalid second argument
                 raise WorkflowValidationError(
@@ -188,10 +364,10 @@ class WorkflowBuilder:
         if node_id is None:
             node_id = f"node_{uuid.uuid4().hex[:8]}"
 
+        # Generate context-aware warning based on node type
+        warning_message = self._generate_intelligent_node_warning(node_class, node_id)
         warnings.warn(
-            f"Alternative API usage detected. Consider using preferred pattern:\n"
-            f"  CURRENT: add_node({node_class.__name__}, '{node_id}', {list(config.keys())})\n"
-            f"  PREFERRED: add_node('{node_class.__name__}', '{node_id}', {config})",
+            warning_message,
             UserWarning,
             stacklevel=3,
         )
@@ -287,8 +463,52 @@ class WorkflowBuilder:
             )
 
         logger.info(f"Added node '{node_id}' of type '{type_name}'")
+
+        # Detect edge nodes
+        if self._is_edge_node(type_name):
+            self._has_edge_nodes = True
+            logger.debug(f"Detected edge node: {type_name}")
+
         return node_id
 
+    def _is_edge_node(self, node_type: str) -> bool:
+        """Check if a node type is an edge node.
+
+        Args:
+            node_type: The node type to check
+
+        Returns:
+            True if the node is an edge node
+        """
+        # Use the same logic as EdgeInfrastructure if available
+        if self._edge_infrastructure:
+            return self._edge_infrastructure.is_edge_node(node_type)
+
+        # Otherwise use local logic
+        # Check exact matches and subclasses
+        edge_prefixes = ["Edge", "edge"]
+        edge_suffixes = [
+            "EdgeNode",
+            "EdgeDataNode",
+            "EdgeStateMachine",
+            "EdgeCacheNode",
+        ]
+
+        # Exact match
+        if node_type in edge_suffixes:
+            return True
+
+        # Check if it starts with Edge/edge
+        for prefix in edge_prefixes:
+            if node_type.startswith(prefix):
+                return True
+
+        # Check if it ends with EdgeNode (for custom edge nodes)
+        if node_type.endswith("EdgeNode"):
+            return True
+
+        return False
+
     # Fluent API methods for backward compatibility
     def add_node_fluent(
         self, node_id: str, node_class_or_type: Any, **config
@@ -539,6 +759,142 @@ class WorkflowBuilder:
         self._metadata.update(kwargs)
         return self
 
+    def add_typed_connection(
+        self,
+        from_node: str,
+        from_output: str,
+        to_node: str,
+        to_input: str,
+        contract: Union[str, ConnectionContract],
+        validate_immediately: bool = False,
+    ) -> "WorkflowBuilder":
+        """
+        Add a typed connection with contract validation.
+
+        This is the new contract-based connection method that enforces
+        validation contracts on data flowing between nodes.
+
+        Args:
+            from_node: Source node ID
+            from_output: Output field from source
+            to_node: Target node ID
+            to_input: Input field on target
+            contract: Contract name (string) or ConnectionContract instance
+            validate_immediately: Whether to validate contract definitions now
+
+        Returns:
+            Self for chaining
+
+        Raises:
+            WorkflowValidationError: If contract is invalid or nodes don't exist
+            ConnectionError: If connection setup fails
+
+        Example:
+            # Using predefined contract
+            workflow.add_typed_connection(
+                "csv_reader", "data", "processor", "input_data",
+                contract="string_data"
+            )
+
+            # Using custom contract
+            custom_contract = ConnectionContract(
+                name="user_data_flow",
+                source_schema={"type": "object", "properties": {"id": {"type": "string"}}},
+                target_schema={"type": "object", "properties": {"id": {"type": "string"}}},
+                security_policies=[SecurityPolicy.NO_PII]
+            )
+            workflow.add_typed_connection(
+                "user_source", "user", "user_processor", "user_data",
+                contract=custom_contract
+            )
+        """
+        # Resolve contract
+        if isinstance(contract, str):
+            contract_obj = self._contract_registry.get(contract)
+            if not contract_obj:
+                available_contracts = self._contract_registry.list_contracts()
+                raise WorkflowValidationError(
+                    f"Contract '{contract}' not found. Available contracts: {available_contracts}"
+                )
+            contract = contract_obj
+
+        # Add the standard connection first
+        self.add_connection(from_node, from_output, to_node, to_input)
+
+        # Store the contract for this connection
+        connection_id = f"{from_node}.{from_output} → {to_node}.{to_input}"
+        self.connection_contracts[connection_id] = contract
+
+        # Immediate validation if requested
+        if validate_immediately:
+            # Validate that contract schemas are valid
+            try:
+                if contract.source_schema:
+                    from jsonschema import Draft7Validator
+
+                    Draft7Validator.check_schema(contract.source_schema)
+                if contract.target_schema:
+                    Draft7Validator.check_schema(contract.target_schema)
+            except Exception as e:
+                raise WorkflowValidationError(
+                    f"Invalid contract schema for connection {connection_id}: {e}"
+                )
+
+        logger.info(
+            f"Added typed connection '{connection_id}' with contract '{contract.name}'"
+        )
+
+        return self
+
+    def get_connection_contract(
+        self, connection_id: str
+    ) -> Optional[ConnectionContract]:
+        """
+        Get the contract for a specific connection.
+
+        Args:
+            connection_id: Connection identifier in format "from.output → to.input"
+
+        Returns:
+            ConnectionContract if found, None otherwise
+        """
+        return self.connection_contracts.get(connection_id)
+
+    def list_connection_contracts(self) -> dict[str, str]:
+        """
+        List all connection contracts in this workflow.
+
+        Returns:
+            Dict mapping connection IDs to contract names
+        """
+        return {
+            conn_id: contract.name
+            for conn_id, contract in self.connection_contracts.items()
+        }
+
+    def validate_all_contracts(self) -> tuple[bool, list[str]]:
+        """
+        Validate all connection contracts in the workflow.
+
+        Returns:
+            Tuple of (all_valid, list_of_errors)
+        """
+        errors = []
+
+        for connection_id, contract in self.connection_contracts.items():
+            try:
+                # Validate contract schemas
+                if contract.source_schema:
+                    from jsonschema import Draft7Validator
+
+                    Draft7Validator.check_schema(contract.source_schema)
+                if contract.target_schema:
+                    Draft7Validator.check_schema(contract.target_schema)
+            except Exception as e:
+                errors.append(f"Contract '{contract.name}' for {connection_id}: {e}")
+
+        return len(errors) == 0, errors
+
     def add_workflow_inputs(
         self, input_node_id: str, input_mappings: dict
     ) -> "WorkflowBuilder":
@@ -623,6 +979,13 @@ class WorkflowBuilder:
         version = metadata.pop("version", "1.0.0")
         author = metadata.pop("author", "")
 
+        # Initialize edge infrastructure if needed
+        if self._has_edge_nodes and not self._edge_infrastructure:
+            from kailash.workflow.edge_infrastructure import EdgeInfrastructure
+
+            self._edge_infrastructure = EdgeInfrastructure(self.edge_config)
+            logger.info("Initialized edge infrastructure for workflow")
+
         # Create workflow
         workflow = Workflow(
             workflow_id=workflow_id,
@@ -633,6 +996,27 @@ class WorkflowBuilder:
             metadata=metadata,
         )
 
+        # Store edge infrastructure reference in workflow metadata if present
+        if self._edge_infrastructure:
+            workflow.metadata["_edge_infrastructure"] = self._edge_infrastructure
+
+        # Validate parameter declarations before building workflow
+        param_issues = self.validate_parameter_declarations(warn_on_issues=True)
+
+        # Check for critical parameter errors that should block workflow creation
+        critical_errors = [
+            issue for issue in param_issues if issue.severity == IssueSeverity.ERROR
+        ]
+        if critical_errors:
+            error_messages = [
+                f"{issue.node_id}: {issue.message}" for issue in critical_errors
+            ]
+            raise WorkflowValidationError(
+                "Cannot build workflow due to parameter declaration errors:\n"
+                + "\n".join(f"  - {msg}" for msg in error_messages)
+                + "\n\nSee: sdk-users/7-gold-standards/enterprise-parameter-passing-gold-standard.md"
+            )
+
         # Add nodes to workflow
         for node_id, node_info in self.nodes.items():
             try:
@@ -645,6 +1029,16 @@ class WorkflowBuilder:
                     # Node class was provided
                     node_class = node_info["class"]
                     node_config = node_info.get("config", {})
+
+                    # Inject edge infrastructure if this is an edge node
+                    if self._edge_infrastructure and self._is_edge_node(
+                        node_class.__name__
+                    ):
+                        node_config["_edge_infrastructure"] = self._edge_infrastructure
+                        logger.debug(
+                            f"Injected edge infrastructure into {node_class.__name__}"
+                        )
+
                     workflow.add_node(
                         node_id=node_id, node_or_type=node_class, **node_config
                     )
@@ -652,6 +1046,12 @@ class WorkflowBuilder:
                     # String node type
                     node_type = node_info["type"]
                     node_config = node_info.get("config", {})
+
+                    # Inject edge infrastructure if this is an edge node
+                    if self._edge_infrastructure and self._is_edge_node(node_type):
+                        node_config["_edge_infrastructure"] = self._edge_infrastructure
+                        logger.debug(f"Injected edge infrastructure into {node_type}")
+
                     workflow.add_node(
                         node_id=node_id, node_or_type=node_type, **node_config
                     )
@@ -711,9 +1111,13 @@ class WorkflowBuilder:
                                             self.workflow_parameters[workflow_param]
                                         )
 
-        # Store workflow parameters in metadata for runtime reference
+        # Store workflow parameters and contracts in metadata for runtime reference
         workflow.metadata["workflow_parameters"] = self.workflow_parameters
         workflow.metadata["parameter_mappings"] = self.parameter_mappings
+        workflow.metadata["connection_contracts"] = {
+            conn_id: contract.to_dict()
+            for conn_id, contract in self.connection_contracts.items()
+        }
 
         logger.info(
             f"Built workflow '{workflow_id}' with "
@@ -789,6 +1193,7 @@ class WorkflowBuilder:
         self._metadata = {}
         self.workflow_parameters = {}
         self.parameter_mappings = {}
+        self.connection_contracts = {}
         return self
 
     @classmethod