quantalogic 0.55.0__py3-none-any.whl → 0.57.0__py3-none-any.whl

quantalogic/flow/flow.py

@@ -12,7 +12,8 @@
 # ///
 
 import asyncio
-import inspect  # Added for accurate parameter detection
+import inspect
+import os
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -47,18 +48,17 @@ class WorkflowEvent:
     transition_from: Optional[str] = None
     transition_to: Optional[str] = None
     sub_workflow_name: Optional[str] = None
-    usage: Optional[Dict[str, Any]] = None  # Added to store token usage and cost
+    usage: Optional[Dict[str, Any]] = None
 
 
 WorkflowObserver = Callable[[WorkflowEvent], None]
 
 
-# Define a class for sub-workflow nodes with updated inputs handling
 class SubWorkflowNode:
     def __init__(self, sub_workflow: "Workflow", inputs: Dict[str, Any], output: str):
         """Initialize a sub-workflow node with flexible inputs mapping."""
         self.sub_workflow = sub_workflow
-        self.inputs = inputs  # Maps sub_key to main_key, callable, or value
+        self.inputs = inputs
         self.output = output
 
     async def __call__(self, engine: "WorkflowEngine"):
@@ -70,7 +70,7 @@ class SubWorkflowNode:
             elif isinstance(mapping, str):
                 sub_context[sub_key] = engine.context.get(mapping)
             else:
-                sub_context[sub_key] = mapping  # Direct value
+                sub_context[sub_key] = mapping
         sub_engine = self.sub_workflow.build(parent_engine=engine)
         result = await sub_engine.run(sub_context)
         return result.get(self.output)
@@ -82,7 +82,7 @@ class WorkflowEngine:
         self.workflow = workflow
         self.context: Dict[str, Any] = {}
         self.observers: List[WorkflowObserver] = []
-        self.parent_engine = parent_engine  # Link to parent engine for sub-workflow observer propagation
+        self.parent_engine = parent_engine
 
     def add_observer(self, observer: WorkflowObserver) -> None:
         """Register an event observer callback."""
@@ -90,7 +90,7 @@ class WorkflowEngine:
             self.observers.append(observer)
             logger.debug(f"Added observer: {observer}")
         if self.parent_engine:
-            self.parent_engine.add_observer(observer)  # Propagate to parent for global visibility
+            self.parent_engine.add_observer(observer)
 
     def remove_observer(self, observer: WorkflowObserver) -> None:
         """Remove an event observer callback."""
@@ -140,18 +140,15 @@ class WorkflowEngine:
                 )
                 break
 
-            # Prepare inputs with mappings
             input_mappings = self.workflow.node_input_mappings.get(current_node, {})
             inputs = {}
-            # Add all mapped inputs
             for key, mapping in input_mappings.items():
                 if callable(mapping):
                     inputs[key] = mapping(self.context)
                 elif isinstance(mapping, str):
                     inputs[key] = self.context.get(mapping)
                 else:
-                    inputs[key] = mapping  # Direct value
-            # For parameters in node_inputs that are not mapped, get from context
+                    inputs[key] = mapping
             for param in self.workflow.node_inputs[current_node]:
                 if param not in inputs:
                     inputs[param] = self.context.get(param)
@@ -159,7 +156,6 @@ class WorkflowEngine:
             result = None
             exception = None
 
-            # Handle sub-workflow nodes
             if isinstance(node_func, SubWorkflowNode):
                 await self._notify_observers(
                     WorkflowEvent(
@@ -172,21 +168,24 @@ class WorkflowEngine:
 
             try:
                 if isinstance(node_func, SubWorkflowNode):
-                    result = await node_func(self)  # Sub-workflow handles its own inputs
-                    usage = None  # Sub-workflow usage is handled by its own nodes
+                    result = await node_func(self)
+                    usage = None
                 else:
                     result = await node_func(**inputs)
-                    usage = getattr(node_func, "usage", None)  # Extract usage if set by LLM nodes
+                    usage = getattr(node_func, "usage", None)
                 output_key = self.workflow.node_outputs[current_node]
                 if output_key:
                     self.context[output_key] = result
+                elif isinstance(result, dict):
+                    self.context.update(result)
+                    logger.debug(f"Updated context with {result} from node {current_node}")
                 await self._notify_observers(
                     WorkflowEvent(
                         event_type=WorkflowEventType.NODE_COMPLETED,
                         node_name=current_node,
                         context=self.context,
                         result=result,
-                        usage=usage,  # Include usage data in the event
+                        usage=usage,
                     )
                 )
             except Exception as e:
@@ -239,17 +238,21 @@ class WorkflowEngine:
 
 class Workflow:
     def __init__(self, start_node: str):
-        """Initialize a workflow with a starting node."""
+        """Initialize a workflow with a starting node.
+
+        Args:
+            start_node: The name of the initial node in the workflow.
+        """
         self.start_node = start_node
         self.nodes: Dict[str, Callable] = {}
         self.node_inputs: Dict[str, List[str]] = {}
         self.node_outputs: Dict[str, Optional[str]] = {}
         self.transitions: Dict[str, List[Tuple[str, Optional[Callable]]]] = {}
-        self.node_input_mappings: Dict[str, Dict[str, Any]] = {}  # Store input mappings for nodes
+        self.node_input_mappings: Dict[str, Dict[str, Any]] = {}
         self.current_node = None
         self._observers: List[WorkflowObserver] = []
-        self._register_node(start_node)  # Register the start node without setting current_node
-        self.current_node = start_node  # Set current_node explicitly after registration
+        self._register_node(start_node)
+        self.current_node = start_node
 
     def _register_node(self, name: str):
         """Register a node without modifying the current node."""
@@ -261,7 +264,15 @@ class Workflow:
         self.node_outputs[name] = output
 
     def node(self, name: str, inputs_mapping: Optional[Dict[str, Any]] = None):
-        """Add a node to the workflow chain with an optional inputs mapping."""
+        """Add a node to the workflow chain with an optional inputs mapping.
+
+        Args:
+            name: The name of the node to add.
+            inputs_mapping: Optional dictionary mapping node inputs to context keys or callables.
+
+        Returns:
+            Self for method chaining.
+        """
         self._register_node(name)
         if inputs_mapping:
             self.node_input_mappings[name] = inputs_mapping
@@ -270,7 +281,14 @@ class Workflow:
         return self
 
     def sequence(self, *nodes: str):
-        """Add a sequence of nodes to execute in order."""
+        """Add a sequence of nodes to execute in order.
+
+        Args:
+            *nodes: Variable number of node names to execute sequentially.
+
+        Returns:
+            Self for method chaining.
+        """
         if not nodes:
             return self
         for node in nodes:
@@ -286,9 +304,17 @@ class Workflow:
         return self
 
     def then(self, next_node: str, condition: Optional[Callable] = None):
-        """Add a transition to the next node with an optional condition."""
+        """Add a transition to the next node with an optional condition.
+
+        Args:
+            next_node: Name of the node to transition to.
+            condition: Optional callable taking context and returning a boolean.
+
+        Returns:
+            Self for method chaining.
+        """
         if next_node not in self.nodes:
-            self._register_node(next_node)  # Register without changing current_node
+            self._register_node(next_node)
         if self.current_node:
             self.transitions.setdefault(self.current_node, []).append((next_node, condition))
             logger.debug(f"Added transition from {self.current_node} to {next_node} with condition {condition}")
@@ -297,23 +323,49 @@ class Workflow:
         self.current_node = next_node
         return self
 
-    def branch(self, branches: List[Tuple[str, Optional[Callable]]]) -> "Workflow":
-        """Add multiple conditional branches from the current node."""
+    def branch(
+        self,
+        branches: List[Tuple[str, Optional[Callable]]],
+        default: Optional[str] = None,
+        next_node: Optional[str] = None,
+    ) -> "Workflow":
+        """Add multiple conditional branches from the current node with an optional default and next node.
+
+        Args:
+            branches: List of tuples (next_node, condition), where condition takes context and returns a boolean.
+            default: Optional node to transition to if no branch conditions are met.
+            next_node: Optional node to set as current_node after branching (e.g., for convergence).
+
+        Returns:
+            Self for method chaining.
+        """
         if not self.current_node:
             logger.warning("No current node set for branching")
             return self
-        for next_node, condition in branches:
-            if next_node not in self.nodes:
-                self._register_node(next_node)
-            self.transitions.setdefault(self.current_node, []).append((next_node, condition))
-            logger.debug(f"Added branch from {self.current_node} to {next_node} with condition {condition}")
+        for next_node_name, condition in branches:
+            if next_node_name not in self.nodes:
+                self._register_node(next_node_name)
+            self.transitions.setdefault(self.current_node, []).append((next_node_name, condition))
+            logger.debug(f"Added branch from {self.current_node} to {next_node_name} with condition {condition}")
+        if default:
+            if default not in self.nodes:
+                self._register_node(default)
+            self.transitions.setdefault(self.current_node, []).append((default, None))
+            logger.debug(f"Added default transition from {self.current_node} to {default}")
+        self.current_node = next_node  # Explicitly set next_node if provided
         return self
 
     def converge(self, convergence_node: str) -> "Workflow":
-        """Set a convergence point for all previous branches."""
+        """Set a convergence point for all previous branches.
+
+        Args:
+            convergence_node: Name of the node where branches converge.
+
+        Returns:
+            Self for method chaining.
+        """
         if convergence_node not in self.nodes:
             self._register_node(convergence_node)
-        # Find all leaf nodes (nodes with no outgoing transitions) and point them to convergence_node
         for node in self.nodes:
             if (node not in self.transitions or not self.transitions[node]) and node != convergence_node:
                 self.transitions.setdefault(node, []).append((convergence_node, None))
@@ -322,32 +374,63 @@ class Workflow:
         return self
 
     def parallel(self, *nodes: str):
-        """Add parallel nodes to execute concurrently."""
+        """Add parallel nodes to execute concurrently.
+
+        Args:
+            *nodes: Variable number of node names to execute in parallel.
+
+        Returns:
+            Self for method chaining.
+        """
         if self.current_node:
             for node in nodes:
                 self.transitions.setdefault(self.current_node, []).append((node, None))
-        self.current_node = None  # Reset after parallel to force explicit next node
+        self.current_node = None
         return self
 
     def add_observer(self, observer: WorkflowObserver) -> "Workflow":
-        """Add an event observer callback to the workflow."""
+        """Add an event observer callback to the workflow.
+
+        Args:
+            observer: Callable to handle workflow events.
+
+        Returns:
+            Self for method chaining.
+        """
         if observer not in self._observers:
             self._observers.append(observer)
             logger.debug(f"Added observer to workflow: {observer}")
-        return self  # Support chaining
+        return self
 
     def add_sub_workflow(self, name: str, sub_workflow: "Workflow", inputs: Dict[str, Any], output: str):
-        """Add a sub-workflow as a node with flexible inputs mapping."""
+        """Add a sub-workflow as a node with flexible inputs mapping.
+
+        Args:
+            name: Name of the sub-workflow node.
+            sub_workflow: The Workflow instance to embed.
+            inputs: Dictionary mapping sub-workflow inputs to context keys or callables.
+            output: Context key for the sub-workflow's result.
+
+        Returns:
+            Self for method chaining.
+        """
         sub_node = SubWorkflowNode(sub_workflow, inputs, output)
         self.nodes[name] = sub_node
-        self.node_inputs[name] = []  # Inputs handled internally by SubWorkflowNode
+        self.node_inputs[name] = []
         self.node_outputs[name] = output
         self.current_node = name
         logger.debug(f"Added sub-workflow {name} with inputs {inputs} and output {output}")
         return self
 
     def build(self, parent_engine: Optional["WorkflowEngine"] = None) -> WorkflowEngine:
-        """Build and return a WorkflowEngine instance with registered observers."""
+        """Build and return a WorkflowEngine instance with registered observers.
+
+        Args:
+            parent_engine: Optional parent WorkflowEngine for sub-workflows.
+
+        Returns:
+            Configured WorkflowEngine instance.
+        """
         engine = WorkflowEngine(self, parent_engine=parent_engine)
         for observer in self._observers:
             engine.add_observer(observer)
@@ -355,11 +438,18 @@ class Workflow:
 
 
 class Nodes:
-    NODE_REGISTRY: Dict[str, Tuple[Callable, List[str], Optional[str]]] = {}  # Registry to hold node functions and metadata
+    NODE_REGISTRY: Dict[str, Tuple[Callable, List[str], Optional[str]]] = {}
 
     @classmethod
     def define(cls, output: Optional[str] = None):
-        """Decorator for defining simple workflow nodes."""
+        """Decorator for defining simple workflow nodes.
+
+        Args:
+            output: Optional context key for the node's result.
+
+        Returns:
+            Decorator function wrapping the node logic.
+        """
         def decorator(func: Callable) -> Callable:
             async def wrapped_func(**kwargs):
                 try:
@@ -372,8 +462,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in node {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature
             sig = inspect.signature(func)
             inputs = [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
@@ -383,7 +471,14 @@ class Nodes:
 
     @classmethod
     def validate_node(cls, output: str):
-        """Decorator for nodes that validate inputs."""
+        """Decorator for nodes that validate inputs and return a string.
+
+        Args:
+            output: Context key for the validation result.
+
+        Returns:
+            Decorator function wrapping the validation logic.
+        """
         def decorator(func: Callable) -> Callable:
             async def wrapped_func(**kwargs):
                 try:
@@ -398,8 +493,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Validation error in {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature
             sig = inspect.signature(func)
             inputs = [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
@@ -409,11 +502,18 @@ class Nodes:
 
     @classmethod
     def transform_node(cls, output: str, transformer: Callable[[Any], Any]):
-        """Decorator for nodes that transform their inputs."""
+        """Decorator for nodes that transform their inputs.
+
+        Args:
+            output: Context key for the transformed result.
+            transformer: Callable to transform the input.
+
+        Returns:
+            Decorator function wrapping the transformation logic.
+        """
         def decorator(func: Callable) -> Callable:
             async def wrapped_func(**kwargs):
                 try:
-                    # Apply transformer to the first input value
                     input_key = list(kwargs.keys())[0] if kwargs else None
                     if input_key:
                         transformed_input = transformer(kwargs[input_key])
@@ -427,8 +527,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in transform node {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature
             sig = inspect.signature(func)
             inputs = [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
@@ -467,8 +565,9 @@ class Nodes:
     @classmethod
     def llm_node(
         cls,
-        system_prompt: str,
-        output: str,
+        system_prompt: str = "",
+        system_prompt_file: Optional[str] = None,
+        output: str = "",
         prompt_template: str = "",
         prompt_file: Optional[str] = None,
         temperature: float = 0.7,
@@ -476,13 +575,38 @@ class Nodes:
         top_p: float = 1.0,
         presence_penalty: float = 0.0,
         frequency_penalty: float = 0.0,
+        model: Callable[[Dict[str, Any]], str] = lambda ctx: "gpt-3.5-turbo",
         **kwargs,
     ):
-        """Decorator for creating LLM nodes with plain text output, supporting dynamic parameters via input mappings."""
+        """Decorator for creating LLM nodes with plain text output, supporting dynamic parameters.
+
+        Args:
+            system_prompt: Inline system prompt defining LLM behavior.
+            system_prompt_file: Path to a system prompt template file (overrides system_prompt).
+            output: Context key for the LLM's result.
+            prompt_template: Inline Jinja2 template for the user prompt.
+            prompt_file: Path to a user prompt template file (overrides prompt_template).
+            temperature: Randomness control (0.0 to 1.0).
+            max_tokens: Maximum response length.
+            top_p: Nucleus sampling parameter (0.0 to 1.0).
+            presence_penalty: Penalty for repetition (-2.0 to 2.0).
+            frequency_penalty: Penalty for frequent words (-2.0 to 2.0).
+            model: Callable or string to determine the LLM model dynamically from context.
+            **kwargs: Additional parameters for the LLM call.
+
+        Returns:
+            Decorator function wrapping the LLM logic.
+        """
         def decorator(func: Callable) -> Callable:
-            async def wrapped_func(model: str, **func_kwargs):
-                # Extract parameters from func_kwargs if provided, else use defaults
+            async def wrapped_func(model_param: str = None, **func_kwargs):
                 system_prompt_to_use = func_kwargs.pop("system_prompt", system_prompt)
+                system_prompt_file_to_use = func_kwargs.pop("system_prompt_file", system_prompt_file)
+                
+                if system_prompt_file_to_use:
+                    system_content = cls._load_prompt_from_file(system_prompt_file_to_use, func_kwargs)
+                else:
+                    system_content = system_prompt_to_use
+                
                 prompt_template_to_use = func_kwargs.pop("prompt_template", prompt_template)
                 prompt_file_to_use = func_kwargs.pop("prompt_file", prompt_file)
                 temperature_to_use = func_kwargs.pop("temperature", temperature)
@@ -490,25 +614,27 @@ class Nodes:
                 top_p_to_use = func_kwargs.pop("top_p", top_p)
                 presence_penalty_to_use = func_kwargs.pop("presence_penalty", presence_penalty)
                 frequency_penalty_to_use = func_kwargs.pop("frequency_penalty", frequency_penalty)
+                
+                # Prioritize model from func_kwargs (workflow mapping), then model_param, then default
+                model_to_use = func_kwargs.get("model", model_param if model_param is not None else model(func_kwargs))
+                logger.debug(f"Selected model for {func.__name__}: {model_to_use}")
 
-                # Use only signature parameters for template rendering
                 sig = inspect.signature(func)
                 template_vars = {k: v for k, v in func_kwargs.items() if k in sig.parameters}
                 prompt = cls._render_template(prompt_template_to_use, prompt_file_to_use, template_vars)
                 messages = [
-                    {"role": "system", "content": system_prompt_to_use},
+                    {"role": "system", "content": system_content},
                     {"role": "user", "content": prompt},
                 ]
                 
-                # Log the model and a preview of the prompt
                 truncated_prompt = prompt[:200] + "..." if len(prompt) > 200 else prompt
-                logger.info(f"LLM node {func.__name__} using model: {model}")
-                logger.debug(f"System prompt: {system_prompt_to_use[:100]}...")
+                logger.info(f"LLM node {func.__name__} using model: {model_to_use}")
+                logger.debug(f"System prompt: {system_content[:100]}...")
                 logger.debug(f"User prompt preview: {truncated_prompt}")
                 
                 try:
                     response = await acompletion(
-                        model=model,
+                        model=model_to_use,
                         messages=messages,
                         temperature=temperature_to_use,
                         max_tokens=max_tokens_to_use,
@@ -530,8 +656,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in LLM node {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature and add 'model'
             sig = inspect.signature(func)
             inputs = ['model'] + [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
@@ -542,9 +666,10 @@ class Nodes:
     @classmethod
     def structured_llm_node(
         cls,
-        system_prompt: str,
-        output: str,
-        response_model: Type[BaseModel],
+        system_prompt: str = "",
+        system_prompt_file: Optional[str] = None,
+        output: str = "",
+        response_model: Type[BaseModel] = None,
         prompt_template: str = "",
         prompt_file: Optional[str] = None,
         temperature: float = 0.7,
@@ -552,9 +677,29 @@ class Nodes:
         top_p: float = 1.0,
         presence_penalty: float = 0.0,
         frequency_penalty: float = 0.0,
+        model: Callable[[Dict[str, Any]], str] = lambda ctx: "gpt-3.5-turbo",
         **kwargs,
     ):
-        """Decorator for creating LLM nodes with structured output, supporting dynamic parameters via input mappings."""
+        """Decorator for creating LLM nodes with structured output, supporting dynamic parameters.
+
+        Args:
+            system_prompt: Inline system prompt defining LLM behavior.
+            system_prompt_file: Path to a system prompt template file (overrides system_prompt).
+            output: Context key for the LLM's structured result.
+            response_model: Pydantic model class for structured output.
+            prompt_template: Inline Jinja2 template for the user prompt.
+            prompt_file: Path to a user prompt template file (overrides prompt_template).
+            temperature: Randomness control (0.0 to 1.0).
+            max_tokens: Maximum response length.
+            top_p: Nucleus sampling parameter (0.0 to 1.0).
+            presence_penalty: Penalty for repetition (-2.0 to 2.0).
+            frequency_penalty: Penalty for frequent words (-2.0 to 2.0).
+            model: Callable or string to determine the LLM model dynamically from context.
+            **kwargs: Additional parameters for the LLM call.
+
+        Returns:
+            Decorator function wrapping the structured LLM logic.
+        """
         try:
             client = instructor.from_litellm(acompletion)
         except ImportError:
@@ -562,9 +707,15 @@ class Nodes:
             raise ImportError("Instructor is required for structured_llm_node")
 
         def decorator(func: Callable) -> Callable:
-            async def wrapped_func(model: str, **func_kwargs):
-                # Extract parameters from func_kwargs if provided, else use defaults
+            async def wrapped_func(model_param: str = None, **func_kwargs):
                 system_prompt_to_use = func_kwargs.pop("system_prompt", system_prompt)
+                system_prompt_file_to_use = func_kwargs.pop("system_prompt_file", system_prompt_file)
+                
+                if system_prompt_file_to_use:
+                    system_content = cls._load_prompt_from_file(system_prompt_file_to_use, func_kwargs)
+                else:
+                    system_content = system_prompt_to_use
+                
                 prompt_template_to_use = func_kwargs.pop("prompt_template", prompt_template)
                 prompt_file_to_use = func_kwargs.pop("prompt_file", prompt_file)
                 temperature_to_use = func_kwargs.pop("temperature", temperature)
@@ -572,26 +723,28 @@ class Nodes:
                 top_p_to_use = func_kwargs.pop("top_p", top_p)
                 presence_penalty_to_use = func_kwargs.pop("presence_penalty", presence_penalty)
                 frequency_penalty_to_use = func_kwargs.pop("frequency_penalty", frequency_penalty)
+                
+                # Prioritize model from func_kwargs (workflow mapping), then model_param, then default
+                model_to_use = func_kwargs.get("model", model_param if model_param is not None else model(func_kwargs))
+                logger.debug(f"Selected model for {func.__name__}: {model_to_use}")
 
-                # Use only signature parameters for template rendering
                 sig = inspect.signature(func)
                 template_vars = {k: v for k, v in func_kwargs.items() if k in sig.parameters}
                 prompt = cls._render_template(prompt_template_to_use, prompt_file_to_use, template_vars)
                 messages = [
-                    {"role": "system", "content": system_prompt_to_use},
+                    {"role": "system", "content": system_content},
                     {"role": "user", "content": prompt},
                 ]
                 
-                # Log the model and a preview of the prompt
                 truncated_prompt = prompt[:200] + "..." if len(prompt) > 200 else prompt
-                logger.info(f"Structured LLM node {func.__name__} using model: {model}")
-                logger.debug(f"System prompt: {system_prompt_to_use[:100]}...")
+                logger.info(f"Structured LLM node {func.__name__} using model: {model_to_use}")
+                logger.debug(f"System prompt: {system_content[:100]}...")
                 logger.debug(f"User prompt preview: {truncated_prompt}")
                 logger.debug(f"Expected response model: {response_model.__name__}")
                 
                 try:
                     structured_response, raw_response = await client.chat.completions.create_with_completion(
-                        model=model,
+                        model=model_to_use,
                         messages=messages,
                         response_model=response_model,
                         temperature=temperature_to_use,
@@ -616,8 +769,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in structured LLM node {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature and add 'model'
             sig = inspect.signature(func)
             inputs = ['model'] + [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
@@ -632,20 +783,26 @@ class Nodes:
         template: str = "",
         template_file: Optional[str] = None,
     ):
-        """Decorator for creating nodes that apply a Jinja2 template to inputs, supporting dynamic parameters."""
+        """Decorator for creating nodes that apply a Jinja2 template to inputs.
+
+        Args:
+            output: Context key for the rendered result.
+            template: Inline Jinja2 template string.
+            template_file: Path to a template file (overrides template).
+
+        Returns:
+            Decorator function wrapping the template logic.
+        """
         def decorator(func: Callable) -> Callable:
             async def wrapped_func(**func_kwargs):
-                # Extract template parameters from func_kwargs if provided, else use defaults
                 template_to_use = func_kwargs.pop("template", template)
                 template_file_to_use = func_kwargs.pop("template_file", template_file)
 
-                # Use only signature parameters (excluding rendered_content) for template rendering
                 sig = inspect.signature(func)
                 expected_params = [p.name for p in sig.parameters.values() if p.name != 'rendered_content']
                 template_vars = {k: v for k, v in func_kwargs.items() if k in expected_params}
                 rendered_content = cls._render_template(template_to_use, template_file_to_use, template_vars)
 
-                # Filter func_kwargs for the function call
                 filtered_kwargs = {k: v for k, v in func_kwargs.items() if k in expected_params}
 
                 try:
@@ -658,8 +815,6 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in template node {func.__name__}: {e}")
                     raise
-
-            # Get parameter names from function signature and add 'rendered_content' if not present
             sig = inspect.signature(func)
             inputs = [param.name for param in sig.parameters.values()]
             if 'rendered_content' not in inputs:
@@ -670,15 +825,20 @@ class Nodes:
         return decorator
 
 
-# Example workflow with observer integration, updated nodes, input mappings, and dynamic parameters
+# Add a templates directory path at the module level
+TEMPLATES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates")
+
+# Helper function to get template paths
+def get_template_path(template_name):
+    return os.path.join(TEMPLATES_DIR, template_name)
+
+
 async def example_workflow():
-    # Define Pydantic model for structured output
     class OrderDetails(BaseModel):
         order_id: str
         items_in_stock: List[str]
         items_out_of_stock: List[str]
 
-    # Define an example observer for progress
     async def progress_monitor(event: WorkflowEvent):
         print(f"[{event.event_type.value}] {event.node_name or 'Workflow'}")
         if event.result is not None:
@@ -686,7 +846,6 @@ async def example_workflow():
         if event.exception is not None:
             print(f"Exception: {event.exception}")
 
-    # Define an observer for token usage
     class TokenUsageObserver:
         def __init__(self):
             self.total_prompt_tokens = 0
@@ -709,16 +868,15 @@ async def example_workflow():
                 for node, usage in self.node_usages.items():
                     print(f"Node {node}: {usage}")
 
-    # Define nodes
     @Nodes.validate_node(output="validation_result")
     async def validate_order(order: Dict[str, Any]) -> str:
         return "Order validated" if order.get("items") else "Invalid order"
 
     @Nodes.structured_llm_node(
-        system_prompt="You are an inventory checker. Respond with a JSON object containing 'order_id', 'items_in_stock', and 'items_out_of_stock'.",
+        system_prompt_file=get_template_path("system_check_inventory.j2"),
         output="inventory_status",
         response_model=OrderDetails,
-        prompt_template="Check if the following items are in stock: {{ items }}. Return the result in JSON format with 'order_id' set to '123'.",
+        prompt_file=get_template_path("prompt_check_inventory.j2"),
     )
     async def check_inventory(items: List[str]) -> OrderDetails:
         return OrderDetails(order_id="123", items_in_stock=["item1"], items_out_of_stock=[])
@@ -754,13 +912,10 @@ async def example_workflow():
     async def format_order_message(rendered_content: str, items: List[str]) -> str:
         return rendered_content
 
-    # Sub-workflow for payment and shipping
     payment_shipping_sub_wf = Workflow("process_payment").sequence("process_payment", "arrange_shipping")
 
-    # Instantiate token usage observer
     token_observer = TokenUsageObserver()
 
-    # Main workflow with dynamic parameter overrides
     workflow = (
         Workflow("validate_order")
         .add_observer(progress_monitor)
@@ -769,13 +924,13 @@ async def example_workflow():
         .node("transform_items")
         .node("format_order_message", inputs_mapping={
             "items": "items",
-            "template": "Custom order: {{ items | join(', ') }}"  # Dynamic override
+            "template": "Custom order: {{ items | join(', ') }}"
         })
         .node("check_inventory", inputs_mapping={
             "model": lambda ctx: "gemini/gemini-2.0-flash",
             "items": "transformed_items",
-            "temperature": 0.5,  # Dynamic override
-            "max_tokens": 1000   # Dynamic override
+            "temperature": 0.5,
+            "max_tokens": 1000
         })
         .add_sub_workflow(
             "payment_shipping",
@@ -783,15 +938,17 @@ async def example_workflow():
             inputs={"order": lambda ctx: {"items": ctx["items"]}},
             output="shipping_confirmation"
         )
-        .branch([
-            ("payment_shipping", lambda ctx: len(ctx.get("inventory_status").items_out_of_stock) == 0 if ctx.get("inventory_status") else False),
-            ("notify_customer_out_of_stock", lambda ctx: len(ctx.get("inventory_status").items_out_of_stock) > 0 if ctx.get("inventory_status") else True)
-        ])
+        .branch(
+            [
+                ("payment_shipping", lambda ctx: len(ctx.get("inventory_status").items_out_of_stock) == 0 if ctx.get("inventory_status") else False),
+                ("notify_customer_out_of_stock", lambda ctx: len(ctx.get("inventory_status").items_out_of_stock) > 0 if ctx.get("inventory_status") else True)
+            ],
+            next_node="update_order_status"
+        )
         .converge("update_order_status")
         .sequence("update_order_status", "send_confirmation_email")
     )
 
-    # Execute workflow
     initial_context = {"customer_order": {"items": ["item1", "item2"]}, "items": ["item1", "item2"]}
     engine = workflow.build()
     result = await engine.run(initial_context)