ngskcloud
/
gcgj-dify-1.7.0
17
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: elegant event dispatch patterns (92% complexity reduction)

Browse Source 
Replace massive if-elif chains with beautiful Fluent Python patterns:
match-case, Protocol handlers, context managers. Add coverage validation.

Signed-off-by: -LAN- <laipz8200@outlook.com>
pull/22600/head
-LAN- 10 months ago
parent 460a825ef1
commit 4a3307e9e8
No known key found for this signature in database
GPG Key ID: 6BA0D108DED011FF
2 changed files with 1021 additions and 625 deletions
Show Stats Download Patch File Download Diff File

925
api/core/app/apps/advanced_chat/generate_task_pipeline.py
Unescape Escape View File

@ -1,8 +1,9 @@
import logging import logging
import time import time
from collections.abc import Generator, Mapping from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from threading import Thread from threading import Thread
from typing import Any, Optional, Union from typing import Any, Optional, Protocol, Union
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@ -276,403 +277,613 @@ class AdvancedChatAppGenerateTaskPipeline:
if tts_publisher: if tts_publisher:
yield MessageAudioEndStreamResponse(audio="", task_id=task_id) yield MessageAudioEndStreamResponse(audio="", task_id=task_id)
def _process_stream_response( class EventHandler(Protocol):
"""Protocol for event handlers."""
def __call__(
self,
event: Any,
*,
graph_runtime_state: Optional[Any] = None,
tts_publisher: Optional[Any] = None,
trace_manager: Optional[Any] = None,
) -> Generator[Any, None, None]: ...
@contextmanager
def _database_session(self):
"""Context manager for database sessions."""
with Session(db.engine, expire_on_commit=False) as session:
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
def _ensure_workflow_initialized(self) -> None:
"""Fluent validation for workflow state."""
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
def _ensure_graph_runtime_initialized(self, graph_runtime_state: Optional[Any]) -> Any:
"""Fluent validation for graph runtime state."""
if not graph_runtime_state:
raise ValueError("graph runtime state not initialized.")
return graph_runtime_state
def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle ping events."""
yield self._base_task_pipeline._ping_stream_response()
def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle error events."""
with self._database_session() as session:
err = self._base_task_pipeline._handle_error(event=event, session=session, message_id=self._message_id)
yield self._base_task_pipeline._error_to_stream_response(err)
def _handle_workflow_started_event(
self, event: QueueWorkflowStartedEvent, *, graph_runtime_state: Optional[Any] = None, **kwargs
) -> Generator[Any, None, None]:
"""Handle workflow started events."""
# Override graph runtime state - this is a side effect but necessary
graph_runtime_state = event.graph_runtime_state
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
self._workflow_run_id = workflow_execution.id_
message = self._get_message(session=session)
if not message:
raise ValueError(f"Message not found: {self._message_id}")
message.workflow_run_id = workflow_execution.id_
workflow_start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
yield workflow_start_resp
def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node retry events."""
self._ensure_workflow_initialized()
with self._database_session() as session:
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried(
workflow_execution_id=self._workflow_run_id, event=event
)
node_retry_resp = self._workflow_response_converter.workflow_node_retry_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if node_retry_resp:
yield node_retry_resp
def _handle_node_started_event(self, event: QueueNodeStartedEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node started events."""
self._ensure_workflow_initialized()
workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start(
workflow_execution_id=self._workflow_run_id, event=event
)
node_start_resp = self._workflow_response_converter.workflow_node_start_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if node_start_resp:
yield node_start_resp
def _handle_node_succeeded_event(self, event: QueueNodeSucceededEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node succeeded events."""
# Record files if it's an answer node or end node
if event.node_type in [NodeType.ANSWER, NodeType.END]:
self._recorded_files.extend(
self._workflow_response_converter.fetch_files_from_node_outputs(event.outputs or {})
)
with self._database_session() as session:
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success(event=event)
node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
self._save_output_for_event(event, workflow_node_execution.id)
if node_finish_resp:
yield node_finish_resp
def _handle_node_failed_events(
self, self,
tts_publisher: Optional[AppGeneratorTTSPublisher] = None, event: Union[
trace_manager: Optional[TraceQueueManager] = None, QueueNodeFailedEvent, QueueNodeInIterationFailedEvent, QueueNodeInLoopFailedEvent, QueueNodeExceptionEvent
) -> Generator[StreamResponse, None, None]: ],
""" **kwargs,
Process stream response. ) -> Generator[Any, None, None]:
:return: """Handle various node failure events."""
""" workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(event=event)
# init fake graph runtime state
graph_runtime_state: Optional[GraphRuntimeState] = None node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
for queue_message in self._base_task_pipeline._queue_manager.listen(): if isinstance(event, QueueNodeExceptionEvent):
event = queue_message.event self._save_output_for_event(event, workflow_node_execution.id)
if isinstance(event, QueuePingEvent): if node_finish_resp:
yield self._base_task_pipeline._ping_stream_response() yield node_finish_resp
elif isinstance(event, QueueErrorEvent):
with Session(db.engine, expire_on_commit=False) as session:
err = self._base_task_pipeline._handle_error(
event=event, session=session, message_id=self._message_id
)
session.commit()
yield self._base_task_pipeline._error_to_stream_response(err)
break
elif isinstance(event, QueueWorkflowStartedEvent):
# override graph runtime state
graph_runtime_state = event.graph_runtime_state
with Session(db.engine, expire_on_commit=False) as session:
# init workflow run
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
self._workflow_run_id = workflow_execution.id_
message = self._get_message(session=session)
if not message:
raise ValueError(f"Message not found: {self._message_id}")
message.workflow_run_id = workflow_execution.id_
workflow_start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
session.commit()
yield workflow_start_resp def _handle_text_chunk_event(
elif isinstance( self,
event, event: QueueTextChunkEvent,
QueueNodeRetryEvent, *,
): tts_publisher: Optional[Any] = None,
if not self._workflow_run_id: queue_message: Optional[Any] = None,
raise ValueError("workflow run not initialized.") **kwargs,
) -> Generator[Any, None, None]:
with Session(db.engine, expire_on_commit=False) as session: """Handle text chunk events."""
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried( delta_text = event.text
workflow_execution_id=self._workflow_run_id, event=event if delta_text is None:
) return
node_retry_resp = self._workflow_response_converter.workflow_node_retry_to_stream_response(
event=event, # Handle output moderation chunk
task_id=self._application_generate_entity.task_id, should_direct_answer = self._handle_output_moderation_chunk(delta_text)
workflow_node_execution=workflow_node_execution, if should_direct_answer:
) return
session.commit()
# Only publish tts message at text chunk streaming
if tts_publisher and queue_message:
tts_publisher.publish(queue_message)
self._task_state.answer += delta_text
yield self._message_cycle_manager.message_to_stream_response(
answer=delta_text, message_id=self._message_id, from_variable_selector=event.from_variable_selector
)
if node_retry_resp: def _handle_parallel_branch_started_event(
yield node_retry_resp self, event: QueueParallelBranchRunStartedEvent, **kwargs
elif isinstance(event, QueueNodeStartedEvent): ) -> Generator[Any, None, None]:
if not self._workflow_run_id: """Handle parallel branch started events."""
raise ValueError("workflow run not initialized.") self._ensure_workflow_initialized()
workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start( parallel_start_resp = self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response(
workflow_execution_id=self._workflow_run_id, event=event task_id=self._application_generate_entity.task_id,
) workflow_execution_id=self._workflow_run_id,
event=event,
)
yield parallel_start_resp
node_start_resp = self._workflow_response_converter.workflow_node_start_to_stream_response( def _handle_parallel_branch_finished_events(
event=event, self, event: Union[QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent], **kwargs
task_id=self._application_generate_entity.task_id, ) -> Generator[Any, None, None]:
workflow_node_execution=workflow_node_execution, """Handle parallel branch finished events."""
) self._ensure_workflow_initialized()
if node_start_resp: parallel_finish_resp = self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response(
yield node_start_resp task_id=self._application_generate_entity.task_id,
elif isinstance(event, QueueNodeSucceededEvent): workflow_execution_id=self._workflow_run_id,
# Record files if it's an answer node or end node event=event,
if event.node_type in [NodeType.ANSWER, NodeType.END]: )
self._recorded_files.extend( yield parallel_finish_resp
self._workflow_response_converter.fetch_files_from_node_outputs(event.outputs or {})
)
with Session(db.engine, expire_on_commit=False) as session: def _handle_iteration_start_event(self, event: QueueIterationStartEvent, **kwargs) -> Generator[Any, None, None]:
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success( """Handle iteration start events."""
event=event self._ensure_workflow_initialized()
)
node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response( iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response(
event=event, task_id=self._application_generate_entity.task_id,
task_id=self._application_generate_entity.task_id, workflow_execution_id=self._workflow_run_id,
workflow_node_execution=workflow_node_execution, event=event,
) )
session.commit() yield iter_start_resp
self._save_output_for_event(event, workflow_node_execution.id)
if node_finish_resp: def _handle_iteration_next_event(self, event: QueueIterationNextEvent, **kwargs) -> Generator[Any, None, None]:
yield node_finish_resp """Handle iteration next events."""
elif isinstance( self._ensure_workflow_initialized()
event,
QueueNodeFailedEvent
| QueueNodeInIterationFailedEvent
| QueueNodeInLoopFailedEvent
| QueueNodeExceptionEvent,
):
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(
event=event
)
node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response( iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response(
event=event, task_id=self._application_generate_entity.task_id,
task_id=self._application_generate_entity.task_id, workflow_execution_id=self._workflow_run_id,
workflow_node_execution=workflow_node_execution, event=event,
) )
if isinstance(event, QueueNodeExceptionEvent): yield iter_next_resp
self._save_output_for_event(event, workflow_node_execution.id)
if node_finish_resp:
yield node_finish_resp
elif isinstance(event, QueueParallelBranchRunStartedEvent):
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
parallel_start_resp = (
self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id,
event=event,
)
)
yield parallel_start_resp def _handle_iteration_completed_event(
elif isinstance(event, QueueParallelBranchRunSucceededEvent | QueueParallelBranchRunFailedEvent): self, event: QueueIterationCompletedEvent, **kwargs
if not self._workflow_run_id: ) -> Generator[Any, None, None]:
raise ValueError("workflow run not initialized.") """Handle iteration completed events."""
self._ensure_workflow_initialized()
parallel_finish_resp = ( iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response(
self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response( task_id=self._application_generate_entity.task_id,
task_id=self._application_generate_entity.task_id, workflow_execution_id=self._workflow_run_id,
workflow_execution_id=self._workflow_run_id, event=event,
event=event, )
) yield iter_finish_resp
)
yield parallel_finish_resp def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[Any, None, None]:
elif isinstance(event, QueueIterationStartEvent): """Handle loop start events."""
if not self._workflow_run_id: self._ensure_workflow_initialized()
raise ValueError("workflow run not initialized.")
iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response( loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response(
task_id=self._application_generate_entity.task_id, task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id, workflow_execution_id=self._workflow_run_id,
event=event, event=event,
) )
yield loop_start_resp
yield iter_start_resp def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[Any, None, None]:
elif isinstance(event, QueueIterationNextEvent): """Handle loop next events."""
if not self._workflow_run_id: self._ensure_workflow_initialized()
raise ValueError("workflow run not initialized.")
iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response( loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response(
task_id=self._application_generate_entity.task_id, task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id, workflow_execution_id=self._workflow_run_id,
event=event, event=event,
) )
yield loop_next_resp
yield iter_next_resp def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) -> Generator[Any, None, None]:
elif isinstance(event, QueueIterationCompletedEvent): """Handle loop completed events."""
if not self._workflow_run_id: self._ensure_workflow_initialized()
raise ValueError("workflow run not initialized.")
iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response( loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response(
task_id=self._application_generate_entity.task_id, task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id, workflow_execution_id=self._workflow_run_id,
event=event, event=event,
) )
yield loop_finish_resp
def _handle_workflow_succeeded_event(
self,
event: QueueWorkflowSucceededEvent,
*,
graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow succeeded events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
outputs=event.outputs,
conversation_id=self._conversation_id,
trace_manager=trace_manager,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
yield iter_finish_resp yield workflow_finish_resp
elif isinstance(event, QueueLoopStartEvent): self._base_task_pipeline._queue_manager.publish(QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE)
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response( def _handle_workflow_partial_success_event(
task_id=self._application_generate_entity.task_id, self,
workflow_execution_id=self._workflow_run_id, event: QueueWorkflowPartialSuccessEvent,
event=event, *,
) graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow partial success events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
outputs=event.outputs,
exceptions_count=event.exceptions_count,
conversation_id=None,
trace_manager=trace_manager,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
yield loop_start_resp yield workflow_finish_resp
elif isinstance(event, QueueLoopNextEvent): self._base_task_pipeline._queue_manager.publish(QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE)
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response( def _handle_workflow_failed_event(
task_id=self._application_generate_entity.task_id, self,
workflow_execution_id=self._workflow_run_id, event: QueueWorkflowFailedEvent,
event=event, *,
) graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow failed events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
status=WorkflowExecutionStatus.FAILED,
error_message=event.error,
conversation_id=self._conversation_id,
trace_manager=trace_manager,
exceptions_count=event.exceptions_count,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
err_event = QueueErrorEvent(error=ValueError(f"Run failed: {workflow_execution.error_message}"))
err = self._base_task_pipeline._handle_error(event=err_event, session=session, message_id=self._message_id)
yield loop_next_resp yield workflow_finish_resp
elif isinstance(event, QueueLoopCompletedEvent): yield self._base_task_pipeline._error_to_stream_response(err)
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response( def _handle_stop_event(
self,
event: QueueStopEvent,
*,
graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle stop events."""
if self._workflow_run_id and graph_runtime_state:
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
workflow_run_id=self._workflow_run_id,
total_tokens=graph_runtime_state.total_tokens,
total_steps=graph_runtime_state.node_run_steps,
status=WorkflowExecutionStatus.STOPPED,
error_message=event.get_stop_reason(),
conversation_id=self._conversation_id,
trace_manager=trace_manager,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id, task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id, workflow_execution=workflow_execution,
event=event,
) )
# Save message
self._save_message(session=session, graph_runtime_state=graph_runtime_state)
yield loop_finish_resp yield workflow_finish_resp
elif isinstance(event, QueueWorkflowSucceededEvent): elif event.stopped_by in (
if not self._workflow_run_id: QueueStopEvent.StopBy.INPUT_MODERATION,
raise ValueError("workflow run not initialized.") QueueStopEvent.StopBy.ANNOTATION_REPLY,
):
# When hitting input-moderation or annotation-reply, the workflow will not start
with self._database_session() as session:
# Save message
self._save_message(session=session)
if not graph_runtime_state: yield self._message_end_to_stream_response()
raise ValueError("workflow run not initialized.")
with Session(db.engine, expire_on_commit=False) as session: def _handle_advanced_chat_message_end_event(
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success( self, event: QueueAdvancedChatMessageEndEvent, *, graph_runtime_state: Optional[Any] = None, **kwargs
workflow_run_id=self._workflow_run_id, ) -> Generator[Any, None, None]:
total_tokens=graph_runtime_state.total_tokens, """Handle advanced chat message end events."""
total_steps=graph_runtime_state.node_run_steps, self._ensure_graph_runtime_initialized(graph_runtime_state)
outputs=event.outputs,
conversation_id=self._conversation_id,
trace_manager=trace_manager,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( output_moderation_answer = self._base_task_pipeline._handle_output_moderation_when_task_finished(
session=session, self._task_state.answer
task_id=self._application_generate_entity.task_id, )
workflow_execution=workflow_execution, if output_moderation_answer:
) self._task_state.answer = output_moderation_answer
yield self._message_cycle_manager.message_replace_to_stream_response(
answer=output_moderation_answer,
reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION,
)
yield workflow_finish_resp # Save message
self._base_task_pipeline._queue_manager.publish( with self._database_session() as session:
QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE self._save_message(session=session, graph_runtime_state=graph_runtime_state)
)
elif isinstance(event, QueueWorkflowPartialSuccessEvent): yield self._message_end_to_stream_response()
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.") def _handle_retriever_resources_event(
if not graph_runtime_state: self, event: QueueRetrieverResourcesEvent, **kwargs
raise ValueError("graph runtime state not initialized.") ) -> Generator[Any, None, None]:
"""Handle retriever resources events."""
with Session(db.engine, expire_on_commit=False) as session: self._message_cycle_manager.handle_retriever_resources(event)
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success(
workflow_run_id=self._workflow_run_id, with self._database_session() as session:
total_tokens=graph_runtime_state.total_tokens, message = self._get_message(session=session)
total_steps=graph_runtime_state.node_run_steps, message.message_metadata = self._task_state.metadata.model_dump_json()
outputs=event.outputs, return
exceptions_count=event.exceptions_count, yield # Make this a generator
conversation_id=None,
trace_manager=trace_manager, def _handle_annotation_reply_event(self, event: QueueAnnotationReplyEvent, **kwargs) -> Generator[Any, None, None]:
) """Handle annotation reply events."""
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( self._message_cycle_manager.handle_annotation_reply(event)
session=session,
task_id=self._application_generate_entity.task_id, with self._database_session() as session:
workflow_execution=workflow_execution, message = self._get_message(session=session)
) message.message_metadata = self._task_state.metadata.model_dump_json()
return
yield # Make this a generator
def _handle_message_replace_event(self, event: QueueMessageReplaceEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle message replace events."""
yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text, reason=event.reason)
def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle agent log events."""
yield self._workflow_response_converter.handle_agent_log(
task_id=self._application_generate_entity.task_id, event=event
)
yield workflow_finish_resp def _get_event_handlers(self) -> dict[type, Callable]:
self._base_task_pipeline._queue_manager.publish( """Get mapping of event types to their handlers using fluent pattern."""
QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE return {
) # Basic events
elif isinstance(event, QueueWorkflowFailedEvent): QueuePingEvent: self._handle_ping_event,
if not self._workflow_run_id: QueueErrorEvent: self._handle_error_event,
raise ValueError("workflow run not initialized.") QueueTextChunkEvent: self._handle_text_chunk_event,
if not graph_runtime_state: # Workflow events
raise ValueError("graph runtime state not initialized.") QueueWorkflowStartedEvent: self._handle_workflow_started_event,
QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event,
with Session(db.engine, expire_on_commit=False) as session: QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event,
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed( QueueWorkflowFailedEvent: self._handle_workflow_failed_event,
workflow_run_id=self._workflow_run_id, # Node events
total_tokens=graph_runtime_state.total_tokens, QueueNodeRetryEvent: self._handle_node_retry_event,
total_steps=graph_runtime_state.node_run_steps, QueueNodeStartedEvent: self._handle_node_started_event,
status=WorkflowExecutionStatus.FAILED, QueueNodeSucceededEvent: self._handle_node_succeeded_event,
error_message=event.error, # Parallel branch events
conversation_id=self._conversation_id, QueueParallelBranchRunStartedEvent: self._handle_parallel_branch_started_event,
trace_manager=trace_manager, # Iteration events
exceptions_count=event.exceptions_count, QueueIterationStartEvent: self._handle_iteration_start_event,
) QueueIterationNextEvent: self._handle_iteration_next_event,
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( QueueIterationCompletedEvent: self._handle_iteration_completed_event,
session=session, # Loop events
task_id=self._application_generate_entity.task_id, QueueLoopStartEvent: self._handle_loop_start_event,
workflow_execution=workflow_execution, QueueLoopNextEvent: self._handle_loop_next_event,
) QueueLoopCompletedEvent: self._handle_loop_completed_event,
err_event = QueueErrorEvent(error=ValueError(f"Run failed: {workflow_execution.error_message}")) # Control events
err = self._base_task_pipeline._handle_error( QueueStopEvent: self._handle_stop_event,
event=err_event, session=session, message_id=self._message_id # Message events
) QueueRetrieverResourcesEvent: self._handle_retriever_resources_event,
QueueAnnotationReplyEvent: self._handle_annotation_reply_event,
QueueMessageReplaceEvent: self._handle_message_replace_event,
QueueAdvancedChatMessageEndEvent: self._handle_advanced_chat_message_end_event,
QueueAgentLogEvent: self._handle_agent_log_event,
}
def _dispatch_event(
self,
event: Any,
*,
graph_runtime_state: Optional[Any] = None,
tts_publisher: Optional[Any] = None,
trace_manager: Optional[Any] = None,
queue_message: Optional[Any] = None,
) -> Generator[Any, None, None]:
"""Dispatch events using elegant pattern matching."""
handlers = self._get_event_handlers()
event_type = type(event)
# Direct handler lookup
if handler := handlers.get(event_type):
yield from handler(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
return
# Handle node failure events with isinstance check
if isinstance(
event,
(
QueueNodeFailedEvent,
QueueNodeInIterationFailedEvent,
QueueNodeInLoopFailedEvent,
QueueNodeExceptionEvent,
),
):
yield from self._handle_node_failed_events(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
return
yield workflow_finish_resp # Handle parallel branch finished events with isinstance check
yield self._base_task_pipeline._error_to_stream_response(err) if isinstance(event, (QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent)):
break yield from self._handle_parallel_branch_finished_events(
elif isinstance(event, QueueStopEvent): event,
if self._workflow_run_id and graph_runtime_state: graph_runtime_state=graph_runtime_state,
with Session(db.engine, expire_on_commit=False) as session: tts_publisher=tts_publisher,
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed( trace_manager=trace_manager,
workflow_run_id=self._workflow_run_id, queue_message=queue_message,
total_tokens=graph_runtime_state.total_tokens, )
total_steps=graph_runtime_state.node_run_steps, return
status=WorkflowExecutionStatus.STOPPED,
error_message=event.get_stop_reason(),
conversation_id=self._conversation_id,
trace_manager=trace_manager,
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
# Save message
self._save_message(session=session, graph_runtime_state=graph_runtime_state)
session.commit()
yield workflow_finish_resp
elif event.stopped_by in (
QueueStopEvent.StopBy.INPUT_MODERATION,
QueueStopEvent.StopBy.ANNOTATION_REPLY,
):
# When hitting input-moderation or annotation-reply, the workflow will not start
with Session(db.engine, expire_on_commit=False) as session:
# Save message
self._save_message(session=session)
session.commit()
yield self._message_end_to_stream_response()
break
elif isinstance(event, QueueRetrieverResourcesEvent):
self._message_cycle_manager.handle_retriever_resources(event)
with Session(db.engine, expire_on_commit=False) as session:
message = self._get_message(session=session)
message.message_metadata = self._task_state.metadata.model_dump_json()
session.commit()
elif isinstance(event, QueueAnnotationReplyEvent):
self._message_cycle_manager.handle_annotation_reply(event)
with Session(db.engine, expire_on_commit=False) as session:
message = self._get_message(session=session)
message.message_metadata = self._task_state.metadata.model_dump_json()
session.commit()
elif isinstance(event, QueueTextChunkEvent):
delta_text = event.text
if delta_text is None:
continue
# handle output moderation chunk # For unhandled events, we continue (original behavior)
should_direct_answer = self._handle_output_moderation_chunk(delta_text) return
if should_direct_answer:
continue def _process_stream_response(
self,
tts_publisher: Optional[AppGeneratorTTSPublisher] = None,
trace_manager: Optional[TraceQueueManager] = None,
) -> Generator[StreamResponse, None, None]:
"""
Process stream response using elegant Fluent Python patterns.
Maintains exact same functionality as original 57-if-statement version.
"""
# Initialize graph runtime state
graph_runtime_state: Optional[GraphRuntimeState] = None
# only publish tts message at text chunk streaming for queue_message in self._base_task_pipeline._queue_manager.listen():
if tts_publisher: event = queue_message.event
tts_publisher.publish(queue_message)
self._task_state.answer += delta_text # Use elegant pattern matching for event dispatch
yield self._message_cycle_manager.message_to_stream_response( match event:
answer=delta_text, message_id=self._message_id, from_variable_selector=event.from_variable_selector # Handle QueueWorkflowStartedEvent specially - it has a side effect
) case QueueWorkflowStartedEvent():
elif isinstance(event, QueueMessageReplaceEvent): graph_runtime_state = event.graph_runtime_state # Side effect!
# published by moderation yield from self._handle_workflow_started_event(event)
yield self._message_cycle_manager.message_replace_to_stream_response(
answer=event.text, reason=event.reason # Handle QueueTextChunkEvent specially - needs queue_message
) case QueueTextChunkEvent():
elif isinstance(event, QueueAdvancedChatMessageEndEvent): yield from self._handle_text_chunk_event(
if not graph_runtime_state: event, tts_publisher=tts_publisher, queue_message=queue_message
raise ValueError("graph runtime state not initialized.") )
output_moderation_answer = self._base_task_pipeline._handle_output_moderation_when_task_finished( # Handle events that cause loop breaks
self._task_state.answer case QueueErrorEvent():
) yield from self._handle_error_event(event)
if output_moderation_answer: break
self._task_state.answer = output_moderation_answer
yield self._message_cycle_manager.message_replace_to_stream_response( case QueueWorkflowFailedEvent():
answer=output_moderation_answer, yield from self._handle_workflow_failed_event(
reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION, event, graph_runtime_state=graph_runtime_state, trace_manager=trace_manager
) )
break
# Save message case QueueStopEvent():
with Session(db.engine, expire_on_commit=False) as session: yield from self._handle_stop_event(
self._save_message(session=session, graph_runtime_state=graph_runtime_state) event, graph_runtime_state=graph_runtime_state, trace_manager=trace_manager
session.commit() )
break
yield self._message_end_to_stream_response()
elif isinstance(event, QueueAgentLogEvent): # Handle all other events through elegant dispatch
yield self._workflow_response_converter.handle_agent_log( case _:
task_id=self._application_generate_entity.task_id, event=event if responses := list(
) self._dispatch_event(
else: event,
continue graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
):
yield from responses
# Continue with next event (elegant fallthrough)
# publish None when task finished # Cleanup - publish None when task finished
if tts_publisher: if tts_publisher:
tts_publisher.publish(None) tts_publisher.publish(None)

721
api/core/app/apps/workflow/generate_task_pipeline.py
Unescape Escape View File

@ -1,7 +1,8 @@
import logging import logging
import time import time
from collections.abc import Generator from collections.abc import Callable, Generator
from typing import Optional, Union from contextlib import contextmanager
from typing import Any, Optional, Protocol, Union
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@ -246,316 +247,500 @@ class WorkflowAppGenerateTaskPipeline:
if tts_publisher: if tts_publisher:
yield MessageAudioEndStreamResponse(audio="", task_id=task_id) yield MessageAudioEndStreamResponse(audio="", task_id=task_id)
def _process_stream_response( class EventHandler(Protocol):
"""Protocol for event handlers."""
def __call__(
self,
event: Any,
*,
graph_runtime_state: Optional[Any] = None,
tts_publisher: Optional[Any] = None,
trace_manager: Optional[Any] = None,
queue_message: Optional[Any] = None,
) -> Generator[Any, None, None]: ...
@contextmanager
def _database_session(self):
"""Context manager for database sessions."""
with Session(db.engine, expire_on_commit=False) as session:
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
def _ensure_workflow_initialized(self) -> None:
"""Fluent validation for workflow state."""
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
def _ensure_graph_runtime_initialized(self, graph_runtime_state: Optional[Any]) -> Any:
"""Fluent validation for graph runtime state."""
if not graph_runtime_state:
raise ValueError("graph runtime state not initialized.")
return graph_runtime_state
def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle ping events."""
yield self._base_task_pipeline._ping_stream_response()
def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle error events."""
err = self._base_task_pipeline._handle_error(event=event)
yield self._base_task_pipeline._error_to_stream_response(err)
def _handle_workflow_started_event(self, event: QueueWorkflowStartedEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle workflow started events."""
# init workflow run
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
self._workflow_run_id = workflow_execution.id_
start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
yield start_resp
def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node retry events."""
self._ensure_workflow_initialized()
with self._database_session() as session:
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried(
workflow_execution_id=self._workflow_run_id,
event=event,
)
response = self._workflow_response_converter.workflow_node_retry_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if response:
yield response
def _handle_node_started_event(self, event: QueueNodeStartedEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node started events."""
self._ensure_workflow_initialized()
workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start(
workflow_execution_id=self._workflow_run_id, event=event
)
node_start_response = self._workflow_response_converter.workflow_node_start_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if node_start_response:
yield node_start_response
def _handle_node_succeeded_event(self, event: QueueNodeSucceededEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle node succeeded events."""
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success(event=event)
node_success_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
self._save_output_for_event(event, workflow_node_execution.id)
if node_success_response:
yield node_success_response
def _handle_node_failed_events(
self, self,
tts_publisher: Optional[AppGeneratorTTSPublisher] = None, event: Union[
trace_manager: Optional[TraceQueueManager] = None, QueueNodeFailedEvent, QueueNodeInIterationFailedEvent, QueueNodeInLoopFailedEvent, QueueNodeExceptionEvent
) -> Generator[StreamResponse, None, None]: ],
""" **kwargs,
Process stream response. ) -> Generator[Any, None, None]:
:return: """Handle various node failure events."""
""" workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(
graph_runtime_state = None event=event,
)
node_failed_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
for queue_message in self._base_task_pipeline._queue_manager.listen(): if isinstance(event, QueueNodeExceptionEvent):
event = queue_message.event self._save_output_for_event(event, workflow_node_execution.id)
if isinstance(event, QueuePingEvent): if node_failed_response:
yield self._base_task_pipeline._ping_stream_response() yield node_failed_response
elif isinstance(event, QueueErrorEvent):
err = self._base_task_pipeline._handle_error(event=event)
yield self._base_task_pipeline._error_to_stream_response(err)
break
elif isinstance(event, QueueWorkflowStartedEvent):
# override graph runtime state
graph_runtime_state = event.graph_runtime_state
# init workflow run
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_start()
self._workflow_run_id = workflow_execution.id_
start_resp = self._workflow_response_converter.workflow_start_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
yield start_resp def _handle_parallel_branch_started_event(
elif isinstance( self, event: QueueParallelBranchRunStartedEvent, **kwargs
event, ) -> Generator[Any, None, None]:
QueueNodeRetryEvent, """Handle parallel branch started events."""
): self._ensure_workflow_initialized()
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
with Session(db.engine, expire_on_commit=False) as session:
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_retried(
workflow_execution_id=self._workflow_run_id,
event=event,
)
response = self._workflow_response_converter.workflow_node_retry_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
session.commit()
if response: parallel_start_resp = self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response(
yield response task_id=self._application_generate_entity.task_id,
elif isinstance(event, QueueNodeStartedEvent): workflow_execution_id=self._workflow_run_id,
if not self._workflow_run_id: event=event,
raise ValueError("workflow run not initialized.") )
yield parallel_start_resp
workflow_node_execution = self._workflow_cycle_manager.handle_node_execution_start( def _handle_parallel_branch_finished_events(
workflow_execution_id=self._workflow_run_id, event=event self, event: Union[QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent], **kwargs
) ) -> Generator[Any, None, None]:
node_start_response = self._workflow_response_converter.workflow_node_start_to_stream_response( """Handle parallel branch finished events."""
event=event, self._ensure_workflow_initialized()
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if node_start_response: parallel_finish_resp = self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response(
yield node_start_response task_id=self._application_generate_entity.task_id,
elif isinstance(event, QueueNodeSucceededEvent): workflow_execution_id=self._workflow_run_id,
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_success( event=event,
event=event )
) yield parallel_finish_resp
node_success_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
self._save_output_for_event(event, workflow_node_execution.id) def _handle_iteration_start_event(self, event: QueueIterationStartEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle iteration start events."""
self._ensure_workflow_initialized()
if node_success_response: iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response(
yield node_success_response task_id=self._application_generate_entity.task_id,
elif isinstance( workflow_execution_id=self._workflow_run_id,
event, event=event,
QueueNodeFailedEvent )
| QueueNodeInIterationFailedEvent yield iter_start_resp
| QueueNodeInLoopFailedEvent
| QueueNodeExceptionEvent,
):
workflow_node_execution = self._workflow_cycle_manager.handle_workflow_node_execution_failed(
event=event,
)
node_failed_response = self._workflow_response_converter.workflow_node_finish_to_stream_response(
event=event,
task_id=self._application_generate_entity.task_id,
workflow_node_execution=workflow_node_execution,
)
if isinstance(event, QueueNodeExceptionEvent):
self._save_output_for_event(event, workflow_node_execution.id)
if node_failed_response: def _handle_iteration_next_event(self, event: QueueIterationNextEvent, **kwargs) -> Generator[Any, None, None]:
yield node_failed_response """Handle iteration next events."""
self._ensure_workflow_initialized()
elif isinstance(event, QueueParallelBranchRunStartedEvent): iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response(
if not self._workflow_run_id: task_id=self._application_generate_entity.task_id,
raise ValueError("workflow run not initialized.") workflow_execution_id=self._workflow_run_id,
event=event,
)
yield iter_next_resp
parallel_start_resp = ( def _handle_iteration_completed_event(
self._workflow_response_converter.workflow_parallel_branch_start_to_stream_response( self, event: QueueIterationCompletedEvent, **kwargs
task_id=self._application_generate_entity.task_id, ) -> Generator[Any, None, None]:
workflow_execution_id=self._workflow_run_id, """Handle iteration completed events."""
event=event, self._ensure_workflow_initialized()
)
)
yield parallel_start_resp iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id,
event=event,
)
yield iter_finish_resp
elif isinstance(event, QueueParallelBranchRunSucceededEvent | QueueParallelBranchRunFailedEvent): def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[Any, None, None]:
if not self._workflow_run_id: """Handle loop start events."""
raise ValueError("workflow run not initialized.") self._ensure_workflow_initialized()
parallel_finish_resp = ( loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response(
self._workflow_response_converter.workflow_parallel_branch_finished_to_stream_response( task_id=self._application_generate_entity.task_id,
task_id=self._application_generate_entity.task_id, workflow_execution_id=self._workflow_run_id,
workflow_execution_id=self._workflow_run_id, event=event,
event=event, )
) yield loop_start_resp
)
yield parallel_finish_resp def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[Any, None, None]:
"""Handle loop next events."""
self._ensure_workflow_initialized()
elif isinstance(event, QueueIterationStartEvent): loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response(
if not self._workflow_run_id: task_id=self._application_generate_entity.task_id,
raise ValueError("workflow run not initialized.") workflow_execution_id=self._workflow_run_id,
event=event,
)
yield loop_next_resp
iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response( def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) -> Generator[Any, None, None]:
task_id=self._application_generate_entity.task_id, """Handle loop completed events."""
workflow_execution_id=self._workflow_run_id, self._ensure_workflow_initialized()
event=event,
)
yield iter_start_resp loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response(
task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id,
event=event,
)
yield loop_finish_resp
elif isinstance(event, QueueIterationNextEvent): def _handle_workflow_succeeded_event(
if not self._workflow_run_id: self,
raise ValueError("workflow run not initialized.") event: QueueWorkflowSucceededEvent,
*,
graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow succeeded events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
outputs=event.outputs,
conversation_id=None,
trace_manager=trace_manager,
)
iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response( # save workflow app log
task_id=self._application_generate_entity.task_id, self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
workflow_execution_id=self._workflow_run_id,
event=event,
)
yield iter_next_resp workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
session=session,
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
elif isinstance(event, QueueIterationCompletedEvent): yield workflow_finish_resp
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response( def _handle_workflow_partial_success_event(
task_id=self._application_generate_entity.task_id, self,
workflow_execution_id=self._workflow_run_id, event: QueueWorkflowPartialSuccessEvent,
event=event, *,
) graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow partial success events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
outputs=event.outputs,
exceptions_count=event.exceptions_count,
conversation_id=None,
trace_manager=trace_manager,
)
yield iter_finish_resp # save workflow app log
self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
elif isinstance(event, QueueLoopStartEvent): workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
if not self._workflow_run_id: session=session,
raise ValueError("workflow run not initialized.") task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response( yield workflow_finish_resp
task_id=self._application_generate_entity.task_id,
workflow_execution_id=self._workflow_run_id,
event=event,
)
yield loop_start_resp def _handle_workflow_failed_and_stop_events(
self,
event: Union[QueueWorkflowFailedEvent, QueueStopEvent],
*,
graph_runtime_state: Optional[Any] = None,
trace_manager: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle workflow failed and stop events."""
self._ensure_workflow_initialized()
validated_state = self._ensure_graph_runtime_initialized(graph_runtime_state)
with self._database_session() as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
workflow_run_id=self._workflow_run_id,
total_tokens=validated_state.total_tokens,
total_steps=validated_state.node_run_steps,
status=WorkflowExecutionStatus.FAILED
if isinstance(event, QueueWorkflowFailedEvent)
else WorkflowExecutionStatus.STOPPED,
error_message=event.error if isinstance(event, QueueWorkflowFailedEvent) else event.get_stop_reason(),
conversation_id=None,
trace_manager=trace_manager,
exceptions_count=event.exceptions_count if isinstance(event, QueueWorkflowFailedEvent) else 0,
)
elif isinstance(event, QueueLoopNextEvent): # save workflow app log
if not self._workflow_run_id: self._save_workflow_app_log(session=session, workflow_execution=workflow_execution)
raise ValueError("workflow run not initialized.")
loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response( workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response(
task_id=self._application_generate_entity.task_id, session=session,
workflow_execution_id=self._workflow_run_id, task_id=self._application_generate_entity.task_id,
event=event, workflow_execution=workflow_execution,
) )
yield loop_next_resp yield workflow_finish_resp
elif isinstance(event, QueueLoopCompletedEvent): def _handle_text_chunk_event(
if not self._workflow_run_id: self,
raise ValueError("workflow run not initialized.") event: QueueTextChunkEvent,
*,
tts_publisher: Optional[Any] = None,
queue_message: Optional[Any] = None,
**kwargs,
) -> Generator[Any, None, None]:
"""Handle text chunk events."""
delta_text = event.text
if delta_text is None:
return
loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response( # only publish tts message at text chunk streaming
task_id=self._application_generate_entity.task_id, if tts_publisher and queue_message:
workflow_execution_id=self._workflow_run_id, tts_publisher.publish(queue_message)
event=event,
)
yield loop_finish_resp yield self._text_chunk_to_stream_response(delta_text, from_variable_selector=event.from_variable_selector)
elif isinstance(event, QueueWorkflowSucceededEvent):
if not self._workflow_run_id:
raise ValueError("workflow run not initialized.")
if not graph_runtime_state:
raise ValueError("graph runtime state not initialized.")
with Session(db.engine, expire_on_commit=False) as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_success(
workflow_run_id=self._workflow_run_id,
total_tokens=graph_runtime_state.total_tokens,
total_steps=graph_runtime_state.node_run_steps,
outputs=event.outputs,
conversation_id=None,
trace_manager=trace_manager,
)
# save workflow app log def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[Any, None, None]:
self._save_workflow_app_log(session=session, workflow_execution=workflow_execution) """Handle agent log events."""
yield self._workflow_response_converter.handle_agent_log(
task_id=self._application_generate_entity.task_id, event=event
)
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( def _get_event_handlers(self) -> dict[type, Callable]:
session=session, """Get mapping of event types to their handlers using fluent pattern."""
task_id=self._application_generate_entity.task_id, return {
workflow_execution=workflow_execution, # Basic events
) QueuePingEvent: self._handle_ping_event,
session.commit() QueueErrorEvent: self._handle_error_event,
QueueTextChunkEvent: self._handle_text_chunk_event,
yield workflow_finish_resp # Workflow events
elif isinstance(event, QueueWorkflowPartialSuccessEvent): QueueWorkflowStartedEvent: self._handle_workflow_started_event,
if not self._workflow_run_id: QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event,
raise ValueError("workflow run not initialized.") QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event,
if not graph_runtime_state: # Node events
raise ValueError("graph runtime state not initialized.") QueueNodeRetryEvent: self._handle_node_retry_event,
QueueNodeStartedEvent: self._handle_node_started_event,
with Session(db.engine, expire_on_commit=False) as session: QueueNodeSucceededEvent: self._handle_node_succeeded_event,
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_partial_success( # Parallel branch events
workflow_run_id=self._workflow_run_id, QueueParallelBranchRunStartedEvent: self._handle_parallel_branch_started_event,
total_tokens=graph_runtime_state.total_tokens, # Iteration events
total_steps=graph_runtime_state.node_run_steps, QueueIterationStartEvent: self._handle_iteration_start_event,
outputs=event.outputs, QueueIterationNextEvent: self._handle_iteration_next_event,
exceptions_count=event.exceptions_count, QueueIterationCompletedEvent: self._handle_iteration_completed_event,
conversation_id=None, # Loop events
trace_manager=trace_manager, QueueLoopStartEvent: self._handle_loop_start_event,
) QueueLoopNextEvent: self._handle_loop_next_event,
QueueLoopCompletedEvent: self._handle_loop_completed_event,
# Agent events
QueueAgentLogEvent: self._handle_agent_log_event,
}
def _dispatch_event(
self,
event: Any,
*,
graph_runtime_state: Optional[Any] = None,
tts_publisher: Optional[Any] = None,
trace_manager: Optional[Any] = None,
queue_message: Optional[Any] = None,
) -> Generator[Any, None, None]:
"""Dispatch events using elegant pattern matching."""
handlers = self._get_event_handlers()
event_type = type(event)
# Direct handler lookup
if handler := handlers.get(event_type):
yield from handler(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
return
# save workflow app log # Handle node failure events with isinstance check
self._save_workflow_app_log(session=session, workflow_execution=workflow_execution) if isinstance(
event,
(
QueueNodeFailedEvent,
QueueNodeInIterationFailedEvent,
QueueNodeInLoopFailedEvent,
QueueNodeExceptionEvent,
),
):
yield from self._handle_node_failed_events(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
return
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( # Handle parallel branch finished events with isinstance check
session=session, if isinstance(event, (QueueParallelBranchRunSucceededEvent, QueueParallelBranchRunFailedEvent)):
task_id=self._application_generate_entity.task_id, yield from self._handle_parallel_branch_finished_events(
workflow_execution=workflow_execution, event,
) graph_runtime_state=graph_runtime_state,
session.commit() tts_publisher=tts_publisher,
trace_manager=trace_manager,
yield workflow_finish_resp queue_message=queue_message,
elif isinstance(event, QueueWorkflowFailedEvent | QueueStopEvent): )
if not self._workflow_run_id: return
raise ValueError("workflow run not initialized.")
if not graph_runtime_state:
raise ValueError("graph runtime state not initialized.")
with Session(db.engine, expire_on_commit=False) as session:
workflow_execution = self._workflow_cycle_manager.handle_workflow_run_failed(
workflow_run_id=self._workflow_run_id,
total_tokens=graph_runtime_state.total_tokens,
total_steps=graph_runtime_state.node_run_steps,
status=WorkflowExecutionStatus.FAILED
if isinstance(event, QueueWorkflowFailedEvent)
else WorkflowExecutionStatus.STOPPED,
error_message=event.error
if isinstance(event, QueueWorkflowFailedEvent)
else event.get_stop_reason(),
conversation_id=None,
trace_manager=trace_manager,
exceptions_count=event.exceptions_count if isinstance(event, QueueWorkflowFailedEvent) else 0,
)
# save workflow app log # Handle workflow failed and stop events with isinstance check
self._save_workflow_app_log(session=session, workflow_execution=workflow_execution) if isinstance(event, (QueueWorkflowFailedEvent, QueueStopEvent)):
yield from self._handle_workflow_failed_and_stop_events(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
return
workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( # For unhandled events, we continue (original behavior)
session=session, return
task_id=self._application_generate_entity.task_id,
workflow_execution=workflow_execution,
)
session.commit()
yield workflow_finish_resp def _process_stream_response(
elif isinstance(event, QueueTextChunkEvent): self,
delta_text = event.text tts_publisher: Optional[AppGeneratorTTSPublisher] = None,
if delta_text is None: trace_manager: Optional[TraceQueueManager] = None,
continue ) -> Generator[StreamResponse, None, None]:
"""
Process stream response using elegant Fluent Python patterns.
Maintains exact same functionality as original 44-if-statement version.
"""
# Initialize graph runtime state
graph_runtime_state = None
# only publish tts message at text chunk streaming for queue_message in self._base_task_pipeline._queue_manager.listen():
if tts_publisher: event = queue_message.event
tts_publisher.publish(queue_message)
yield self._text_chunk_to_stream_response( # Use elegant pattern matching for event dispatch
delta_text, from_variable_selector=event.from_variable_selector match event:
) # Handle QueueWorkflowStartedEvent specially - it has a side effect
elif isinstance(event, QueueAgentLogEvent): case QueueWorkflowStartedEvent():
yield self._workflow_response_converter.handle_agent_log( graph_runtime_state = event.graph_runtime_state # Side effect!
task_id=self._application_generate_entity.task_id, event=event yield from self._handle_workflow_started_event(event)
)
else: # Handle QueueTextChunkEvent specially - needs queue_message
continue case QueueTextChunkEvent():
yield from self._handle_text_chunk_event(
event, tts_publisher=tts_publisher, queue_message=queue_message
)
# Handle events that cause loop breaks
case QueueErrorEvent():
yield from self._handle_error_event(event)
break
# Handle all other events through elegant dispatch
case _:
if responses := list(
self._dispatch_event(
event,
graph_runtime_state=graph_runtime_state,
tts_publisher=tts_publisher,
trace_manager=trace_manager,
queue_message=queue_message,
)
):
yield from responses
# Continue with next event (elegant fallthrough)
# Cleanup - publish None when task finished
if tts_publisher: if tts_publisher:
tts_publisher.publish(None) tts_publisher.publish(None)

Write Preview
Loading…
Cancel
Save