| 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) |
|
|
