entities.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from typing import Any, Optional
  2. from pydantic import Field
  3. from core.workflow.nodes.base import BaseIterationNodeData, BaseIterationState, BaseNodeData
  4. class IterationNodeData(BaseIterationNodeData):
  5. """
  6. Iteration Node Data.
  7. """
  8. parent_loop_id: Optional[str] = None # redundant field, not used currently
  9. iterator_selector: list[str] # variable selector
  10. output_selector: list[str] # output selector
  11. class IterationStartNodeData(BaseNodeData):
  12. """
  13. Iteration Start Node Data.
  14. """
  15. pass
  16. class IterationState(BaseIterationState):
  17. """
  18. Iteration State.
  19. """
  20. outputs: list[Any] = Field(default_factory=list)
  21. current_output: Optional[Any] = None
  22. class MetaData(BaseIterationState.MetaData):
  23. """
  24. Data.
  25. """
  26. iterator_length: int
  27. def get_last_output(self) -> Optional[Any]:
  28. """
  29. Get last output.
  30. """
  31. if self.outputs:
  32. return self.outputs[-1]
  33. return None
  34. def get_current_output(self) -> Optional[Any]:
  35. """
  36. Get current output.
  37. """
  38. return self.current_output