| import logging |
| import threading |
| import time |
| from typing import Any, Optional |
|
|
| from flask import Flask, current_app |
| from pydantic import BaseModel, ConfigDict |
|
|
| from configs import dify_config |
| from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom |
| from core.app.entities.queue_entities import QueueMessageReplaceEvent |
| from core.moderation.base import ModerationAction, ModerationOutputsResult |
| from core.moderation.factory import ModerationFactory |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class ModerationRule(BaseModel): |
| type: str |
| config: dict[str, Any] |
|
|
|
|
| class OutputModeration(BaseModel): |
| tenant_id: str |
| app_id: str |
|
|
| rule: ModerationRule |
| queue_manager: AppQueueManager |
|
|
| thread: Optional[threading.Thread] = None |
| thread_running: bool = True |
| buffer: str = "" |
| is_final_chunk: bool = False |
| final_output: Optional[str] = None |
| model_config = ConfigDict(arbitrary_types_allowed=True) |
|
|
| def should_direct_output(self) -> bool: |
| return self.final_output is not None |
|
|
| def get_final_output(self) -> str: |
| return self.final_output or "" |
|
|
| def append_new_token(self, token: str) -> None: |
| self.buffer += token |
|
|
| if not self.thread: |
| self.thread = self.start_thread() |
|
|
| def moderation_completion(self, completion: str, public_event: bool = False) -> str: |
| self.buffer = completion |
| self.is_final_chunk = True |
|
|
| result = self.moderation(tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=completion) |
|
|
| if not result or not result.flagged: |
| return completion |
|
|
| if result.action == ModerationAction.DIRECT_OUTPUT: |
| final_output = result.preset_response |
| else: |
| final_output = result.text |
|
|
| if public_event: |
| self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) |
|
|
| return final_output |
|
|
| def start_thread(self) -> threading.Thread: |
| buffer_size = dify_config.MODERATION_BUFFER_SIZE |
| thread = threading.Thread( |
| target=self.worker, |
| kwargs={ |
| "flask_app": current_app._get_current_object(), |
| "buffer_size": buffer_size if buffer_size > 0 else dify_config.MODERATION_BUFFER_SIZE, |
| }, |
| ) |
|
|
| thread.start() |
|
|
| return thread |
|
|
| def stop_thread(self): |
| if self.thread and self.thread.is_alive(): |
| self.thread_running = False |
|
|
| def worker(self, flask_app: Flask, buffer_size: int): |
| with flask_app.app_context(): |
| current_length = 0 |
| while self.thread_running: |
| moderation_buffer = self.buffer |
| buffer_length = len(moderation_buffer) |
| if not self.is_final_chunk: |
| chunk_length = buffer_length - current_length |
| if 0 <= chunk_length < buffer_size: |
| time.sleep(1) |
| continue |
|
|
| current_length = buffer_length |
|
|
| result = self.moderation( |
| tenant_id=self.tenant_id, app_id=self.app_id, moderation_buffer=moderation_buffer |
| ) |
|
|
| if not result or not result.flagged: |
| continue |
|
|
| if result.action == ModerationAction.DIRECT_OUTPUT: |
| final_output = result.preset_response |
| self.final_output = final_output |
| else: |
| final_output = result.text + self.buffer[len(moderation_buffer) :] |
|
|
| |
| if self.thread_running: |
| self.queue_manager.publish(QueueMessageReplaceEvent(text=final_output), PublishFrom.TASK_PIPELINE) |
|
|
| if result.action == ModerationAction.DIRECT_OUTPUT: |
| break |
|
|
| def moderation(self, tenant_id: str, app_id: str, moderation_buffer: str) -> Optional[ModerationOutputsResult]: |
| try: |
| moderation_factory = ModerationFactory( |
| name=self.rule.type, app_id=app_id, tenant_id=tenant_id, config=self.rule.config |
| ) |
|
|
| result: ModerationOutputsResult = moderation_factory.moderation_for_outputs(moderation_buffer) |
| return result |
| except Exception as e: |
| logger.error("Moderation Output error: %s", e) |
|
|
| return None |
|
|
