utils.py

import contextvars
from typing import Optional, NamedTuple

class ExecutionContext(NamedTuple):
    """
    Context information about the currently executing node.

    Attributes:
        node_id: The ID of the currently executing node
        list_index: The index in a list being processed (for operations on batches/lists)
    """
    prompt_id: str
    node_id: str
    list_index: Optional[int]

current_executing_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar("current_executing_context", default=None)

def get_executing_context() -> Optional[ExecutionContext]:
    return current_executing_context.get(None)

class CurrentNodeContext:
    """
    Context manager for setting the current executing node context.

    Sets the current_executing_context on enter and resets it on exit.

    Example:
        with CurrentNodeContext(node_id="123", list_index=0):
            # Code that should run with the current node context set
            process_image()
    """
    def __init__(self, prompt_id: str, node_id: str, list_index: Optional[int] = None):
        self.context = ExecutionContext(
            prompt_id= prompt_id,
            node_id= node_id,
            list_index= list_index
        )
        self.token = None

    def __enter__(self):
        self.token = current_executing_context.set(self.context)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.token is not None:
            current_executing_context.reset(self.token)