queue_entities.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570
  1. from collections.abc import Mapping
  2. from datetime import datetime
  3. from enum import Enum, StrEnum
  4. from typing import Any, Optional
  5. from pydantic import BaseModel
  6. from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk
  7. from core.workflow.entities.node_entities import NodeRunMetadataKey
  8. from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
  9. from core.workflow.nodes import NodeType
  10. from core.workflow.nodes.base import BaseNodeData
  11. class QueueEvent(StrEnum):
  12. """
  13. QueueEvent enum
  14. """
  15. LLM_CHUNK = "llm_chunk"
  16. TEXT_CHUNK = "text_chunk"
  17. AGENT_MESSAGE = "agent_message"
  18. MESSAGE_REPLACE = "message_replace"
  19. MESSAGE_END = "message_end"
  20. ADVANCED_CHAT_MESSAGE_END = "advanced_chat_message_end"
  21. WORKFLOW_STARTED = "workflow_started"
  22. WORKFLOW_SUCCEEDED = "workflow_succeeded"
  23. WORKFLOW_FAILED = "workflow_failed"
  24. WORKFLOW_PARTIAL_SUCCEEDED = "workflow_partial_succeeded"
  25. ITERATION_START = "iteration_start"
  26. ITERATION_NEXT = "iteration_next"
  27. ITERATION_COMPLETED = "iteration_completed"
  28. NODE_STARTED = "node_started"
  29. NODE_SUCCEEDED = "node_succeeded"
  30. NODE_FAILED = "node_failed"
  31. NODE_EXCEPTION = "node_exception"
  32. RETRIEVER_RESOURCES = "retriever_resources"
  33. ANNOTATION_REPLY = "annotation_reply"
  34. AGENT_THOUGHT = "agent_thought"
  35. MESSAGE_FILE = "message_file"
  36. PARALLEL_BRANCH_RUN_STARTED = "parallel_branch_run_started"
  37. PARALLEL_BRANCH_RUN_SUCCEEDED = "parallel_branch_run_succeeded"
  38. PARALLEL_BRANCH_RUN_FAILED = "parallel_branch_run_failed"
  39. AGENT_LOG = "agent_log"
  40. ERROR = "error"
  41. PING = "ping"
  42. STOP = "stop"
  43. class AppQueueEvent(BaseModel):
  44. """
  45. QueueEvent abstract entity
  46. """
  47. event: QueueEvent
  48. class QueueLLMChunkEvent(AppQueueEvent):
  49. """
  50. QueueLLMChunkEvent entity
  51. Only for basic mode apps
  52. """
  53. event: QueueEvent = QueueEvent.LLM_CHUNK
  54. chunk: LLMResultChunk
  55. class QueueIterationStartEvent(AppQueueEvent):
  56. """
  57. QueueIterationStartEvent entity
  58. """
  59. event: QueueEvent = QueueEvent.ITERATION_START
  60. node_execution_id: str
  61. node_id: str
  62. node_type: NodeType
  63. node_data: BaseNodeData
  64. parallel_id: Optional[str] = None
  65. """parallel id if node is in parallel"""
  66. parallel_start_node_id: Optional[str] = None
  67. """parallel start node id if node is in parallel"""
  68. parent_parallel_id: Optional[str] = None
  69. """parent parallel id if node is in parallel"""
  70. parent_parallel_start_node_id: Optional[str] = None
  71. """parent parallel start node id if node is in parallel"""
  72. start_at: datetime
  73. node_run_index: int
  74. inputs: Optional[dict[str, Any]] = None
  75. predecessor_node_id: Optional[str] = None
  76. metadata: Optional[dict[str, Any]] = None
  77. class QueueIterationNextEvent(AppQueueEvent):
  78. """
  79. QueueIterationNextEvent entity
  80. """
  81. event: QueueEvent = QueueEvent.ITERATION_NEXT
  82. index: int
  83. node_execution_id: str
  84. node_id: str
  85. node_type: NodeType
  86. node_data: BaseNodeData
  87. parallel_id: Optional[str] = None
  88. """parallel id if node is in parallel"""
  89. parallel_start_node_id: Optional[str] = None
  90. """parallel start node id if node is in parallel"""
  91. parent_parallel_id: Optional[str] = None
  92. """parent parallel id if node is in parallel"""
  93. parent_parallel_start_node_id: Optional[str] = None
  94. """parent parallel start node id if node is in parallel"""
  95. parallel_mode_run_id: Optional[str] = None
  96. """iteratoin run in parallel mode run id"""
  97. node_run_index: int
  98. output: Optional[Any] = None # output for the current iteration
  99. duration: Optional[float] = None
  100. class QueueIterationCompletedEvent(AppQueueEvent):
  101. """
  102. QueueIterationCompletedEvent entity
  103. """
  104. event: QueueEvent = QueueEvent.ITERATION_COMPLETED
  105. node_execution_id: str
  106. node_id: str
  107. node_type: NodeType
  108. node_data: BaseNodeData
  109. parallel_id: Optional[str] = None
  110. """parallel id if node is in parallel"""
  111. parallel_start_node_id: Optional[str] = None
  112. """parallel start node id if node is in parallel"""
  113. parent_parallel_id: Optional[str] = None
  114. """parent parallel id if node is in parallel"""
  115. parent_parallel_start_node_id: Optional[str] = None
  116. """parent parallel start node id if node is in parallel"""
  117. start_at: datetime
  118. node_run_index: int
  119. inputs: Optional[dict[str, Any]] = None
  120. outputs: Optional[dict[str, Any]] = None
  121. metadata: Optional[dict[str, Any]] = None
  122. steps: int = 0
  123. error: Optional[str] = None
  124. class QueueTextChunkEvent(AppQueueEvent):
  125. """
  126. QueueTextChunkEvent entity
  127. """
  128. event: QueueEvent = QueueEvent.TEXT_CHUNK
  129. text: str
  130. from_variable_selector: Optional[list[str]] = None
  131. """from variable selector"""
  132. in_iteration_id: Optional[str] = None
  133. """iteration id if node is in iteration"""
  134. class QueueAgentMessageEvent(AppQueueEvent):
  135. """
  136. QueueMessageEvent entity
  137. """
  138. event: QueueEvent = QueueEvent.AGENT_MESSAGE
  139. chunk: LLMResultChunk
  140. class QueueMessageReplaceEvent(AppQueueEvent):
  141. """
  142. QueueMessageReplaceEvent entity
  143. """
  144. event: QueueEvent = QueueEvent.MESSAGE_REPLACE
  145. text: str
  146. class QueueRetrieverResourcesEvent(AppQueueEvent):
  147. """
  148. QueueRetrieverResourcesEvent entity
  149. """
  150. event: QueueEvent = QueueEvent.RETRIEVER_RESOURCES
  151. retriever_resources: list[dict]
  152. in_iteration_id: Optional[str] = None
  153. """iteration id if node is in iteration"""
  154. class QueueAnnotationReplyEvent(AppQueueEvent):
  155. """
  156. QueueAnnotationReplyEvent entity
  157. """
  158. event: QueueEvent = QueueEvent.ANNOTATION_REPLY
  159. message_annotation_id: str
  160. class QueueMessageEndEvent(AppQueueEvent):
  161. """
  162. QueueMessageEndEvent entity
  163. """
  164. event: QueueEvent = QueueEvent.MESSAGE_END
  165. llm_result: Optional[LLMResult] = None
  166. class QueueAdvancedChatMessageEndEvent(AppQueueEvent):
  167. """
  168. QueueAdvancedChatMessageEndEvent entity
  169. """
  170. event: QueueEvent = QueueEvent.ADVANCED_CHAT_MESSAGE_END
  171. class QueueWorkflowStartedEvent(AppQueueEvent):
  172. """
  173. QueueWorkflowStartedEvent entity
  174. """
  175. event: QueueEvent = QueueEvent.WORKFLOW_STARTED
  176. graph_runtime_state: GraphRuntimeState
  177. class QueueWorkflowSucceededEvent(AppQueueEvent):
  178. """
  179. QueueWorkflowSucceededEvent entity
  180. """
  181. event: QueueEvent = QueueEvent.WORKFLOW_SUCCEEDED
  182. outputs: Optional[dict[str, Any]] = None
  183. class QueueWorkflowFailedEvent(AppQueueEvent):
  184. """
  185. QueueWorkflowFailedEvent entity
  186. """
  187. event: QueueEvent = QueueEvent.WORKFLOW_FAILED
  188. error: str
  189. exceptions_count: int
  190. class QueueWorkflowPartialSuccessEvent(AppQueueEvent):
  191. """
  192. QueueWorkflowFailedEvent entity
  193. """
  194. event: QueueEvent = QueueEvent.WORKFLOW_PARTIAL_SUCCEEDED
  195. exceptions_count: int
  196. outputs: Optional[dict[str, Any]] = None
  197. class QueueNodeStartedEvent(AppQueueEvent):
  198. """
  199. QueueNodeStartedEvent entity
  200. """
  201. event: QueueEvent = QueueEvent.NODE_STARTED
  202. node_execution_id: str
  203. node_id: str
  204. node_type: NodeType
  205. node_data: BaseNodeData
  206. node_run_index: int = 1
  207. predecessor_node_id: Optional[str] = None
  208. parallel_id: Optional[str] = None
  209. """parallel id if node is in parallel"""
  210. parallel_start_node_id: Optional[str] = None
  211. """parallel start node id if node is in parallel"""
  212. parent_parallel_id: Optional[str] = None
  213. """parent parallel id if node is in parallel"""
  214. parent_parallel_start_node_id: Optional[str] = None
  215. """parent parallel start node id if node is in parallel"""
  216. in_iteration_id: Optional[str] = None
  217. """iteration id if node is in iteration"""
  218. start_at: datetime
  219. parallel_mode_run_id: Optional[str] = None
  220. """iteratoin run in parallel mode run id"""
  221. class QueueNodeSucceededEvent(AppQueueEvent):
  222. """
  223. QueueNodeSucceededEvent entity
  224. """
  225. event: QueueEvent = QueueEvent.NODE_SUCCEEDED
  226. node_execution_id: str
  227. node_id: str
  228. node_type: NodeType
  229. node_data: BaseNodeData
  230. parallel_id: Optional[str] = None
  231. """parallel id if node is in parallel"""
  232. parallel_start_node_id: Optional[str] = None
  233. """parallel start node id if node is in parallel"""
  234. parent_parallel_id: Optional[str] = None
  235. """parent parallel id if node is in parallel"""
  236. parent_parallel_start_node_id: Optional[str] = None
  237. """parent parallel start node id if node is in parallel"""
  238. in_iteration_id: Optional[str] = None
  239. """iteration id if node is in iteration"""
  240. start_at: datetime
  241. inputs: Optional[dict[str, Any]] = None
  242. process_data: Optional[dict[str, Any]] = None
  243. outputs: Optional[dict[str, Any]] = None
  244. execution_metadata: Optional[dict[NodeRunMetadataKey, Any]] = None
  245. error: Optional[str] = None
  246. """single iteration duration map"""
  247. iteration_duration_map: Optional[dict[str, float]] = None
  248. class QueueAgentLogEvent(AppQueueEvent):
  249. """
  250. QueueAgentLogEvent entity
  251. """
  252. event: QueueEvent = QueueEvent.AGENT_LOG
  253. id: str
  254. label: str
  255. node_execution_id: str
  256. parent_id: str | None
  257. error: str | None
  258. status: str
  259. data: Mapping[str, Any]
  260. class QueueNodeInIterationFailedEvent(AppQueueEvent):
  261. """
  262. QueueNodeInIterationFailedEvent entity
  263. """
  264. event: QueueEvent = QueueEvent.NODE_FAILED
  265. node_execution_id: str
  266. node_id: str
  267. node_type: NodeType
  268. node_data: BaseNodeData
  269. parallel_id: Optional[str] = None
  270. """parallel id if node is in parallel"""
  271. parallel_start_node_id: Optional[str] = None
  272. """parallel start node id if node is in parallel"""
  273. parent_parallel_id: Optional[str] = None
  274. """parent parallel id if node is in parallel"""
  275. parent_parallel_start_node_id: Optional[str] = None
  276. """parent parallel start node id if node is in parallel"""
  277. in_iteration_id: Optional[str] = None
  278. """iteration id if node is in iteration"""
  279. start_at: datetime
  280. inputs: Optional[dict[str, Any]] = None
  281. process_data: Optional[dict[str, Any]] = None
  282. outputs: Optional[dict[str, Any]] = None
  283. execution_metadata: Optional[dict[NodeRunMetadataKey, Any]] = None
  284. error: str
  285. class QueueNodeExceptionEvent(AppQueueEvent):
  286. """
  287. QueueNodeExceptionEvent entity
  288. """
  289. event: QueueEvent = QueueEvent.NODE_EXCEPTION
  290. node_execution_id: str
  291. node_id: str
  292. node_type: NodeType
  293. node_data: BaseNodeData
  294. parallel_id: Optional[str] = None
  295. """parallel id if node is in parallel"""
  296. parallel_start_node_id: Optional[str] = None
  297. """parallel start node id if node is in parallel"""
  298. parent_parallel_id: Optional[str] = None
  299. """parent parallel id if node is in parallel"""
  300. parent_parallel_start_node_id: Optional[str] = None
  301. """parent parallel start node id if node is in parallel"""
  302. in_iteration_id: Optional[str] = None
  303. """iteration id if node is in iteration"""
  304. start_at: datetime
  305. inputs: Optional[dict[str, Any]] = None
  306. process_data: Optional[dict[str, Any]] = None
  307. outputs: Optional[dict[str, Any]] = None
  308. execution_metadata: Optional[dict[NodeRunMetadataKey, Any]] = None
  309. error: str
  310. class QueueNodeFailedEvent(AppQueueEvent):
  311. """
  312. QueueNodeFailedEvent entity
  313. """
  314. event: QueueEvent = QueueEvent.NODE_FAILED
  315. node_execution_id: str
  316. node_id: str
  317. node_type: NodeType
  318. node_data: BaseNodeData
  319. parallel_id: Optional[str] = None
  320. """parallel id if node is in parallel"""
  321. parallel_start_node_id: Optional[str] = None
  322. """parallel start node id if node is in parallel"""
  323. parent_parallel_id: Optional[str] = None
  324. """parent parallel id if node is in parallel"""
  325. parent_parallel_start_node_id: Optional[str] = None
  326. """parent parallel start node id if node is in parallel"""
  327. in_iteration_id: Optional[str] = None
  328. """iteration id if node is in iteration"""
  329. start_at: datetime
  330. inputs: Optional[dict[str, Any]] = None
  331. process_data: Optional[dict[str, Any]] = None
  332. outputs: Optional[dict[str, Any]] = None
  333. execution_metadata: Optional[dict[NodeRunMetadataKey, Any]] = None
  334. error: str
  335. class QueueAgentThoughtEvent(AppQueueEvent):
  336. """
  337. QueueAgentThoughtEvent entity
  338. """
  339. event: QueueEvent = QueueEvent.AGENT_THOUGHT
  340. agent_thought_id: str
  341. class QueueMessageFileEvent(AppQueueEvent):
  342. """
  343. QueueAgentThoughtEvent entity
  344. """
  345. event: QueueEvent = QueueEvent.MESSAGE_FILE
  346. message_file_id: str
  347. class QueueErrorEvent(AppQueueEvent):
  348. """
  349. QueueErrorEvent entity
  350. """
  351. event: QueueEvent = QueueEvent.ERROR
  352. error: Any = None
  353. class QueuePingEvent(AppQueueEvent):
  354. """
  355. QueuePingEvent entity
  356. """
  357. event: QueueEvent = QueueEvent.PING
  358. class QueueStopEvent(AppQueueEvent):
  359. """
  360. QueueStopEvent entity
  361. """
  362. class StopBy(Enum):
  363. """
  364. Stop by enum
  365. """
  366. USER_MANUAL = "user-manual"
  367. ANNOTATION_REPLY = "annotation-reply"
  368. OUTPUT_MODERATION = "output-moderation"
  369. INPUT_MODERATION = "input-moderation"
  370. event: QueueEvent = QueueEvent.STOP
  371. stopped_by: StopBy
  372. def get_stop_reason(self) -> str:
  373. """
  374. To stop reason
  375. """
  376. reason_mapping = {
  377. QueueStopEvent.StopBy.USER_MANUAL: "Stopped by user.",
  378. QueueStopEvent.StopBy.ANNOTATION_REPLY: "Stopped by annotation reply.",
  379. QueueStopEvent.StopBy.OUTPUT_MODERATION: "Stopped by output moderation.",
  380. QueueStopEvent.StopBy.INPUT_MODERATION: "Stopped by input moderation.",
  381. }
  382. return reason_mapping.get(self.stopped_by, "Stopped by unknown reason.")
  383. class QueueMessage(BaseModel):
  384. """
  385. QueueMessage abstract entity
  386. """
  387. task_id: str
  388. app_mode: str
  389. event: AppQueueEvent
  390. class MessageQueueMessage(QueueMessage):
  391. """
  392. MessageQueueMessage entity
  393. """
  394. message_id: str
  395. conversation_id: str
  396. class WorkflowQueueMessage(QueueMessage):
  397. """
  398. WorkflowQueueMessage entity
  399. """
  400. pass
  401. class QueueParallelBranchRunStartedEvent(AppQueueEvent):
  402. """
  403. QueueParallelBranchRunStartedEvent entity
  404. """
  405. event: QueueEvent = QueueEvent.PARALLEL_BRANCH_RUN_STARTED
  406. parallel_id: str
  407. parallel_start_node_id: str
  408. parent_parallel_id: Optional[str] = None
  409. """parent parallel id if node is in parallel"""
  410. parent_parallel_start_node_id: Optional[str] = None
  411. """parent parallel start node id if node is in parallel"""
  412. in_iteration_id: Optional[str] = None
  413. """iteration id if node is in iteration"""
  414. class QueueParallelBranchRunSucceededEvent(AppQueueEvent):
  415. """
  416. QueueParallelBranchRunSucceededEvent entity
  417. """
  418. event: QueueEvent = QueueEvent.PARALLEL_BRANCH_RUN_SUCCEEDED
  419. parallel_id: str
  420. parallel_start_node_id: str
  421. parent_parallel_id: Optional[str] = None
  422. """parent parallel id if node is in parallel"""
  423. parent_parallel_start_node_id: Optional[str] = None
  424. """parent parallel start node id if node is in parallel"""
  425. in_iteration_id: Optional[str] = None
  426. """iteration id if node is in iteration"""
  427. class QueueParallelBranchRunFailedEvent(AppQueueEvent):
  428. """
  429. QueueParallelBranchRunFailedEvent entity
  430. """
  431. event: QueueEvent = QueueEvent.PARALLEL_BRANCH_RUN_FAILED
  432. parallel_id: str
  433. parallel_start_node_id: str
  434. parent_parallel_id: Optional[str] = None
  435. """parent parallel id if node is in parallel"""
  436. parent_parallel_start_node_id: Optional[str] = None
  437. """parent parallel start node id if node is in parallel"""
  438. in_iteration_id: Optional[str] = None
  439. """iteration id if node is in iteration"""
  440. error: str