FlowEngine

llmfy.flow_engine.flow_engine

FlowEngine

A workflow engine that manages state transitions through nodes and edges.

Attributes:

Name	Type	Description
state_schema		TypedDict class defining the state structure
nodes	Dict[str, Node]	Dictionary of node name to node function
edges	List[Edge]	Dictionary of node name to list of target nodes
conditional_edges	List[Edge]	Dictionary of node name to conditional routing info
state	Dict[str, Any]	Current state of the workflow
checkpointer		Optional checkpointer for state persistence

Source code in llmfy/flow_engine/flow_engine.py

class FlowEngine:
    """
    A workflow engine that manages state transitions through nodes and edges.

    Attributes:
        state_schema: TypedDict class defining the state structure
        nodes: Dictionary of node name to node function
        edges: Dictionary of node name to list of target nodes
        conditional_edges: Dictionary of node name to conditional routing info
        state: Current state of the workflow
        checkpointer: Optional checkpointer for state persistence
    """

    def __init__(
        self,
        state_schema: type,
        checkpointer: Optional[BaseCheckpointer] = None,
    ):
        """
        Initialize the FlowEngine with a state schema.

        Args:
            state_schema: A TypedDict class defining the state structure
            checkpointer: Optional checkpointer for state persistence
        """
        self.is_built = False
        self.state_schema = state_schema
        self.nodes: Dict[str, Node] = {}
        self.edges: List[Edge] = []
        self.state: Dict[str, Any] = {}
        self._reducers = {}
        self._type_hints = {}  # Store type hints for deserialization

        # Checkpointer configuration
        self.checkpointer = checkpointer
        self._session_id: Optional[str] = None
        self._step_counter: int = 0
        self._checkpoint_enabled: bool = checkpointer is not None

        # Add special START and END nodes
        self.nodes[START] = Node(name=START, node_type=NodeType.START)
        self.nodes[END] = Node(name=END, node_type=NodeType.END)

        # Extract annotations from the state schema and validate reducers
        self._extract_state_annotations()

        # Visualizer
        self.visualizer = WorkflowVisualizer()

    def _extract_state_annotations(self):
        """Extract and validate reducer functions from the state schema."""
        hints = get_type_hints(self.state_schema, include_extras=True)

        for f, hint in hints.items():
            origin = get_origin(hint)

            # Check if it's an Annotated type
            if origin is Annotated:
                args = get_args(hint)
                if len(args) >= 2:
                    # Store the actual type (first arg)
                    self._type_hints[f] = args[0]
                    # Store the reducer (second arg)
                    reducer = args[1]
                    # Validate the reducer function
                    self._validate_reducer(f, reducer)
                    self._reducers[f] = reducer
                else:
                    self._type_hints[f] = args[0] if args else hint
                    self._reducers[f] = None
            else:
                # Store the type hint directly
                self._type_hints[f] = hint
                self._reducers[f] = None

    def _validate_reducer(self, field_name: str, reducer: Callable):
        """
        Validate that a reducer function has the correct signature.

        Args:
            field_name: Name of the field this reducer is for
            reducer: The reducer function to validate

        Raises:
            LLMfyException: If the reducer doesn't have exactly 2 parameters
        """
        if not callable(reducer):
            raise LLMfyException(
                f"Reducer for field '{field_name}' must be callable, got {type(reducer)}"
            )

        sig = inspect.signature(reducer)
        params = list(sig.parameters.values())

        if len(params) != 2:
            raise LLMfyException(
                f"Reducer for field '{field_name}' must have exactly 2 parameters "
                f"(old_value, new_value), but has {len(params)} parameters"
            )

    def _deserialize_state(self, state_dict: Dict[str, Any]) -> Dict[str, Any]:
        """
        Deserialize state dictionary and reconstruct objects based on TypedDict schema.

        Args:
            state_dict: Dictionary with potentially serialized objects

        Returns:
            Dictionary with objects reconstructed to their proper types
        """
        deserialized = {}

        for field_name, value in state_dict.items():
            # Get the expected type for this field
            expected_type = self._type_hints.get(field_name)

            if expected_type is None:
                # No type hint available, keep as is
                deserialized[field_name] = value
                continue

            # Reconstruct the value based on expected type
            deserialized[field_name] = self._reconstruct_value(value, expected_type)

        return deserialized

    def _reconstruct_value(self, value: Any, expected_type: Any) -> Any:
        """
        Reconstruct a value to match the expected type.

        Args:
            value: The value to reconstruct (may be dict, list, or primitive)
            expected_type: The expected type from TypedDict

        Returns:
            Reconstructed value
        """
        if value is None:
            return None

        # Handle Annotated types - extract the actual type
        origin = get_origin(expected_type)
        if origin is Annotated:
            args = get_args(expected_type)
            if args:
                # First arg is the actual type
                expected_type = args[0]
                origin = get_origin(expected_type)

        # Handle List types
        if origin in (list, List):
            if not isinstance(value, list):
                return value

            # Get the element type
            args = get_args(expected_type)
            if args:
                element_type = args[0]
                # Reconstruct each element
                return [self._reconstruct_value(item, element_type) for item in value]
            return value

        # Handle Dict types
        if origin in (dict, Dict):
            if not isinstance(value, dict):
                return value

            args = get_args(expected_type)
            if len(args) >= 2:
                key_type, value_type = args[0], args[1]
                return {
                    self._reconstruct_value(k, key_type): self._reconstruct_value(
                        v, value_type
                    )
                    for k, v in value.items()
                }
            return value

        # Handle custom class objects
        if isinstance(value, dict) and hasattr(expected_type, "__mro__"):
            # Check if this looks like a serialized object
            if "__dict__" in value or self._is_serialized_object(value):
                return self._reconstruct_object(value, expected_type)

        # Handle primitive types and already correct types
        if isinstance(value, expected_type):
            return value

        # Try to cast to expected type if it's a simple type
        try:
            if expected_type in (int, float, str, bool):
                return expected_type(value)  # type: ignore
        except (LLMfyException, TypeError):
            pass

        return value

    def _is_serialized_object(self, value: dict) -> bool:
        """Check if a dict looks like a serialized object."""
        # Common patterns for serialized objects
        return (
            "__type__" in value
            or "__class__" in value
            or "__module__" in value
            or (
                isinstance(value, dict)
                and not any(k.startswith("_") for k in value.keys())
                and len(value) > 0
            )
        )

    def _reconstruct_object(self, data: dict, cls: type) -> Any:
        """
        Reconstruct an object from a dictionary.

        Args:
            data: Dictionary containing object data
            cls: The class to instantiate

        Returns:
            Reconstructed object instance
        """
        try:
            # Handle checkpointer serialization format with __type__, __module__, data
            if "__type__" in data and "__module__" in data and "data" in data:
                obj_data = data["data"]
            # Handle format with __dict__
            elif "__dict__" in data:
                obj_data = data["__dict__"]
            else:
                obj_data = data

            # Try to instantiate the class
            if hasattr(cls, "__init__"):
                # Get __init__ signature
                sig = inspect.signature(cls.__init__)
                params = list(sig.parameters.keys())[1:]  # Skip 'self'

                if len(params) == 0:
                    # No-arg constructor, set attributes after
                    obj = cls()
                    for key, value in obj_data.items():
                        if not key.startswith("_"):
                            setattr(obj, key, value)
                    return obj
                else:
                    # Try to match constructor parameters
                    init_args = {}
                    for param in params:
                        if param in obj_data:
                            init_args[param] = obj_data[param]

                    obj = cls(**init_args)

                    # Set remaining attributes
                    for key, value in obj_data.items():
                        if key not in init_args and not key.startswith("_"):
                            setattr(obj, key, value)

                    return obj
            else:
                # Fallback: create instance and set attributes
                obj = cls()
                for key, value in obj_data.items():
                    if not key.startswith("_"):
                        setattr(obj, key, value)
                return obj

        except Exception as _:
            # If reconstruction fails, return the dict
            # This allows the workflow to continue even if object reconstruction fails
            return data

    def add_node(
        self,
        name: str,
        func: Callable,
        stream: bool = False,
    ):
        """
        Add a node to the workflow.

        Args:
            name (str): Name of the node
            func (Callable): Function to execute (can be sync or async)
            stream (bool): Node is use stream or not, if node use streaming set to True. Defaults to False.
        """
        if name in [START, END]:
            raise LLMfyException(f"Cannot add node with reserved name: {name}")

        # Determine if this is a conditional node (will be set when conditional edge is added)
        node = Node(name=name, node_type=NodeType.FUNCTION, func=func, stream=stream)
        self.nodes[name] = node

    def add_edge(self, source: str, target: str):
        """
        Add an edge connecting two nodes.

        Args:
            source: Source node name (can be START)
            target: Target node name (can be END)
        """
        # Validation: START cannot be a target
        if target == START:
            raise LLMfyException("START cannot be a target node")

        # Validation: END cannot be a source
        if source == END:
            raise LLMfyException("END cannot be a source node")

        # Validation: edge cannot target itself
        if source == target:
            raise LLMfyException("Source same as target, edge cannot target itself")

        # Create edge
        edge = Edge(source=source, targets=target)
        self.edges.append(edge)

        # Update node connections
        if source in self.nodes:
            self.nodes[source].targets.append(target)
        if target in self.nodes:
            self.nodes[target].sources.append(source)

    def add_conditional_edge(
        self,
        source: str,
        targets: List[str],
        condition_func: Callable,
    ):
        """
        Add a conditional edge that routes to different nodes based on a condition.

        Args:
            source: Source node name
            targets: List of possible target nodes (can include END)
            condition_func: Function that takes state and returns target node name
        """
        # Validation: START cannot be in targets
        if START in targets:
            raise LLMfyException("START cannot be a target in conditional edges")

        # Validation: END cannot be source
        if source == END:
            raise LLMfyException("END cannot be a source node")

        # Create conditional edge
        edge = Edge(source=source, targets=targets, condition=condition_func)
        self.edges.append(edge)

        # Mark source node as conditional
        if source in self.nodes and source not in [START, END]:
            self.nodes[source].node_type = NodeType.CONDITIONAL

        # Update node connections
        if source in self.nodes:
            self.nodes[source].targets.extend(targets)
        for to_node in targets:
            if to_node in self.nodes:
                self.nodes[to_node].sources.append(source)

    def _update_state(self, updates: Dict[str, Any]):
        """
        Update the workflow state with new values.

        Uses reducer functions if available, otherwise replaces values.

        Args:
            updates: Dictionary of state updates
        """
        for key, new_value in updates.items():
            if key in self._reducers and self._reducers[key] is not None:
                # Use the reducer function
                reducer = self._reducers[key]
                old_value = self.state.get(key)
                self.state[key] = reducer(old_value, new_value)
            else:
                # Replace the value
                self.state[key] = new_value

    def _validate_workflow(self):
        """
        Validate the workflow structure before execution.

        Raises:
            LLMfyException: If the workflow has structural issues
        """
        # Validation 1: START must have at least one outgoing edge
        start_edges = [e for e in self.edges if e.source == START]
        if not start_edges:
            raise LLMfyException(
                "No edge from START node. Use flow.add_edge(START, 'node_name')"
            )

        # Validation 2: At least one path must lead to END
        end_edges = [e for e in self.edges if END in e.targets]
        if not end_edges:
            raise LLMfyException(
                "No edge to END node. At least one execution path must reach END. "
                "Use flow.add_edge('node_name', END) or include END in conditional targets."
            )

        # Collect all referenced nodes
        all_referenced_nodes = set()

        for edge in self.edges:
            # Add source (except special nodes)
            if edge.source not in [START, END]:
                all_referenced_nodes.add(edge.source)

            # Add targets (except special nodes)
            for target in edge.targets:
                if target not in [START, END]:
                    all_referenced_nodes.add(target)

        # Validation 3: All referenced nodes must be defined
        defined_nodes = set(self.nodes.keys()) - {START, END}
        undefined_nodes = all_referenced_nodes - defined_nodes
        if undefined_nodes:
            raise LLMfyException(
                f"Referenced nodes are not defined: {', '.join(sorted(undefined_nodes))}. "
                f"Use flow.add_node() to define them."
            )

        # Validation 4: Conditional edges - validate that condition function returns valid targets
        for edge in self.edges:
            if edge.condition is not None:
                # This is a conditional edge
                source_node_name = edge.source
                # valid_targets = set(edge.targets)

                # Check that all targets exist (except END)
                for target in edge.targets:
                    if target != END and target not in self.nodes:
                        raise LLMfyException(
                            f"Conditional edge from '{source_node_name}' references "
                            f"undefined target '{target}'"
                        )

                # Can't validate the return value until runtime, but we document it
                # The runtime validation happens in _get_next_node

        # Validation 5: Detect nodes with multiple non-conditional edges
        edge_counts = {}
        for edge in self.edges:
            if edge.condition is None:  # Only check non-conditional edges
                if edge.source not in edge_counts:
                    edge_counts[edge.source] = []
                edge_counts[edge.source].extend(edge.targets)

        for source, targets in edge_counts.items():
            if len(targets) > 1:
                raise LLMfyException(
                    f"Node '{source}' has multiple outgoing edges ({len(targets)}) but no conditional logic. "
                    f"Use add_conditional_edge() instead of multiple add_edge() calls."
                )

        # Warning: Detect unreachable nodes
        unreachable_nodes = defined_nodes - all_referenced_nodes
        if unreachable_nodes:
            import warnings

            warnings.warn(
                f"Some nodes are defined but not reachable: {', '.join(sorted(unreachable_nodes))}",
                UserWarning,
            )

    async def _save_checkpoint(self, node_name: str):
        """
        Save current state as a checkpoint.

        Args:
            node_name: Name of the node that just executed
        """
        if not self._checkpoint_enabled or self.checkpointer is None:
            return

        checkpoint_id = str(uuid.uuid4())
        metadata = CheckpointMetadata(
            checkpoint_id=checkpoint_id,
            session_id=self._session_id,  # type: ignore
            timestamp=datetime.now(timezone.utc),
            node_name=node_name,
            step=self._step_counter,
        )

        checkpoint = Checkpoint(metadata=metadata, state=deepcopy(self.state))

        await self.checkpointer.save(checkpoint)

    async def _execute_node(self, node_name: str) -> Dict[str, Any]:
        """
        Execute a node function with the current state.

        Args:
            node_name: Name of the node to execute

        Returns:
            Dictionary of state updates from the node
        """
        if node_name not in self.nodes:
            raise LLMfyException(f"Node '{node_name}' not found")

        node = self.nodes[node_name]
        func = node.func

        if func is None:
            raise LLMfyException(f"Node '{node_name}' has no function defined")

        # Check if the function is async or sync
        if inspect.iscoroutinefunction(func):
            result = await func(self.state)
        else:
            result = func(self.state)

        # Return empty dict if node doesn't return anything
        if result is None:
            return {}

        return result

    async def _execute_stream_node(self, node_name: str, func: Optional[Callable]):
        """
        Execute a stream node function with the current state.

        Args:
            node_name: Name of the node to execute
        """
        if func is None:
            raise LLMfyException(f"Node '{node_name}' has no function defined")

        # Check if the function is async generator or generator
        if inspect.isasyncgenfunction(func):
            async for chunk in func(self.state):
                if isinstance(chunk, NodeStreamResponse):
                    yield chunk
                else:
                    raise LLMfyException(
                        f"Stream response in node: '{node_name}' must use `NodeStreamResponse`"
                    )

        elif inspect.isgeneratorfunction(func):
            for chunk in func(self.state):
                if isinstance(chunk, NodeStreamResponse):
                    yield chunk
                else:
                    raise LLMfyException(
                        f"Stream response in node: '{node_name}' must use `NodeStreamResponse`"
                    )

        else:
            raise LLMfyException(
                f"Function in node: '{node_name}' is not stream. Please yield `NodeStreamResponse`."
            )

    async def _get_next_node(self, current_node: str) -> Union[str, None]:
        """
        Determine the next node to execute.

        Args:
            current_node: Current node name

        Returns:
            Next node name or None if END
        """
        # Find edges from current node
        outgoing_edges = [e for e in self.edges if e.source == current_node]

        if not outgoing_edges:
            return None

        # Should only be one edge per node (validated in _validate_workflow)
        edge = outgoing_edges[0]

        # Check if this is a conditional edge
        if edge.condition is not None:
            condition_func = edge.condition

            # Execute condition function (can be sync or async)
            if inspect.iscoroutinefunction(condition_func):
                next_node = await condition_func(self.state)
            else:
                next_node = condition_func(self.state)

            # Validate that the returned node is in the targets
            if next_node not in edge.targets:
                raise LLMfyException(
                    f"Condition function returned '{next_node}' which is not in targets: {edge.targets}"
                )

            return next_node if next_node != END else None

        # Regular edge - single target
        next_node = edge.targets[0]
        return next_node if next_node != END else None

    def build(self):
        """
        Build FlowEngine workflow.

        Returns:
            FlowEngine: Workflow built.
        """
        # Validate workflow structure before execution
        self._validate_workflow()

        # Set is built true
        self.is_built = True

        return self

    async def invoke(
        self,
        apply_state: Optional[Dict[str, Any]] = None,
        session_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Execute the workflow starting from START node or continue from last checkpoint.

        Args:
            apply_state: Optional state updates to apply. If continuing from checkpoint,
                these are merged with the checkpoint state using reducers.
            session_id: Session ID for checkpoint management. If provided and a checkpoint
                exists, continues from last checkpoint. If None, always starts fresh.

        Returns:
            Final state after workflow execution
        """
        # Check build
        if not self.is_built:
            raise LLMfyException("Build first. Use `your_flow.build()`")

        # Set thread ID
        self._session_id = session_id or str(uuid.uuid4())
        self._step_counter = 0

        # Initialize state
        if apply_state is None:
            apply_state = {}

        # Try to load from last checkpoint if session_id is provided
        loaded_checkpoint = None
        if session_id and self.checkpointer:
            loaded_checkpoint = await self.checkpointer.load(session_id)

        if loaded_checkpoint:
            # Continue from checkpoint - deserialize objects
            raw_state = loaded_checkpoint.state
            self.state = self._deserialize_state(raw_state)
            self._step_counter = loaded_checkpoint.metadata.step

            # Apply initial_state as updates to checkpoint state
            if apply_state:
                self._update_state(apply_state)

            # Find where to resume (next node after last completed node)
            current_node = await self._get_next_node(
                loaded_checkpoint.metadata.node_name
            )

            # If no next node (workflow was completed), start from beginning
            if current_node is None:
                start_edges = [e for e in self.edges if e.source == START]
                current_node = start_edges[0].targets[0]
        else:
            # Start fresh - no checkpoint found or no session_id provided
            self.state = deepcopy(apply_state)

            # Find the starting node from START edges
            start_edges = [e for e in self.edges if e.source == START]
            current_node = start_edges[0].targets[0]

        # Save initial checkpoint
        if self._checkpoint_enabled:
            await self._save_checkpoint(START)

        # Execute workflow
        while current_node is not None:
            # Increment step counter
            self._step_counter += 1

            if current_node not in self.nodes:
                raise LLMfyException(f"Node '{current_node}' not found")

            node = self.nodes[current_node]

            if node.stream:
                # Handle stream node
                async for chunk in self._execute_stream_node(
                    current_node,
                    func=node.func,
                ):
                    # NodeStreamType.RESULT is always send at last stream
                    if chunk.type == NodeStreamType.RESULT:
                        # In invoke update state is in last stream result.
                        # Update state with results
                        if chunk.state:
                            self._update_state(chunk.state)

                        # Save checkpoint
                        if self._checkpoint_enabled:
                            await self._save_checkpoint(current_node)

                # Determine next node (now async)
                current_node = await self._get_next_node(current_node)

            else:
                # Handle non-streaming node
                # Execute current node
                updates = await self._execute_node(current_node)

                # Update state with results
                if updates:
                    self._update_state(updates)

                # Save checkpoint after node execution
                if self._checkpoint_enabled:
                    await self._save_checkpoint(current_node)

                # Determine next node (now async)
                current_node = await self._get_next_node(current_node)

        return self.state

    async def stream(
        self,
        apply_state: Optional[Dict[str, Any]] = None,
        session_id: Optional[str] = None,
        stream_callback: Optional[Callable] = None,
    ):
        """
        Execute the workflow in streaming mode, starting from START node or continue from last checkpoint.

        Args:
            apply_state: Optional state updates to apply. If continuing from checkpoint,
                these are merged with the checkpoint state using reducers.
            session_id: Session ID for checkpoint management. If provided and a checkpoint
                exists, continues from last checkpoint. If None, always starts fresh.
            stream_callback: Optional callback function for handling streaming chunks (content only)

        Returns:
            Final state after workflow execution
        """
        # Check build
        if not self.is_built:
            raise LLMfyException("Build first. Use `your_flow.build()`")

        # Set thread ID
        self._session_id = session_id or str(uuid.uuid4())
        self._step_counter = 0

        # Initialize state
        if apply_state is None:
            apply_state = {}

        # Try to load from last checkpoint if session_id is provided
        loaded_checkpoint = None
        if session_id and self.checkpointer:
            loaded_checkpoint = await self.checkpointer.load(session_id)

        if loaded_checkpoint:
            # Continue from checkpoint - deserialize objects
            raw_state = loaded_checkpoint.state
            self.state = self._deserialize_state(raw_state)
            self._step_counter = loaded_checkpoint.metadata.step

            # Apply apply_state as updates to checkpoint state
            if apply_state:
                self._update_state(apply_state)

            # Find where to resume (next node after last completed node)
            current_node = await self._get_next_node(
                loaded_checkpoint.metadata.node_name
            )

            # If no next node (workflow was completed), start from beginning
            if current_node is None:
                start_edges = [e for e in self.edges if e.source == START]
                current_node = start_edges[0].targets[0]
        else:
            # Start fresh - no checkpoint found or no session_id provided
            self.state = deepcopy(apply_state)

            # Find the starting node from START edges
            start_edges = [e for e in self.edges if e.source == START]
            current_node = start_edges[0].targets[0]

        # Save initial checkpoint
        if self._checkpoint_enabled:
            await self._save_checkpoint(START)

        # START
        response = FlowEngineStreamResponse()
        response.type = FlowEngineStreamType.START
        response.node = START
        response.content = None
        response.state = self.state
        yield response

        # Execute workflow
        while current_node is not None:
            # Increment step counter
            self._step_counter += 1

            if current_node not in self.nodes:
                raise LLMfyException(f"Node '{current_node}' not found")

            node = self.nodes[current_node]

            if node.stream:
                # Handle stream node
                async for chunk in self._execute_stream_node(
                    current_node,
                    func=node.func,
                ):
                    # NodeStreamType.RESULT is always send at last stream
                    if chunk.type == NodeStreamType.RESULT:
                        # Update state with results
                        if chunk.state:
                            self._update_state(chunk.state)

                        # Save checkpoint
                        if self._checkpoint_enabled:
                            await self._save_checkpoint(current_node)

                        # NODE RESULT
                        response.type = FlowEngineStreamType.RESULT
                        response.node = current_node
                        response.content = chunk.content
                        response.state = self.state
                        yield response

                    else:
                        # Not NodeStreamType.RESULT
                        # NODE STREAM
                        response.type = FlowEngineStreamType.STREAM
                        response.node = current_node
                        response.content = chunk.content
                        response.state = self.state
                        yield response

                # Determine next node (now async)
                current_node = await self._get_next_node(current_node)

            else:
                # Handle non-streaming node
                updates = await self._execute_node(current_node)

                # Update state with results
                if updates:
                    self._update_state(updates)

                # Save checkpoint after node execution
                if self._checkpoint_enabled:
                    await self._save_checkpoint(current_node)

                # NODE RESULT
                response.type = FlowEngineStreamType.RESULT
                response.node = current_node
                response.content = updates
                response.state = self.state
                yield response

                # Determine next node (now async)
                current_node = await self._get_next_node(current_node)

    async def get_state(self, session_id: str) -> Optional[Dict[str, Any]]:
        """
        Get the current state for a thread from the last checkpoint.

        Args:
            session_id: The session ID

        Returns:
            The state if checkpoint exists, None otherwise
        """
        if not self.checkpointer:
            raise LLMfyException("No checkpointer configured")

        checkpoint = await self.checkpointer.load(session_id)
        if checkpoint:
            # Deserialize the state to reconstruct objects
            return self._deserialize_state(checkpoint.state)
        return None

    async def list_checkpoints(
        self,
        session_id: str,
        limit: int = 10,
    ) -> list[Checkpoint]:
        """
        List checkpoints for a specific thread.

        Args:
            session_id: The session ID
            limit: Maximum number of checkpoints to return

        Returns:
            List of checkpoints, newest first
        """
        if not self.checkpointer:
            raise LLMfyException("No checkpointer configured")

        return await self.checkpointer.list(session_id, limit)

    async def get_checkpoint(
        self,
        session_id: str,
        checkpoint_id: Optional[str] = None,
    ) -> Optional[Checkpoint]:
        """
        Get a specific checkpoint or the latest checkpoint for a thread.

        Args:
            session_id: The session ID
            checkpoint_id: Optional checkpoint ID, or None for latest

        Returns:
            The checkpoint if found, None otherwise
        """
        if not self.checkpointer:
            raise LLMfyException("No checkpointer configured")

        return await self.checkpointer.load(session_id, checkpoint_id)

    async def delete_checkpoints(
        self,
        session_id: str,
        checkpoint_id: Optional[str] = None,
    ):
        """
        Delete checkpoint(s) for a thread.

        Args:
            session_id: The session ID
            checkpoint_id: Optional checkpoint ID to delete, or None to delete all
        """
        if not self.checkpointer:
            raise LLMfyException("No checkpointer configured")

        await self.checkpointer.delete(session_id, checkpoint_id)

    async def reset_session(self, session_id: str):
        """
        Reset a session by deleting all its checkpoints.
        This allows starting fresh with the same session_id.

        Args:
            session: The session ID to reset
        """
        if not self.checkpointer:
            raise LLMfyException("No checkpointer configured")

        await self.checkpointer.delete(session_id)

    def details(self) -> str:
        """
        Generate a simple details text visualization of the workflow.

        Returns:
            String representation of the workflow graph
        """
        # Check build
        if not self.is_built:
            raise LLMfyException("Build first. Use `your_flow.build()`")

        lines = ["Workflow Graph:", "=" * 50]

        # Show START connections
        start_edges = [e for e in self.edges if e.source == START]
        for edge in start_edges:
            for target in edge.targets:
                lines.append(f"START -> {target}")

        # Show all function nodes
        function_nodes = [
            n for n in self.nodes.values() if n.node_type == NodeType.FUNCTION
        ]
        if function_nodes:
            lines.append("\nFunction Nodes:")
            for node in function_nodes:
                lines.append(f"  - {node.name}")

        # Show conditional nodes
        conditional_nodes = [
            n for n in self.nodes.values() if n.node_type == NodeType.CONDITIONAL
        ]
        if conditional_nodes:
            lines.append("\nConditional Nodes:")
            for node in conditional_nodes:
                lines.append(f"  - {node.name}")

        # Show regular edges
        regular_edges = [
            e for e in self.edges if e.condition is None and e.source != START
        ]
        if regular_edges:
            lines.append("\nRegular Edges:")
            for edge in regular_edges:
                for target in edge.targets:
                    lines.append(f"  {edge.source} -> {target}")

        # Show conditional edges
        conditional_edges = [e for e in self.edges if e.condition is not None]
        if conditional_edges:
            lines.append("\nConditional Edges:")
            for edge in conditional_edges:
                targets = ", ".join(edge.targets)
                lines.append(f"  {edge.source} ->? [{targets}]")

        return "\n".join(lines)

    def visualize(self) -> str:
        """
        Visualize workflow diagram.
        Generate Mermaid diagram url.

        Returns:
            str: Mermaid URL.
        """
        # Check build
        if not self.is_built:
            raise LLMfyException("Build first. Use `your_flow.build()`")

        mermaid_code = self.visualizer.create_mermaid_diagram(self)
        return self.visualizer.generate_diagram_url(mermaid_code)

is_built = False instance-attribute

state_schema = state_schema instance-attribute

nodes = {} instance-attribute

edges = [] instance-attribute

state = {} instance-attribute

checkpointer = checkpointer instance-attribute

visualizer = WorkflowVisualizer() instance-attribute

__init__(state_schema, checkpointer=None)

Initialize the FlowEngine with a state schema.

Parameters:

Name	Type	Description	Default
state_schema	type	A TypedDict class defining the state structure	required
checkpointer	Optional[BaseCheckpointer]	Optional checkpointer for state persistence	None

Source code in llmfy/flow_engine/flow_engine.py

def __init__(
    self,
    state_schema: type,
    checkpointer: Optional[BaseCheckpointer] = None,
):
    """
    Initialize the FlowEngine with a state schema.

    Args:
        state_schema: A TypedDict class defining the state structure
        checkpointer: Optional checkpointer for state persistence
    """
    self.is_built = False
    self.state_schema = state_schema
    self.nodes: Dict[str, Node] = {}
    self.edges: List[Edge] = []
    self.state: Dict[str, Any] = {}
    self._reducers = {}
    self._type_hints = {}  # Store type hints for deserialization

    # Checkpointer configuration
    self.checkpointer = checkpointer
    self._session_id: Optional[str] = None
    self._step_counter: int = 0
    self._checkpoint_enabled: bool = checkpointer is not None

    # Add special START and END nodes
    self.nodes[START] = Node(name=START, node_type=NodeType.START)
    self.nodes[END] = Node(name=END, node_type=NodeType.END)

    # Extract annotations from the state schema and validate reducers
    self._extract_state_annotations()

    # Visualizer
    self.visualizer = WorkflowVisualizer()

add_node(name, func, stream=False)

Add a node to the workflow.

Parameters:

Name	Type	Description	Default
name	str	Name of the node	required
func	Callable	Function to execute (can be sync or async)	required
stream	bool	Node is use stream or not, if node use streaming set to True. Defaults to False.	False

Source code in llmfy/flow_engine/flow_engine.py

def add_node(
    self,
    name: str,
    func: Callable,
    stream: bool = False,
):
    """
    Add a node to the workflow.

    Args:
        name (str): Name of the node
        func (Callable): Function to execute (can be sync or async)
        stream (bool): Node is use stream or not, if node use streaming set to True. Defaults to False.
    """
    if name in [START, END]:
        raise LLMfyException(f"Cannot add node with reserved name: {name}")

    # Determine if this is a conditional node (will be set when conditional edge is added)
    node = Node(name=name, node_type=NodeType.FUNCTION, func=func, stream=stream)
    self.nodes[name] = node

add_edge(source, target)

Add an edge connecting two nodes.

Parameters:

Name	Type	Description	Default
source	str	Source node name (can be START)	required
target	str	Target node name (can be END)	required

Source code in llmfy/flow_engine/flow_engine.py

def add_edge(self, source: str, target: str):
    """
    Add an edge connecting two nodes.

    Args:
        source: Source node name (can be START)
        target: Target node name (can be END)
    """
    # Validation: START cannot be a target
    if target == START:
        raise LLMfyException("START cannot be a target node")

    # Validation: END cannot be a source
    if source == END:
        raise LLMfyException("END cannot be a source node")

    # Validation: edge cannot target itself
    if source == target:
        raise LLMfyException("Source same as target, edge cannot target itself")

    # Create edge
    edge = Edge(source=source, targets=target)
    self.edges.append(edge)

    # Update node connections
    if source in self.nodes:
        self.nodes[source].targets.append(target)
    if target in self.nodes:
        self.nodes[target].sources.append(source)

add_conditional_edge(source, targets, condition_func)

Add a conditional edge that routes to different nodes based on a condition.

Parameters:

Name	Type	Description	Default
source	str	Source node name	required
targets	List[str]	List of possible target nodes (can include END)	required
condition_func	Callable	Function that takes state and returns target node name	required

Source code in llmfy/flow_engine/flow_engine.py

def add_conditional_edge(
    self,
    source: str,
    targets: List[str],
    condition_func: Callable,
):
    """
    Add a conditional edge that routes to different nodes based on a condition.

    Args:
        source: Source node name
        targets: List of possible target nodes (can include END)
        condition_func: Function that takes state and returns target node name
    """
    # Validation: START cannot be in targets
    if START in targets:
        raise LLMfyException("START cannot be a target in conditional edges")

    # Validation: END cannot be source
    if source == END:
        raise LLMfyException("END cannot be a source node")

    # Create conditional edge
    edge = Edge(source=source, targets=targets, condition=condition_func)
    self.edges.append(edge)

    # Mark source node as conditional
    if source in self.nodes and source not in [START, END]:
        self.nodes[source].node_type = NodeType.CONDITIONAL

    # Update node connections
    if source in self.nodes:
        self.nodes[source].targets.extend(targets)
    for to_node in targets:
        if to_node in self.nodes:
            self.nodes[to_node].sources.append(source)

build()

Build FlowEngine workflow.

Returns:

Name	Type	Description
FlowEngine		Workflow built.

Source code in llmfy/flow_engine/flow_engine.py

def build(self):
    """
    Build FlowEngine workflow.

    Returns:
        FlowEngine: Workflow built.
    """
    # Validate workflow structure before execution
    self._validate_workflow()

    # Set is built true
    self.is_built = True

    return self

invoke(apply_state=None, session_id=None) async

Execute the workflow starting from START node or continue from last checkpoint.

Parameters:

Name	Type	Description	Default
apply_state	Optional[Dict[str, Any]]	Optional state updates to apply. If continuing from checkpoint, these are merged with the checkpoint state using reducers.	None
session_id	Optional[str]	Session ID for checkpoint management. If provided and a checkpoint exists, continues from last checkpoint. If None, always starts fresh.	None

Returns:

Type	Description
Dict[str, Any]	Final state after workflow execution

Source code in llmfy/flow_engine/flow_engine.py

async def invoke(
    self,
    apply_state: Optional[Dict[str, Any]] = None,
    session_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Execute the workflow starting from START node or continue from last checkpoint.

    Args:
        apply_state: Optional state updates to apply. If continuing from checkpoint,
            these are merged with the checkpoint state using reducers.
        session_id: Session ID for checkpoint management. If provided and a checkpoint
            exists, continues from last checkpoint. If None, always starts fresh.

    Returns:
        Final state after workflow execution
    """
    # Check build
    if not self.is_built:
        raise LLMfyException("Build first. Use `your_flow.build()`")

    # Set thread ID
    self._session_id = session_id or str(uuid.uuid4())
    self._step_counter = 0

    # Initialize state
    if apply_state is None:
        apply_state = {}

    # Try to load from last checkpoint if session_id is provided
    loaded_checkpoint = None
    if session_id and self.checkpointer:
        loaded_checkpoint = await self.checkpointer.load(session_id)

    if loaded_checkpoint:
        # Continue from checkpoint - deserialize objects
        raw_state = loaded_checkpoint.state
        self.state = self._deserialize_state(raw_state)
        self._step_counter = loaded_checkpoint.metadata.step

        # Apply initial_state as updates to checkpoint state
        if apply_state:
            self._update_state(apply_state)

        # Find where to resume (next node after last completed node)
        current_node = await self._get_next_node(
            loaded_checkpoint.metadata.node_name
        )

        # If no next node (workflow was completed), start from beginning
        if current_node is None:
            start_edges = [e for e in self.edges if e.source == START]
            current_node = start_edges[0].targets[0]
    else:
        # Start fresh - no checkpoint found or no session_id provided
        self.state = deepcopy(apply_state)

        # Find the starting node from START edges
        start_edges = [e for e in self.edges if e.source == START]
        current_node = start_edges[0].targets[0]

    # Save initial checkpoint
    if self._checkpoint_enabled:
        await self._save_checkpoint(START)

    # Execute workflow
    while current_node is not None:
        # Increment step counter
        self._step_counter += 1

        if current_node not in self.nodes:
            raise LLMfyException(f"Node '{current_node}' not found")

        node = self.nodes[current_node]

        if node.stream:
            # Handle stream node
            async for chunk in self._execute_stream_node(
                current_node,
                func=node.func,
            ):
                # NodeStreamType.RESULT is always send at last stream
                if chunk.type == NodeStreamType.RESULT:
                    # In invoke update state is in last stream result.
                    # Update state with results
                    if chunk.state:
                        self._update_state(chunk.state)

                    # Save checkpoint
                    if self._checkpoint_enabled:
                        await self._save_checkpoint(current_node)

            # Determine next node (now async)
            current_node = await self._get_next_node(current_node)

        else:
            # Handle non-streaming node
            # Execute current node
            updates = await self._execute_node(current_node)

            # Update state with results
            if updates:
                self._update_state(updates)

            # Save checkpoint after node execution
            if self._checkpoint_enabled:
                await self._save_checkpoint(current_node)

            # Determine next node (now async)
            current_node = await self._get_next_node(current_node)

    return self.state

stream(apply_state=None, session_id=None, stream_callback=None) async

Execute the workflow in streaming mode, starting from START node or continue from last checkpoint.

Parameters:

Name	Type	Description	Default
apply_state	Optional[Dict[str, Any]]	Optional state updates to apply. If continuing from checkpoint, these are merged with the checkpoint state using reducers.	None
session_id	Optional[str]	Session ID for checkpoint management. If provided and a checkpoint exists, continues from last checkpoint. If None, always starts fresh.	None
stream_callback	Optional[Callable]	Optional callback function for handling streaming chunks (content only)	None

Returns:

Type	Description
	Final state after workflow execution

Source code in llmfy/flow_engine/flow_engine.py

async def stream(
    self,
    apply_state: Optional[Dict[str, Any]] = None,
    session_id: Optional[str] = None,
    stream_callback: Optional[Callable] = None,
):
    """
    Execute the workflow in streaming mode, starting from START node or continue from last checkpoint.

    Args:
        apply_state: Optional state updates to apply. If continuing from checkpoint,
            these are merged with the checkpoint state using reducers.
        session_id: Session ID for checkpoint management. If provided and a checkpoint
            exists, continues from last checkpoint. If None, always starts fresh.
        stream_callback: Optional callback function for handling streaming chunks (content only)

    Returns:
        Final state after workflow execution
    """
    # Check build
    if not self.is_built:
        raise LLMfyException("Build first. Use `your_flow.build()`")

    # Set thread ID
    self._session_id = session_id or str(uuid.uuid4())
    self._step_counter = 0

    # Initialize state
    if apply_state is None:
        apply_state = {}

    # Try to load from last checkpoint if session_id is provided
    loaded_checkpoint = None
    if session_id and self.checkpointer:
        loaded_checkpoint = await self.checkpointer.load(session_id)

    if loaded_checkpoint:
        # Continue from checkpoint - deserialize objects
        raw_state = loaded_checkpoint.state
        self.state = self._deserialize_state(raw_state)
        self._step_counter = loaded_checkpoint.metadata.step

        # Apply apply_state as updates to checkpoint state
        if apply_state:
            self._update_state(apply_state)

        # Find where to resume (next node after last completed node)
        current_node = await self._get_next_node(
            loaded_checkpoint.metadata.node_name
        )

        # If no next node (workflow was completed), start from beginning
        if current_node is None:
            start_edges = [e for e in self.edges if e.source == START]
            current_node = start_edges[0].targets[0]
    else:
        # Start fresh - no checkpoint found or no session_id provided
        self.state = deepcopy(apply_state)

        # Find the starting node from START edges
        start_edges = [e for e in self.edges if e.source == START]
        current_node = start_edges[0].targets[0]

    # Save initial checkpoint
    if self._checkpoint_enabled:
        await self._save_checkpoint(START)

    # START
    response = FlowEngineStreamResponse()
    response.type = FlowEngineStreamType.START
    response.node = START
    response.content = None
    response.state = self.state
    yield response

    # Execute workflow
    while current_node is not None:
        # Increment step counter
        self._step_counter += 1

        if current_node not in self.nodes:
            raise LLMfyException(f"Node '{current_node}' not found")

        node = self.nodes[current_node]

        if node.stream:
            # Handle stream node
            async for chunk in self._execute_stream_node(
                current_node,
                func=node.func,
            ):
                # NodeStreamType.RESULT is always send at last stream
                if chunk.type == NodeStreamType.RESULT:
                    # Update state with results
                    if chunk.state:
                        self._update_state(chunk.state)

                    # Save checkpoint
                    if self._checkpoint_enabled:
                        await self._save_checkpoint(current_node)

                    # NODE RESULT
                    response.type = FlowEngineStreamType.RESULT
                    response.node = current_node
                    response.content = chunk.content
                    response.state = self.state
                    yield response

                else:
                    # Not NodeStreamType.RESULT
                    # NODE STREAM
                    response.type = FlowEngineStreamType.STREAM
                    response.node = current_node
                    response.content = chunk.content
                    response.state = self.state
                    yield response

            # Determine next node (now async)
            current_node = await self._get_next_node(current_node)

        else:
            # Handle non-streaming node
            updates = await self._execute_node(current_node)

            # Update state with results
            if updates:
                self._update_state(updates)

            # Save checkpoint after node execution
            if self._checkpoint_enabled:
                await self._save_checkpoint(current_node)

            # NODE RESULT
            response.type = FlowEngineStreamType.RESULT
            response.node = current_node
            response.content = updates
            response.state = self.state
            yield response

            # Determine next node (now async)
            current_node = await self._get_next_node(current_node)

get_state(session_id) async

Get the current state for a thread from the last checkpoint.

Parameters:

Name	Type	Description	Default
session_id	str	The session ID	required

Returns:

Type	Description
Optional[Dict[str, Any]]	The state if checkpoint exists, None otherwise

Source code in llmfy/flow_engine/flow_engine.py

async def get_state(self, session_id: str) -> Optional[Dict[str, Any]]:
    """
    Get the current state for a thread from the last checkpoint.

    Args:
        session_id: The session ID

    Returns:
        The state if checkpoint exists, None otherwise
    """
    if not self.checkpointer:
        raise LLMfyException("No checkpointer configured")

    checkpoint = await self.checkpointer.load(session_id)
    if checkpoint:
        # Deserialize the state to reconstruct objects
        return self._deserialize_state(checkpoint.state)
    return None

list_checkpoints(session_id, limit=10) async

List checkpoints for a specific thread.

Parameters:

Name	Type	Description	Default
session_id	str	The session ID	required
limit	int	Maximum number of checkpoints to return	10

Returns:

Type	Description
list[Checkpoint]	List of checkpoints, newest first

Source code in llmfy/flow_engine/flow_engine.py

async def list_checkpoints(
    self,
    session_id: str,
    limit: int = 10,
) -> list[Checkpoint]:
    """
    List checkpoints for a specific thread.

    Args:
        session_id: The session ID
        limit: Maximum number of checkpoints to return

    Returns:
        List of checkpoints, newest first
    """
    if not self.checkpointer:
        raise LLMfyException("No checkpointer configured")

    return await self.checkpointer.list(session_id, limit)

get_checkpoint(session_id, checkpoint_id=None) async

Get a specific checkpoint or the latest checkpoint for a thread.

Parameters:

Name	Type	Description	Default
session_id	str	The session ID	required
checkpoint_id	Optional[str]	Optional checkpoint ID, or None for latest	None

Returns:

Type	Description
Optional[Checkpoint]	The checkpoint if found, None otherwise

Source code in llmfy/flow_engine/flow_engine.py

async def get_checkpoint(
    self,
    session_id: str,
    checkpoint_id: Optional[str] = None,
) -> Optional[Checkpoint]:
    """
    Get a specific checkpoint or the latest checkpoint for a thread.

    Args:
        session_id: The session ID
        checkpoint_id: Optional checkpoint ID, or None for latest

    Returns:
        The checkpoint if found, None otherwise
    """
    if not self.checkpointer:
        raise LLMfyException("No checkpointer configured")

    return await self.checkpointer.load(session_id, checkpoint_id)

delete_checkpoints(session_id, checkpoint_id=None) async

Delete checkpoint(s) for a thread.

Parameters:

Name	Type	Description	Default
session_id	str	The session ID	required
checkpoint_id	Optional[str]	Optional checkpoint ID to delete, or None to delete all	None

Source code in llmfy/flow_engine/flow_engine.py

async def delete_checkpoints(
    self,
    session_id: str,
    checkpoint_id: Optional[str] = None,
):
    """
    Delete checkpoint(s) for a thread.

    Args:
        session_id: The session ID
        checkpoint_id: Optional checkpoint ID to delete, or None to delete all
    """
    if not self.checkpointer:
        raise LLMfyException("No checkpointer configured")

    await self.checkpointer.delete(session_id, checkpoint_id)

reset_session(session_id) async

Reset a session by deleting all its checkpoints. This allows starting fresh with the same session_id.

Parameters:

Name	Type	Description	Default
session		The session ID to reset	required

Source code in llmfy/flow_engine/flow_engine.py

async def reset_session(self, session_id: str):
    """
    Reset a session by deleting all its checkpoints.
    This allows starting fresh with the same session_id.

    Args:
        session: The session ID to reset
    """
    if not self.checkpointer:
        raise LLMfyException("No checkpointer configured")

    await self.checkpointer.delete(session_id)

details()

Generate a simple details text visualization of the workflow.

Returns:

Type	Description
str	String representation of the workflow graph

Source code in llmfy/flow_engine/flow_engine.py

def details(self) -> str:
    """
    Generate a simple details text visualization of the workflow.

    Returns:
        String representation of the workflow graph
    """
    # Check build
    if not self.is_built:
        raise LLMfyException("Build first. Use `your_flow.build()`")

    lines = ["Workflow Graph:", "=" * 50]

    # Show START connections
    start_edges = [e for e in self.edges if e.source == START]
    for edge in start_edges:
        for target in edge.targets:
            lines.append(f"START -> {target}")

    # Show all function nodes
    function_nodes = [
        n for n in self.nodes.values() if n.node_type == NodeType.FUNCTION
    ]
    if function_nodes:
        lines.append("\nFunction Nodes:")
        for node in function_nodes:
            lines.append(f"  - {node.name}")

    # Show conditional nodes
    conditional_nodes = [
        n for n in self.nodes.values() if n.node_type == NodeType.CONDITIONAL
    ]
    if conditional_nodes:
        lines.append("\nConditional Nodes:")
        for node in conditional_nodes:
            lines.append(f"  - {node.name}")

    # Show regular edges
    regular_edges = [
        e for e in self.edges if e.condition is None and e.source != START
    ]
    if regular_edges:
        lines.append("\nRegular Edges:")
        for edge in regular_edges:
            for target in edge.targets:
                lines.append(f"  {edge.source} -> {target}")

    # Show conditional edges
    conditional_edges = [e for e in self.edges if e.condition is not None]
    if conditional_edges:
        lines.append("\nConditional Edges:")
        for edge in conditional_edges:
            targets = ", ".join(edge.targets)
            lines.append(f"  {edge.source} ->? [{targets}]")

    return "\n".join(lines)

visualize()

Visualize workflow diagram. Generate Mermaid diagram url.

Returns:

Name	Type	Description
str	str	Mermaid URL.

Source code in llmfy/flow_engine/flow_engine.py

def visualize(self) -> str:
    """
    Visualize workflow diagram.
    Generate Mermaid diagram url.

    Returns:
        str: Mermaid URL.
    """
    # Check build
    if not self.is_built:
        raise LLMfyException("Build first. Use `your_flow.build()`")

    mermaid_code = self.visualizer.create_mermaid_diagram(self)
    return self.visualizer.generate_diagram_url(mermaid_code)