hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/builtin/nodes/function_node.py

@@ -0,0 +1,454 @@
+"""Function node factory for creating function-based pipeline nodes."""
+
+import asyncio
+import importlib
+import inspect
+import time
+from collections.abc import Callable
+from typing import Any, get_type_hints
+
+from pydantic import BaseModel
+
+from hexdag.core.domain.dag import NodeSpec
+from hexdag.core.logging import get_logger
+from hexdag.core.protocols import is_schema_type
+
+from .base_node_factory import BaseNodeFactory
+from .mapped_input import MappedInput
+
+logger = get_logger(__name__)
+
+
+class FunctionNode(BaseNodeFactory):
+    """Simple factory for creating function-based nodes with optional Pydantic validation.
+
+    Function nodes are highly dynamic - the function itself defines configuration via its
+    signature and parameters. No static Config class needed (follows YAGNI principle).
+    All configuration is passed dynamically through __call__() parameters.
+    """
+
+    def __call__(
+        self,
+        name: str,
+        fn: Callable[..., Any] | str,
+        input_schema: dict[str, Any] | type[BaseModel] | None = None,
+        output_schema: dict[str, Any] | type[BaseModel] | None = None,
+        deps: list[str] | None = None,
+        input_mapping: dict[str, str] | None = None,
+        unpack_input: bool = False,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a NodeSpec for a function-based node.
+
+        Args:
+        ----
+            name: Node name
+            fn: Function to execute (callable or module path string like 'mymodule.myfunc')
+            input_schema: Input schema for validation (if None, inferred from function)
+            output_schema: Output schema for validation (if None, inferred from function)
+            deps: List of dependency node names
+            input_mapping: Optional field mapping dict {target_field: "source.path"}
+            unpack_input: If True, unpack input_mapping fields as individual **kwargs
+                to the function instead of passing as single input_data object.
+                This allows functions with signatures like `fn(load_id, rate, *, db=None)`
+                instead of `fn(input_data, *, db=None)`.
+            **kwargs: Additional parameters
+
+        Returns
+        -------
+        NodeSpec
+            Complete node specification ready for execution
+        """
+        # Resolve function from string path if needed
+        resolved_fn = self._resolve_function(fn)
+
+        # Validate function can be used properly
+        self._validate_function(resolved_fn, unpack_input=unpack_input)
+
+        if input_mapping is not None:
+            kwargs["input_mapping"] = input_mapping
+
+        if unpack_input:
+            kwargs["unpack_input"] = True
+
+        if input_mapping and not input_schema:
+            # Auto-generate Pydantic model from field mapping
+            input_schema = MappedInput.create_model(
+                f"{name}MappedInput",
+                input_mapping,
+                dependency_models=None,  # Could enhance with dependency introspection
+            )
+
+        # Infer schemas from function annotations if not provided
+        if input_schema is None or output_schema is None:
+            inferred_input, inferred_output = self._infer_schemas_from_function(resolved_fn)
+            input_schema = input_schema or inferred_input
+            output_schema = output_schema or inferred_output
+
+        # For basic types like dict, list, str, etc., use them directly
+
+        if isinstance(input_schema, type) and input_schema.__name__ in {
+            "dict",
+            "list",
+            "str",
+            "int",
+            "float",
+            "bool",
+        }:
+            input_model: type[BaseModel] | type | None = input_schema
+        else:
+            input_model = self.create_pydantic_model(f"{name}Input", input_schema)
+
+        if isinstance(output_schema, type) and output_schema.__name__ in {
+            "dict",
+            "list",
+            "str",
+            "int",
+            "float",
+            "bool",
+        }:
+            output_model: type[BaseModel] | type | None = output_schema
+        else:
+            output_model = self.create_pydantic_model(f"{name}Output", output_schema)
+
+        wrapped_fn = self._create_wrapped_function(
+            name, resolved_fn, input_model, output_model, unpack_input=unpack_input
+        )
+
+        # Extract framework-level parameters from kwargs
+        framework = self.extract_framework_params(kwargs)
+
+        return NodeSpec(
+            name=name,
+            fn=wrapped_fn,
+            in_model=input_model,
+            out_model=output_model,
+            deps=frozenset(deps or []),
+            params=kwargs,
+            timeout=framework["timeout"],
+            max_retries=framework["max_retries"],
+            when=framework["when"],
+        )
+
+    def _create_wrapped_function(
+        self,
+        name: str,
+        fn: Callable[..., Any],
+        input_model: type[BaseModel] | type | None,
+        output_model: type[BaseModel] | type | None,
+        *,
+        unpack_input: bool = False,
+    ) -> Callable[..., Any]:
+        """Create a simple wrapped function with explicit port handling.
+
+        Parameters
+        ----------
+        name : str
+            Node name for logging
+        fn : Callable[..., Any]
+            The function to wrap
+        input_model : type[BaseModel] | type | None
+            Input model for validation
+        output_model : type[BaseModel] | type | None
+            Output model for validation
+        unpack_input : bool, default=False
+            If True, unpack input_data fields as individual **kwargs to the function
+            instead of passing as single input_data object.
+
+        Returns
+        -------
+        Callable[..., Any]
+            Wrapped function that handles orchestrator integration
+        """
+        # Analyze function signature once
+        sig = inspect.signature(fn)
+        accepts_kwargs = any(p.kind == p.VAR_KEYWORD for p in sig.parameters.values())
+        param_names = set(sig.parameters.keys())
+
+        # Get function name for logging
+        fn_name = getattr(fn, "__name__", "anonymous")
+
+        async def wrapped_fn(input_data: Any, **ports: Any) -> Any:
+            """Execute function with explicit port handling."""
+            node_logger = logger.bind(node=name, node_type="function_node")
+
+            # Log input details at debug level
+            if isinstance(input_data, dict):
+                node_logger.debug(
+                    "Input received",
+                    input_keys=list(input_data.keys()),
+                    input_key_count=len(input_data),
+                )
+            else:
+                node_logger.debug(
+                    "Input received",
+                    input_type=type(input_data).__name__,
+                )
+
+            # Log execution start
+            node_logger.info("Executing function", fn_name=fn_name)
+
+            start_time = time.perf_counter()
+
+            # Prepare function call arguments
+            if accepts_kwargs:
+                # Function accepts **kwargs, pass all ports
+                call_kwargs = ports
+            else:
+                # Function has specific parameters, only pass ports that match parameter names
+                call_kwargs = {k: v for k, v in ports.items() if k in param_names}
+
+            try:
+                # Execute function (handle both sync and async)
+                if unpack_input:
+                    # Unpack input_data fields as individual kwargs
+                    # This allows functions with signatures like fn(load_id, rate, *, db=None)
+                    if isinstance(input_data, dict):
+                        unpacked_kwargs = {**input_data, **call_kwargs}
+                    elif isinstance(input_data, BaseModel):
+                        unpacked_kwargs = {**input_data.model_dump(), **call_kwargs}
+                    else:
+                        # Fallback: try to convert to dict
+                        unpacked_kwargs = {**vars(input_data), **call_kwargs}
+
+                    if asyncio.iscoroutinefunction(fn):
+                        result = await fn(**unpacked_kwargs)
+                    else:
+                        result = fn(**unpacked_kwargs)
+                else:
+                    # Standard behavior: pass input_data as first positional argument
+                    if asyncio.iscoroutinefunction(fn):
+                        result = await fn(input_data, **call_kwargs)
+                    else:
+                        result = fn(input_data, **call_kwargs)
+
+                # Log successful completion
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                node_logger.debug(
+                    "Function completed",
+                    fn_name=fn_name,
+                    duration_ms=f"{duration_ms:.2f}",
+                    output_type=type(result).__name__,
+                )
+
+                return result
+
+            except Exception as e:
+                # Log failure
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                node_logger.error(
+                    "Function failed",
+                    fn_name=fn_name,
+                    duration_ms=f"{duration_ms:.2f}",
+                    error=str(e),
+                    error_type=type(e).__name__,
+                )
+                raise
+
+        # Preserve function metadata
+        wrapped_fn.__name__ = getattr(fn, "__name__", f"wrapped_{name}")
+        wrapped_fn.__doc__ = getattr(fn, "__doc__", f"Wrapped function: {name}")
+
+        return wrapped_fn
+
+    def _resolve_function(self, fn: Callable[..., Any] | str) -> Callable[..., Any]:
+        """Resolve function from callable or module path string.
+
+        Args
+        ----
+            fn: Function (callable) or module path string (e.g., 'mymodule.myfunc')
+
+        Returns
+        -------
+        Callable[..., Any]
+            The resolved callable function
+
+        Raises
+        ------
+        TypeError
+            If fn is not a callable or string
+        ValueError
+            If string path cannot be resolved to a callable
+        """
+        if callable(fn):
+            return fn
+
+        # At this point, fn should be a string based on type hints
+        # But we validate at runtime for safety
+        if not isinstance(fn, str):  # pyright: ignore[reportUnnecessaryIsInstance]
+            raise TypeError(
+                f"Expected a callable function or string module path, got {type(fn).__name__}"
+            )
+
+        # Parse module path string
+        if "." not in fn:
+            raise ValueError(f"Function path must be in format 'module.function', got: {fn}")
+
+        # Split the module path
+        module_path, func_name = fn.rsplit(".", 1)
+
+        try:
+            module = importlib.import_module(module_path)
+            resolved_fn = getattr(module, func_name)
+
+            if not callable(resolved_fn):
+                raise ValueError(
+                    f"Resolved '{fn}' is not callable (got {type(resolved_fn).__name__})"
+                )
+
+            return resolved_fn  # type: ignore[no-any-return]
+
+        except ImportError as e:
+            raise ValueError(f"Could not import module from function path '{fn}': {e}") from e
+        except AttributeError as e:
+            raise ValueError(f"Function '{func_name}' not found in module '{module_path}'") from e
+
+    def _validate_function(self, fn: Callable[..., Any], *, unpack_input: bool = False) -> None:
+        """Validate that function can be properly wrapped.
+
+        Args
+        ----
+            fn: Function to validate
+            unpack_input: If True, function receives unpacked kwargs instead of input_data
+
+        Raises
+        ------
+        ValueError
+            If function cannot be used
+        """
+        sig = inspect.signature(fn)
+        params = list(sig.parameters.values())
+
+        if unpack_input:
+            # With unpack_input, function receives fields as **kwargs
+            # It can have any signature, including just **kwargs
+            # No validation needed - fields are passed as keyword arguments
+            return
+
+        # Standard mode: function needs at least one parameter to receive input_data
+        if not params:
+            raise ValueError("Function must have at least one parameter to receive input_data")
+
+        first_param = params[0]
+        if first_param.kind == inspect.Parameter.VAR_KEYWORD:
+            raise ValueError("First parameter cannot be **kwargs - need parameter for input_data")
+
+    def _infer_schemas_from_function(
+        self, fn: Callable[..., Any]
+    ) -> tuple[type[BaseModel] | None, type[BaseModel] | None]:
+        """Infer input and output schemas from function type annotations.
+
+        Args
+        ----
+            fn: Function to analyze
+
+        Returns
+        -------
+        tuple[type[BaseModel] | None, type[BaseModel] | None]
+            Tuple of (input_schema, output_schema) where each can be None if not inferrable
+        """
+        try:
+            type_hints = get_type_hints(fn)
+            sig = inspect.signature(fn)
+
+            # Infer input schema from first parameter (skip 'self' if present)
+            input_schema = None
+            if params := list(sig.parameters.values()):
+                # Skip 'self' parameter if present
+                first_param = (
+                    params[0]
+                    if params[0].name != "self"
+                    else (params[1] if len(params) > 1 else None)
+                )
+                if first_param and first_param.name in type_hints:
+                    param_type = type_hints[first_param.name]
+                    if is_schema_type(param_type):
+                        input_schema = param_type
+
+            # Infer output schema from return annotation
+            output_schema = None
+            if "return" in type_hints:
+                return_type = type_hints["return"]
+                if is_schema_type(return_type):
+                    output_schema = return_type
+
+            return input_schema, output_schema
+
+        except (TypeError, AttributeError, ValueError):
+            # If type hints are malformed or unavailable, skip inference
+            return None, None
+
+    @staticmethod
+    def create_passthrough_mapping(fields: list[str]) -> dict[str, str]:
+        """Create a passthrough mapping where field names are unchanged.
+
+        Args:
+        ----
+            fields: List of field names to pass through
+
+        Returns
+        -------
+        dict[str, str]
+            Mapping dict {field: field} for each field
+        """
+        return {field: field for field in fields}
+
+    @staticmethod
+    def create_rename_mapping(mapping: dict[str, str]) -> dict[str, str]:
+        """Create a simple rename mapping.
+
+        Args:
+        ----
+            mapping: Dict of {new_name: old_name}
+
+        Returns
+        -------
+        dict[str, str]
+            The mapping dict as-is (for consistency with other methods)
+        """
+        return mapping
+
+    @staticmethod
+    def create_prefixed_mapping(fields: list[str], source_node: str, prefix: str) -> dict[str, str]:
+        """Create a mapping with prefixed field names.
+
+        Args:
+        ----
+            fields: List of field names to map
+            source_node: Name of the source node
+            prefix: Prefix to add to field names
+
+        Returns
+        -------
+        dict[str, str]
+            Mapping dict {prefix_field: source_node.field}
+        """
+        return {f"{prefix}{field}": f"{source_node}.{field}" for field in fields}
+
+    def with_input_mapping(self, node: NodeSpec, input_mapping: dict[str, str]) -> NodeSpec:
+        """Enhance an existing node with input mapping.
+
+        Args:
+        ----
+            node: The node to enhance
+            input_mapping: The input mapping to apply
+
+        Returns
+        -------
+        NodeSpec
+            New NodeSpec with the input mapping applied
+        """
+        new_params = dict(node.params) if node.params else {}
+        new_params["input_mapping"] = input_mapping
+
+        return NodeSpec(
+            name=node.name,
+            fn=node.fn,
+            in_model=node.in_model,
+            out_model=node.out_model,
+            deps=node.deps,
+            params=new_params,
+            timeout=node.timeout,
+            max_retries=node.max_retries,
+            when=node.when,
+        )