Co-authored-by: 聂政 <niezheng@pjlab.org.cn>tags/1.3.1
| ) | ) | ||||
| elif isinstance(event, QueueMessageReplaceEvent): | elif isinstance(event, QueueMessageReplaceEvent): | ||||
| # published by moderation | # published by moderation | ||||
| yield self._message_cycle_manager._message_replace_to_stream_response(answer=event.text) | |||||
| yield self._message_cycle_manager._message_replace_to_stream_response( | |||||
| answer=event.text, reason=event.reason | |||||
| ) | |||||
| elif isinstance(event, QueueAdvancedChatMessageEndEvent): | elif isinstance(event, QueueAdvancedChatMessageEndEvent): | ||||
| if not graph_runtime_state: | if not graph_runtime_state: | ||||
| raise ValueError("graph runtime state not initialized.") | raise ValueError("graph runtime state not initialized.") | ||||
| if output_moderation_answer: | if output_moderation_answer: | ||||
| self._task_state.answer = output_moderation_answer | self._task_state.answer = output_moderation_answer | ||||
| yield self._message_cycle_manager._message_replace_to_stream_response( | yield self._message_cycle_manager._message_replace_to_stream_response( | ||||
| answer=output_moderation_answer | |||||
| answer=output_moderation_answer, | |||||
| reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION, | |||||
| ) | ) | ||||
| # Save message | # Save message |
| QueueMessageReplaceEvent entity | QueueMessageReplaceEvent entity | ||||
| """ | """ | ||||
| class MessageReplaceReason(StrEnum): | |||||
| """ | |||||
| Reason for message replace event | |||||
| """ | |||||
| OUTPUT_MODERATION = "output_moderation" | |||||
| event: QueueEvent = QueueEvent.MESSAGE_REPLACE | event: QueueEvent = QueueEvent.MESSAGE_REPLACE | ||||
| text: str | text: str | ||||
| reason: str | |||||
| class QueueRetrieverResourcesEvent(AppQueueEvent): | class QueueRetrieverResourcesEvent(AppQueueEvent): |
| event: StreamEvent = StreamEvent.MESSAGE_REPLACE | event: StreamEvent = StreamEvent.MESSAGE_REPLACE | ||||
| answer: str | answer: str | ||||
| reason: str | |||||
| class AgentThoughtStreamResponse(StreamResponse): | class AgentThoughtStreamResponse(StreamResponse): |
| if self._output_moderation_handler: | if self._output_moderation_handler: | ||||
| self._output_moderation_handler.stop_thread() | self._output_moderation_handler.stop_thread() | ||||
| completion = self._output_moderation_handler.moderation_completion( | |||||
| completion, flagged = self._output_moderation_handler.moderation_completion( | |||||
| completion=completion, public_event=False | completion=completion, public_event=False | ||||
| ) | ) | ||||
| self._output_moderation_handler = None | self._output_moderation_handler = None | ||||
| return completion | |||||
| if flagged: | |||||
| return completion | |||||
| return None | return None |
| from_variable_selector=from_variable_selector, | from_variable_selector=from_variable_selector, | ||||
| ) | ) | ||||
| def _message_replace_to_stream_response(self, answer: str) -> MessageReplaceStreamResponse: | |||||
| def _message_replace_to_stream_response(self, answer: str, reason: str = "") -> MessageReplaceStreamResponse: | |||||
| """ | """ | ||||
| Message replace to stream response. | Message replace to stream response. | ||||
| :param answer: answer | :param answer: answer | ||||
| :return: | :return: | ||||
| """ | """ | ||||
| return MessageReplaceStreamResponse(task_id=self._application_generate_entity.task_id, answer=answer) | |||||
| return MessageReplaceStreamResponse( | |||||
| task_id=self._application_generate_entity.task_id, answer=answer, reason=reason | |||||
| ) |
| if not self.thread: | if not self.thread: | ||||
| self.thread = self.start_thread() | self.thread = self.start_thread() | ||||
| def moderation_completion(self, completion: str, public_event: bool = False) -> str: | |||||
| def moderation_completion(self, completion: str, public_event: bool = False) -> tuple[str, bool]: | |||||
| self.buffer = completion | self.buffer = completion | ||||
| self.is_final_chunk = True | self.is_final_chunk = True | ||||
| result = self.moderation(tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=completion) | result = self.moderation(tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=completion) | ||||
| if not result or not result.flagged: | if not result or not result.flagged: | ||||
| return completion | |||||
| return completion, False | |||||
| if result.action == ModerationAction.DIRECT_OUTPUT: | if result.action == ModerationAction.DIRECT_OUTPUT: | ||||
| final_output = result.preset_response | final_output = result.preset_response | ||||
| final_output = result.text | final_output = result.text | ||||
| if public_event: | if public_event: | ||||
| self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) | |||||
| self.queue_manager.publish( | |||||
| QueueMessageReplaceEvent( | |||||
| text=final_output, reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION | |||||
| ), | |||||
| PublishFrom.TASK_PIPELINE, | |||||
| ) | |||||
| return final_output | |||||
| return final_output, True | |||||
| def start_thread(self) -> threading.Thread: | def start_thread(self) -> threading.Thread: | ||||
| buffer_size = dify_config.MODERATION_BUFFER_SIZE | buffer_size = dify_config.MODERATION_BUFFER_SIZE | ||||
| # trigger replace event | # trigger replace event | ||||
| if self.thread_running: | if self.thread_running: | ||||
| self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) | |||||
| self.queue_manager.publish( | |||||
| QueueMessageReplaceEvent( | |||||
| text=final_output, reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION | |||||
| ), | |||||
| PublishFrom.TASK_PIPELINE, | |||||
| ) | |||||
| if result.action == ModerationAction.DIRECT_OUTPUT: | if result.action == ModerationAction.DIRECT_OUTPUT: | ||||
| break | break |