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 pr
| 19 | return current_executing_context.get(None) |
| 20 | |
| 21 | class CurrentNodeContext: |
| 22 | """ |
| 23 | Context manager for setting the current executing node context. |
| 24 | |
| 25 | Sets the current_executing_context on enter and resets it on exit. |
| 26 | |
| 27 | Example: |
| 28 | with CurrentNodeContext(node_id="123", list_index=0): |
| 29 | # Code that should run with the current node context set |
| 30 | process_image() |
| 31 | """ |
| 32 | def __init__(self, prompt_id: str, node_id: str, list_index: Optional[int] = None): |
| 33 | self.context = ExecutionContext( |
| 34 | prompt_id= prompt_id, |
| 35 | node_id= node_id, |
| 36 | list_index= list_index |
| 37 | ) |
| 38 | self.token = None |
| 39 | |
| 40 | def __enter__(self): |
| 41 | self.token = current_executing_context.set(self.context) |
| 42 | return self |
| 43 | |
| 44 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 45 | if self.token is not None: |
| 46 | current_executing_context.reset(self.token) |
no outgoing calls
no test coverage detected