quantalogic 0.61.3__py3-none-any.whl → 0.80__py3-none-any.whl

quantalogic/flow/flow.py

@@ -17,7 +17,7 @@ import os
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import instructor
 from jinja2 import Environment, FileSystemLoader, Template, TemplateNotFound
@@ -253,6 +253,10 @@ class Workflow:
         self._observers: List[WorkflowObserver] = []
         self._register_node(start_node)
         self.current_node = start_node
+        # Loop-specific attributes
+        self.in_loop = False
+        self.loop_nodes = []
+        self.loop_entry_node = None
 
     def _register_node(self, name: str):
         """Register a node without modifying the current node."""
@@ -274,6 +278,8 @@ class Workflow:
             Self for method chaining.
         """
         self._register_node(name)
+        if self.in_loop:
+            self.loop_nodes.append(name)
         if inputs_mapping:
             self.node_input_mappings[name] = inputs_mapping
             logger.debug(f"Added inputs mapping for node {name}: {inputs_mapping}")
@@ -422,6 +428,66 @@ class Workflow:
         logger.debug(f"Added sub-workflow {name} with inputs {inputs} and output {output}")
         return self
 
+    def start_loop(self):
+        """Begin defining a loop in the workflow.
+
+        Raises:
+            ValueError: If called without a current node.
+
+        Returns:
+            Self for method chaining.
+        """
+        if self.current_node is None:
+            raise ValueError("Cannot start loop without a current node")
+        self.loop_entry_node = self.current_node
+        self.in_loop = True
+        self.loop_nodes = []
+        return self
+
+    def end_loop(self, condition: Callable[[Dict[str, Any]], bool], next_node: str):
+        """End the loop, setting up transitions based on the condition.
+
+        Args:
+            condition: Callable taking context and returning True when the loop should exit.
+            next_node: Name of the node to transition to after the loop exits.
+
+        Raises:
+            ValueError: If no loop nodes are defined.
+
+        Returns:
+            Self for method chaining.
+        """
+        if not self.in_loop or not self.loop_nodes:
+            raise ValueError("No loop nodes defined")
+        
+        first_node = self.loop_nodes[0]
+        last_node = self.loop_nodes[-1]
+        
+        # Transition from the node before the loop to the first loop node
+        self.transitions.setdefault(self.loop_entry_node, []).append((first_node, None))
+        
+        # Transitions within the loop
+        for i in range(len(self.loop_nodes) - 1):
+            self.transitions.setdefault(self.loop_nodes[i], []).append((self.loop_nodes[i + 1], None))
+        
+        # Conditional transitions from the last loop node
+        # If condition is False, loop back to the first node
+        self.transitions.setdefault(last_node, []).append((first_node, lambda ctx: not condition(ctx)))
+        # If condition is True, exit to the next node
+        self.transitions.setdefault(last_node, []).append((next_node, condition))
+        
+        # Register the next_node if not already present
+        if next_node not in self.nodes:
+            self._register_node(next_node)
+        
+        # Update state
+        self.current_node = next_node
+        self.in_loop = False
+        self.loop_nodes = []
+        self.loop_entry_node = None
+        
+        return self
+
     def build(self, parent_engine: Optional["WorkflowEngine"] = None) -> WorkflowEngine:
         """Build and return a WorkflowEngine instance with registered observers.
 
@@ -575,7 +641,7 @@ 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",
+        model: Union[Callable[[Dict[str, Any]], str], str] = lambda ctx: "gpt-3.5-turbo",
         **kwargs,
     ):
         """Decorator for creating LLM nodes with plain text output, supporting dynamic parameters.
@@ -598,27 +664,45 @@ class Nodes:
             Decorator function wrapping the LLM logic.
         """
         def decorator(func: Callable) -> Callable:
-            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)
-                
+            # Store all decorator parameters in a config dictionary
+            config = {
+                "system_prompt": system_prompt,
+                "system_prompt_file": system_prompt_file,
+                "prompt_template": prompt_template,
+                "prompt_file": prompt_file,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "top_p": top_p,
+                "presence_penalty": presence_penalty,
+                "frequency_penalty": frequency_penalty,
+                "model": model,
+                **kwargs,
+            }
+
+            async def wrapped_func(**func_kwargs):
+                # Use func_kwargs to override config values if provided, otherwise use config defaults
+                system_prompt_to_use = func_kwargs.pop("system_prompt", config["system_prompt"])
+                system_prompt_file_to_use = func_kwargs.pop("system_prompt_file", config["system_prompt_file"])
+                prompt_template_to_use = func_kwargs.pop("prompt_template", config["prompt_template"])
+                prompt_file_to_use = func_kwargs.pop("prompt_file", config["prompt_file"])
+                temperature_to_use = func_kwargs.pop("temperature", config["temperature"])
+                max_tokens_to_use = func_kwargs.pop("max_tokens", config["max_tokens"])
+                top_p_to_use = func_kwargs.pop("top_p", config["top_p"])
+                presence_penalty_to_use = func_kwargs.pop("presence_penalty", config["presence_penalty"])
+                frequency_penalty_to_use = func_kwargs.pop("frequency_penalty", config["frequency_penalty"])
+                model_to_use = func_kwargs.pop("model", config["model"])
+
+                # Handle callable model parameter
+                if callable(model_to_use):
+                    model_to_use = model_to_use(func_kwargs)
+
+                # Load system prompt from file if specified
                 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)
-                max_tokens_to_use = func_kwargs.pop("max_tokens", max_tokens)
-                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}")
 
+                # Prepare template variables and render prompt
                 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)
@@ -626,12 +710,14 @@ class Nodes:
                     {"role": "system", "content": system_content},
                     {"role": "user", "content": prompt},
                 ]
-                
+
+                # Logging for debugging
                 truncated_prompt = prompt[:200] + "..." if len(prompt) > 200 else prompt
                 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}")
-                
+
+                # Call the acompletion function with the resolved model
                 try:
                     response = await acompletion(
                         model=model_to_use,
@@ -656,8 +742,10 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in LLM node {func.__name__}: {e}")
                     raise
+
+            # Register the node with its inputs and output
             sig = inspect.signature(func)
-            inputs = ['model'] + [param.name for param in sig.parameters.values()]
+            inputs = [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
             cls.NODE_REGISTRY[func.__name__] = (wrapped_func, inputs, output)
             return wrapped_func
@@ -677,7 +765,7 @@ 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",
+        model: Union[Callable[[Dict[str, Any]], str], str] = lambda ctx: "gpt-3.5-turbo",
         **kwargs,
     ):
         """Decorator for creating LLM nodes with structured output, supporting dynamic parameters.
@@ -707,27 +795,45 @@ class Nodes:
             raise ImportError("Instructor is required for structured_llm_node")
 
         def decorator(func: Callable) -> Callable:
-            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)
-                
+            # Store all decorator parameters in a config dictionary
+            config = {
+                "system_prompt": system_prompt,
+                "system_prompt_file": system_prompt_file,
+                "prompt_template": prompt_template,
+                "prompt_file": prompt_file,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "top_p": top_p,
+                "presence_penalty": presence_penalty,
+                "frequency_penalty": frequency_penalty,
+                "model": model,
+                **kwargs,
+            }
+
+            async def wrapped_func(**func_kwargs):
+                # Resolve parameters, prioritizing func_kwargs over config defaults
+                system_prompt_to_use = func_kwargs.pop("system_prompt", config["system_prompt"])
+                system_prompt_file_to_use = func_kwargs.pop("system_prompt_file", config["system_prompt_file"])
+                prompt_template_to_use = func_kwargs.pop("prompt_template", config["prompt_template"])
+                prompt_file_to_use = func_kwargs.pop("prompt_file", config["prompt_file"])
+                temperature_to_use = func_kwargs.pop("temperature", config["temperature"])
+                max_tokens_to_use = func_kwargs.pop("max_tokens", config["max_tokens"])
+                top_p_to_use = func_kwargs.pop("top_p", config["top_p"])
+                presence_penalty_to_use = func_kwargs.pop("presence_penalty", config["presence_penalty"])
+                frequency_penalty_to_use = func_kwargs.pop("frequency_penalty", config["frequency_penalty"])
+                model_to_use = func_kwargs.pop("model", config["model"])
+
+                # Handle callable model parameter
+                if callable(model_to_use):
+                    model_to_use = model_to_use(func_kwargs)
+
+                # Load system prompt from file if specified
                 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)
-                max_tokens_to_use = func_kwargs.pop("max_tokens", max_tokens)
-                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}")
 
+                # Render prompt using template variables
                 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)
@@ -735,13 +841,15 @@ class Nodes:
                     {"role": "system", "content": system_content},
                     {"role": "user", "content": prompt},
                 ]
-                
+
+                # Logging for debugging
                 truncated_prompt = prompt[:200] + "..." if len(prompt) > 200 else prompt
                 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__}")
-                
+
+                # Generate structured response
                 try:
                     structured_response, raw_response = await client.chat.completions.create_with_completion(
                         model=model_to_use,
@@ -769,8 +877,10 @@ class Nodes:
                 except Exception as e:
                     logger.error(f"Error in structured LLM node {func.__name__}: {e}")
                     raise
+
+            # Register the node
             sig = inspect.signature(func)
-            inputs = ['model'] + [param.name for param in sig.parameters.values()]
+            inputs = [param.name for param in sig.parameters.values()]
             logger.debug(f"Registering node {func.__name__} with inputs {inputs} and output {output}")
             cls.NODE_REGISTRY[func.__name__] = (wrapped_func, inputs, output)
             return wrapped_func

quantalogic/flow/flow_extractor.py

@@ -21,7 +21,7 @@ class WorkflowExtractor(ast.NodeVisitor):
     AST visitor to extract workflow nodes and structure from a Python file.
 
     This class parses Python source code to identify workflow components defined with Nodes decorators
-    and Workflow construction, including branch and converge patterns, building a WorkflowDefinition
+    and Workflow construction, including branch, converge, and loop patterns, building a WorkflowDefinition
     compatible with WorkflowManager. Fully supports input mappings and template nodes.
     """
 
@@ -34,6 +34,10 @@ class WorkflowExtractor(ast.NodeVisitor):
         self.global_vars = {}  # Tracks global variable assignments (e.g., DEFAULT_LLM_PARAMS)
         self.observers = []  # List of observer function names
         self.convergence_nodes = []  # List of convergence nodes
+        # Added for loop support
+        self.in_loop = False           # Flag indicating if we're inside a loop
+        self.loop_nodes = []           # List of nodes within the current loop
+        self.loop_entry_node = None    # Node before the loop starts
 
     def visit_Module(self, node):
         """Log and explicitly process top-level statements in the module."""
@@ -468,6 +472,10 @@ class WorkflowExtractor(ast.NodeVisitor):
                     if previous_node:
                         self.transitions.append(TransitionDefinition(from_node=previous_node, to_node=node_name))
                         logger.debug(f"Added node transition: {previous_node} -> {node_name}")
+                    # Add node to loop_nodes if inside a loop
+                    if self.in_loop:
+                        self.loop_nodes.append(node_name)
+                        logger.debug(f"Added '{node_name}' to loop_nodes in '{var_name}'")
                 return node_name
 
             elif method_name == "add_sub_workflow":
@@ -519,6 +527,58 @@ class WorkflowExtractor(ast.NodeVisitor):
                     logger.warning(f"Unsupported observer argument in 'add_observer' for '{var_name}'")
                 return previous_node
 
+            elif method_name == "start_loop":
+                if previous_node is None:
+                    logger.warning(f"start_loop called without a previous node in '{var_name}'")
+                    return None
+                self.in_loop = True
+                self.loop_entry_node = previous_node
+                self.loop_nodes = []
+                logger.debug(f"Started loop after node '{previous_node}' in '{var_name}'")
+                return previous_node
+
+            elif method_name == "end_loop":
+                cond = None
+                next_node = None
+                for keyword in expr.keywords:
+                    if keyword.arg == "condition":
+                        cond = ast.unparse(keyword.value)
+                    elif keyword.arg == "next_node":
+                        next_node = (keyword.value.value 
+                                   if isinstance(keyword.value, ast.Constant) 
+                                   else ast.unparse(keyword.value))
+                if not cond or not next_node:
+                    logger.warning(f"end_loop in '{var_name}' missing condition or next_node")
+                    return None
+                if not self.loop_nodes:
+                    logger.warning(f"end_loop called without loop nodes in '{var_name}'")
+                    return None
+                first_loop_node = self.loop_nodes[0]
+                last_loop_node = self.loop_nodes[-1]
+                # Loop-back transition: last node to first node when condition is false
+                negated_cond = f"not ({cond})"
+                self.transitions.append(
+                    TransitionDefinition(
+                        from_node=last_loop_node,
+                        to_node=first_loop_node,
+                        condition=negated_cond
+                    )
+                )
+                # Exit transition: last node to next_node when condition is true
+                self.transitions.append(
+                    TransitionDefinition(
+                        from_node=last_loop_node,
+                        to_node=next_node,
+                        condition=cond
+                    )
+                )
+                logger.debug(f"Added loop transitions: '{last_loop_node}' -> '{first_loop_node}' "
+                            f"(not {cond}), '{last_loop_node}' -> '{next_node}' ({cond})")
+                self.in_loop = False
+                self.loop_nodes = []
+                self.loop_entry_node = None
+                return next_node
+
             else:
                 logger.warning(f"Unsupported Workflow method '{method_name}' in variable '{var_name}'")
         return None

quantalogic/flow/flow_generator.py

@@ -3,7 +3,6 @@ import os
 import re
 from typing import Dict, Optional
 
-from quantalogic.flow.flow import Nodes
 from quantalogic.flow.flow_manager_schema import BranchCondition, WorkflowDefinition
 
 
@@ -27,7 +26,7 @@ def generate_executable_script(
     - Metadata specifying the required Python version and dependencies.
     - Global variables from the original script.
     - Functions defined with appropriate Nodes decorators (e.g., @Nodes.define, @Nodes.llm_node).
-    - Workflow instantiation using direct chaining syntax with function names, including branch and converge.
+    - Workflow instantiation using direct chaining syntax with function names, including branch, converge, and loop support.
     - Support for input mappings and template nodes via workflow configuration and decorators.
     - A default initial_context inferred from the workflow with customization guidance.
     """
@@ -91,6 +90,22 @@ def generate_executable_script(
                     if not value.startswith("lambda ctx:"):  # Static mappings only
                         initial_context[value] = ""
 
+    # Detect loops
+    loop_nodes = []
+    loop_condition = None
+    loop_exit_node = None
+    for trans in workflow_def.workflow.transitions:
+        if isinstance(trans.to_node, str) and trans.condition:
+            # Check for loop-back transition
+            if any(t.from_node == trans.to_node and t.to_node == trans.from_node for t in workflow_def.workflow.transitions):
+                loop_nodes.append(trans.from_node)
+                loop_nodes.append(trans.to_node)
+                loop_condition = trans.condition
+            # Check for exit transition
+            elif loop_nodes and trans.from_node == loop_nodes[-1] and f"not ({loop_condition})" in trans.condition:
+                loop_exit_node = trans.to_node
+    loop_nodes = list(dict.fromkeys(loop_nodes))  # Remove duplicates, preserve order
+
     with open(output_file, "w") as f:
         # Shebang and metadata
         f.write("#!/usr/bin/env -S uv run\n")
@@ -159,15 +174,20 @@ def generate_executable_script(
                         decorator = f"@Nodes.define(output={repr(node_def.output or f'{node_name}_result')})\n"
                     f.write(f"{decorator}{func_body}\n\n")
 
-        # Define workflow using chaining syntax
-        f.write("# Define the workflow with branch and converge support\n")
+        # Define workflow using chaining syntax with loop support
+        f.write("# Define the workflow with branch, converge, and loop support\n")
         f.write("workflow = (\n")
         start_node = workflow_def.workflow.start
         start_func = workflow_def.nodes[start_node].function if start_node in workflow_def.nodes and workflow_def.nodes[start_node].function else start_node
         f.write(f'    Workflow("{start_func}")\n')
 
+        added_nodes = set()
         for node_name, node_def in workflow_def.nodes.items():
+            if node_name in added_nodes:
+                continue
             func_name = node_def.function if node_def.function else node_name
+            if loop_nodes and node_name == loop_nodes[0]:
+                f.write("    .start_loop()\n")
             if node_def.sub_workflow:
                 sub_start = node_def.sub_workflow.start or f"{node_name}_start"
                 sub_start_func = workflow_def.nodes[sub_start].function if sub_start in workflow_def.nodes and workflow_def.nodes[sub_start].function else sub_start
@@ -208,15 +228,23 @@ def generate_executable_script(
                     f.write(f'    .node("{func_name}", inputs_mapping={inputs_mapping_str})\n')
                 else:
                     f.write(f'    .node("{func_name}")\n')
+            added_nodes.add(node_name)
 
         for trans in workflow_def.workflow.transitions:
             from_node = trans.from_node
-            from_func = workflow_def.nodes[from_node].function if from_node in workflow_def.nodes and workflow_def.nodes[from_node].function else from_node
+            _from_func = workflow_def.nodes[from_node].function if from_node in workflow_def.nodes and workflow_def.nodes[from_node].function else from_node
+            if from_node not in added_nodes:
+                continue  # Skip if already added via .node()
             to_node = trans.to_node
             if isinstance(to_node, str):
+                if loop_nodes and from_node in loop_nodes and to_node in loop_nodes:
+                    continue  # Skip loop-back transition, handled by end_loop
                 to_func = workflow_def.nodes[to_node].function if to_node in workflow_def.nodes and workflow_def.nodes[to_node].function else to_node
                 condition = f"lambda ctx: {trans.condition}" if trans.condition else "None"
-                f.write(f'    .then("{to_func}", condition={condition})\n')
+                if loop_nodes and from_node == loop_nodes[-1] and to_node == loop_exit_node:
+                    f.write(f'    .end_loop(condition=lambda ctx: {loop_condition}, next_node="{to_func}")\n')
+                else:
+                    f.write(f'    .then("{to_func}", condition={condition})\n')
             elif all(isinstance(tn, str) for tn in to_node):
                 to_funcs = [workflow_def.nodes[tn].function if tn in workflow_def.nodes and workflow_def.nodes[tn].function else tn for tn in to_node]
                 f.write(f'    .parallel({", ".join(f"{n!r}" for n in to_funcs)})\n')

quantalogic/flow/flow_manager.py

@@ -195,6 +195,27 @@ class WorkflowManager:
         )
         self.workflow.workflow.transitions.append(transition)
 
+    def add_loop(self, loop_nodes: List[str], condition: str, exit_node: str) -> None:
+        """Add a loop construct to the workflow.
+
+        Args:
+            loop_nodes: List of node names to execute in the loop.
+            condition: Python expression using 'ctx' that, when True, keeps the loop running.
+            exit_node: Node to transition to when the loop condition is False.
+        """
+        if not loop_nodes:
+            raise ValueError("Loop must contain at least one node")
+        for node in loop_nodes + [exit_node]:
+            if node not in self.workflow.nodes:
+                raise ValueError(f"Node '{node}' does not exist")
+        # Add transitions between loop nodes
+        for i in range(len(loop_nodes) - 1):
+            self.add_transition(from_node=loop_nodes[i], to_node=loop_nodes[i + 1])
+        # Add loop-back transition
+        self.add_transition(from_node=loop_nodes[-1], to_node=loop_nodes[0], condition=condition)
+        # Add exit transition
+        self.add_transition(from_node=loop_nodes[-1], to_node=exit_node, condition=f"not ({condition})")
+
     def set_start_node(self, name: str) -> None:
         """Set the start node of the workflow."""
         if name not in self.workflow.nodes:
@@ -288,7 +309,7 @@ class WorkflowManager:
                 )
 
     def instantiate_workflow(self) -> Workflow:
-        """Instantiate a Workflow object with full support for template_node and inputs_mapping."""
+        """Instantiate a Workflow object with full support for template_node, inputs_mapping, and loops."""
         self._ensure_dependencies()
 
         functions: Dict[str, Callable] = {}
@@ -471,10 +492,30 @@ class WorkflowManager:
             else:
                 wf.node(node_name, inputs_mapping=inputs_mapping if inputs_mapping else None)
 
+        # Detect loops by finding cycles with conditions
+        loop_nodes = []
+        loop_condition = None
+        loop_exit_node = None
+        for trans in self.workflow.workflow.transitions:
+            if isinstance(trans.to_node, str) and trans.condition:
+                if trans.to_node in loop_nodes and trans.from_node in loop_nodes:
+                    continue  # Already identified as part of loop
+                if any(t.from_node == trans.to_node and t.to_node == trans.from_node for t in self.workflow.workflow.transitions):
+                    # Found a potential loop
+                    loop_nodes.append(trans.from_node)
+                    loop_nodes.append(trans.to_node)
+                    loop_condition = trans.condition
+                elif trans.from_node in loop_nodes:
+                    # Check for exit transition
+                    if f"not ({loop_condition})" in trans.condition:
+                        loop_exit_node = trans.to_node
+
         added_nodes = set()
         for trans in self.workflow.workflow.transitions:
             from_node = trans.from_node
             if from_node not in added_nodes and from_node not in sub_workflows:
+                if loop_nodes and from_node == loop_nodes[0]:  # Start of loop
+                    wf.start_loop()
                 wf.node(from_node)
                 added_nodes.add(from_node)
             if isinstance(trans.to_node, str):
@@ -483,7 +524,10 @@ class WorkflowManager:
                 if to_nodes[0] not in added_nodes and to_nodes[0] not in sub_workflows:
                     wf.node(to_nodes[0])
                     added_nodes.add(to_nodes[0])
-                wf.then(to_nodes[0], condition=condition)
+                if loop_nodes and to_nodes[0] == loop_exit_node and loop_condition:  # End of loop
+                    wf.end_loop(condition=eval(f"lambda ctx: {loop_condition}"), next_node=to_nodes[0])
+                else:
+                    wf.then(to_nodes[0], condition=condition)
             elif all(isinstance(tn, str) for tn in trans.to_node):
                 to_nodes = trans.to_node
                 for to_node in to_nodes:
@@ -544,7 +588,7 @@ class WorkflowManager:
 
 
 async def test_workflow():
-    """Test the workflow execution."""
+    """Test the workflow execution with a loop."""
     manager = WorkflowManager()
     manager.workflow.dependencies = ["requests>=2.28.0"]
     manager.add_function(
@@ -553,9 +597,14 @@ async def test_workflow():
         code="def greet(user_name): return f'Hello, {user_name}!'",
     )
     manager.add_function(
-        name="check_condition",
+        name="check_length",
         type_="embedded",
-        code="def check_condition(user_name): return len(user_name) > 3",
+        code="def check_length(user_name): return len(user_name) < 5",
+    )
+    manager.add_function(
+        name="append_char",
+        type_="embedded",
+        code="def append_char(user_name): return user_name + 'x'",
     )
     manager.add_function(
         name="farewell",
@@ -578,19 +627,16 @@ async def test_workflow():
         inputs_mapping={"user_name": "name_input"},
     )
     manager.add_node(
-        name="format_greeting",
-        template_config={"template": "User: {{ user_name }} greeted on {{ date }}"},
-        inputs_mapping={"user_name": "name_input", "date": "lambda ctx: '2025-03-06'"},
-    )
-    manager.add_node(
-        name="branch_true",
-        function="check_condition",
+        name="check",
+        function="check_length",
         inputs_mapping={"user_name": "name_input"},
+        output="continue_loop"
     )
     manager.add_node(
-        name="branch_false",
-        function="check_condition",
+        name="append",
+        function="append_char",
         inputs_mapping={"user_name": "name_input"},
+        output="name_input"
     )
     manager.add_node(
         name="end",
@@ -598,18 +644,11 @@ async def test_workflow():
         inputs_mapping={"user_name": "name_input"},
     )
     manager.set_start_node("start")
-    manager.add_transition(
-        from_node="start",
-        to_node="format_greeting"
-    )
-    manager.add_transition(
-        from_node="format_greeting",
-        to_node=[
-            BranchCondition(to_node="branch_true", condition="ctx.get('user_name') == 'Alice'"),
-            BranchCondition(to_node="branch_false", condition="ctx.get('user_name') != 'Alice'")
-        ]
+    manager.add_loop(
+        loop_nodes=["start", "check", "append"],
+        condition="ctx.get('continue_loop', False)",
+        exit_node="end"
     )
-    manager.add_convergence_node("end")
     manager.add_observer("monitor")
     manager.save_to_yaml("workflow.yaml")