node_entities.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. from enum import Enum
  2. from typing import Any, Optional
  3. from pydantic import BaseModel
  4. from models.workflow import WorkflowNodeExecutionStatus
  5. class NodeType(Enum):
  6. """
  7. Node Types.
  8. """
  9. START = 'start'
  10. END = 'end'
  11. ANSWER = 'answer'
  12. LLM = 'llm'
  13. KNOWLEDGE_RETRIEVAL = 'knowledge-retrieval'
  14. IF_ELSE = 'if-else'
  15. CODE = 'code'
  16. TEMPLATE_TRANSFORM = 'template-transform'
  17. QUESTION_CLASSIFIER = 'question-classifier'
  18. HTTP_REQUEST = 'http-request'
  19. TOOL = 'tool'
  20. VARIABLE_ASSIGNER = 'variable-assigner'
  21. @classmethod
  22. def value_of(cls, value: str) -> 'NodeType':
  23. """
  24. Get value of given node type.
  25. :param value: node type value
  26. :return: node type
  27. """
  28. for node_type in cls:
  29. if node_type.value == value:
  30. return node_type
  31. raise ValueError(f'invalid node type value {value}')
  32. class SystemVariable(Enum):
  33. """
  34. System Variables.
  35. """
  36. QUERY = 'query'
  37. FILES = 'files'
  38. CONVERSATION_ID = 'conversation_id'
  39. USER_ID = 'user_id'
  40. @classmethod
  41. def value_of(cls, value: str) -> 'SystemVariable':
  42. """
  43. Get value of given system variable.
  44. :param value: system variable value
  45. :return: system variable
  46. """
  47. for system_variable in cls:
  48. if system_variable.value == value:
  49. return system_variable
  50. raise ValueError(f'invalid system variable value {value}')
  51. class NodeRunMetadataKey(Enum):
  52. """
  53. Node Run Metadata Key.
  54. """
  55. TOTAL_TOKENS = 'total_tokens'
  56. TOTAL_PRICE = 'total_price'
  57. CURRENCY = 'currency'
  58. TOOL_INFO = 'tool_info'
  59. class NodeRunResult(BaseModel):
  60. """
  61. Node Run Result.
  62. """
  63. status: WorkflowNodeExecutionStatus = WorkflowNodeExecutionStatus.RUNNING
  64. inputs: Optional[dict] = None # node inputs
  65. process_data: Optional[dict] = None # process data
  66. outputs: Optional[dict] = None # node outputs
  67. metadata: Optional[dict[NodeRunMetadataKey, Any]] = None # node metadata
  68. edge_source_handle: Optional[str] = None # source handle id of node with multiple branches
  69. error: Optional[str] = None # error message if status is failed