feat: add api-based extension & external data tool & moderation backend (#1403)

Co-authored-by: takatost <takatost@gmail.com>

api/.vscode/launch.json

@@ -0,0 +1,27 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Flask",
+            "type": "python",
+            "request": "launch",
+            "module": "flask",
+            "env": {
+                "FLASK_APP": "app.py",
+                "FLASK_DEBUG": "1",
+                "GEVENT_SUPPORT": "True"
+            },
+            "args": [
+                "run",
+                "--host=0.0.0.0",
+                "--port=5001",
+                "--debug"
+            ],
+            "jinja": true,
+            "justMyCode": true
+        }
+    ]
+}

api/app.py

@@ -19,7 +19,7 @@ from flask_cors import CORS
 
 from core.model_providers.providers import hosted
 from extensions import ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
-    ext_database, ext_storage, ext_mail, ext_stripe
+    ext_database, ext_storage, ext_mail, ext_stripe, ext_code_based_extension
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
 
@@ -79,6 +79,7 @@ def create_app(test_config=None) -> Flask:
 def initialize_extensions(app):
     # Since the application instance is now created, pass it to each Flask
     # extension instance to bind it to the Flask application instance (app)
+    ext_code_based_extension.init()
     ext_database.init_app(app)
     ext_migrate.init(app, db)
     ext_redis.init_app(app)

api/config.py

@@ -57,6 +57,7 @@ DEFAULTS = {
     'CLEAN_DAY_SETTING': 30,
     'UPLOAD_FILE_SIZE_LIMIT': 15,
     'UPLOAD_FILE_BATCH_LIMIT': 5,
+    'OUTPUT_MODERATION_BUFFER_SIZE': 300
 }
 
 
@@ -228,6 +229,9 @@ class Config:
         self.UPLOAD_FILE_SIZE_LIMIT = int(get_env('UPLOAD_FILE_SIZE_LIMIT'))
         self.UPLOAD_FILE_BATCH_LIMIT = int(get_env('UPLOAD_FILE_BATCH_LIMIT'))
 
+        # moderation settings
+        self.OUTPUT_MODERATION_BUFFER_SIZE = int(get_env('OUTPUT_MODERATION_BUFFER_SIZE'))
+
 
 class CloudEditionConfig(Config):
 

api/controllers/console/__init__.py

@@ -6,7 +6,7 @@ bp = Blueprint('console', __name__, url_prefix='/console/api')
 api = ExternalApi(bp)
 
 # Import other controllers
-from . import setup, version, apikey, admin
+from . import extension, setup, version, apikey, admin
 
 # Import app controllers
 from .app import advanced_prompt_template, app, site, completion, model_config, statistic, conversation, message, generator, audio

api/controllers/console/explore/parameter.py

@@ -27,6 +27,7 @@ class AppParameterApi(InstalledAppResource):
         'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
+        'sensitive_word_avoidance': fields.Raw
     }
 
     @marshal_with(parameters_fields)
@@ -42,7 +43,8 @@ class AppParameterApi(InstalledAppResource):
             'speech_to_text': app_model_config.speech_to_text_dict,
             'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
-            'user_input_form': app_model_config.user_input_form_list
+            'user_input_form': app_model_config.user_input_form_list,
+            'sensitive_word_avoidance': app_model_config.sensitive_word_avoidance_dict
         }
 
 

api/controllers/console/extension.py

@@ -0,0 +1,114 @@
+from flask_restful import Resource, reqparse, marshal_with
+from flask_login import current_user
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from libs.login import login_required
+from models.api_based_extension import APIBasedExtension
+from fields.api_based_extension_fields import api_based_extension_fields
+from services.code_based_extension_service import CodeBasedExtensionService
+from services.api_based_extension_service import APIBasedExtensionService
+
+
+class CodeBasedExtensionAPI(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('module', type=str, required=True, location='args')
+        args = parser.parse_args()
+
+        return {
+            'module': args['module'],
+            'data': CodeBasedExtensionService.get_code_based_extension(args['module'])
+        }
+
+
+class APIBasedExtensionAPI(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_based_extension_fields)
+    def get(self):
+        tenant_id = current_user.current_tenant_id
+        return APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_based_extension_fields)
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=True, location='json')
+        parser.add_argument('api_endpoint', type=str, required=True, location='json')
+        parser.add_argument('api_key', type=str, required=True, location='json')
+        args = parser.parse_args()
+
+        extension_data = APIBasedExtension(
+            tenant_id=current_user.current_tenant_id,
+            name=args['name'],
+            api_endpoint=args['api_endpoint'],
+            api_key=args['api_key']
+        )
+
+        return APIBasedExtensionService.save(extension_data)
+
+
+class APIBasedExtensionDetailAPI(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_based_extension_fields)
+    def get(self, id):
+        api_based_extension_id = str(id)
+        tenant_id = current_user.current_tenant_id
+
+        return APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_based_extension_fields)
+    def post(self, id):
+        api_based_extension_id = str(id)
+        tenant_id = current_user.current_tenant_id
+
+        extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=True, location='json')
+        parser.add_argument('api_endpoint', type=str, required=True, location='json')
+        parser.add_argument('api_key', type=str, required=True, location='json')
+        args = parser.parse_args()
+
+        extension_data_from_db.name = args['name']
+        extension_data_from_db.api_endpoint = args['api_endpoint']
+
+        if args['api_key'] != '[__HIDDEN__]':
+            extension_data_from_db.api_key = args['api_key']
+
+        return APIBasedExtensionService.save(extension_data_from_db)
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, id):
+        api_based_extension_id = str(id)
+        tenant_id = current_user.current_tenant_id
+
+        extension_data_from_db = APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+
+        APIBasedExtensionService.delete(extension_data_from_db)
+
+        return {'result': 'success'}
+
+
+api.add_resource(CodeBasedExtensionAPI, '/code-based-extension')
+
+api.add_resource(APIBasedExtensionAPI, '/api-based-extension')
+api.add_resource(APIBasedExtensionDetailAPI, '/api-based-extension/<uuid:id>')

api/controllers/service_api/app/app.py

@@ -28,6 +28,7 @@ class AppParameterApi(AppApiResource):
         'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
+        'sensitive_word_avoidance': fields.Raw
     }
 
     @marshal_with(parameters_fields)
@@ -42,7 +43,8 @@ class AppParameterApi(AppApiResource):
             'speech_to_text': app_model_config.speech_to_text_dict,
             'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
-            'user_input_form': app_model_config.user_input_form_list
+            'user_input_form': app_model_config.user_input_form_list,
+            'sensitive_word_avoidance': app_model_config.sensitive_word_avoidance_dict
         }
 
 

api/controllers/service_api/app/completion.py

@@ -183,4 +183,3 @@ api.add_resource(CompletionApi, '/completion-messages')
 api.add_resource(CompletionStopApi, '/completion-messages/<string:task_id>/stop')
 api.add_resource(ChatApi, '/chat-messages')
 api.add_resource(ChatStopApi, '/chat-messages/<string:task_id>/stop')
-

api/controllers/web/app.py

@@ -27,6 +27,7 @@ class AppParameterApi(WebApiResource):
         'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
+        'sensitive_word_avoidance': fields.Raw
     }
 
     @marshal_with(parameters_fields)
@@ -41,7 +42,8 @@ class AppParameterApi(WebApiResource):
             'speech_to_text': app_model_config.speech_to_text_dict,
             'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
-            'user_input_form': app_model_config.user_input_form_list
+            'user_input_form': app_model_config.user_input_form_list,
+            'sensitive_word_avoidance': app_model_config.sensitive_word_avoidance_dict
         }
 
 

api/controllers/web/completion.py

@@ -139,7 +139,7 @@ class ChatStopApi(WebApiResource):
         return {'result': 'success'}, 200
 
 
-def compact_response(response: Union[dict | Generator]) -> Response:
+def compact_response(response: Union[dict, Generator]) -> Response:
     if isinstance(response, dict):
         return Response(response=json.dumps(response), status=200, mimetype='application/json')
     else:

api/core/__init__.py

@@ -0,0 +1 @@
+import core.moderation.base

api/core/callback_handler/llm_callback_handler.py

@@ -1,13 +1,25 @@
 import logging
-from typing import Any, Dict, List, Union
+import threading
+import time
+from typing import Any, Dict, List, Union, Optional
 
+from flask import Flask, current_app
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import LLMResult, BaseMessage
+from pydantic import BaseModel
 
 from core.callback_handler.entity.llm_message import LLMMessage
-from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
+from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException, \
+    ConversationTaskInterruptException
 from core.model_providers.models.entity.message import to_prompt_messages, PromptMessage
 from core.model_providers.models.llm.base import BaseLLM
+from core.moderation.base import ModerationOutputsResult, ModerationAction
+from core.moderation.factory import ModerationFactory
+
+
+class ModerationRule(BaseModel):
+    type: str
+    config: Dict[str, Any]
 
 
 class LLMCallbackHandler(BaseCallbackHandler):
@@ -20,6 +32,24 @@ class LLMCallbackHandler(BaseCallbackHandler):
         self.start_at = None
         self.conversation_message_task = conversation_message_task
 
+        self.output_moderation_handler = None
+        self.init_output_moderation()
+
+    def init_output_moderation(self):
+        app_model_config = self.conversation_message_task.app_model_config
+        sensitive_word_avoidance_dict = app_model_config.sensitive_word_avoidance_dict
+
+        if sensitive_word_avoidance_dict and sensitive_word_avoidance_dict.get("enabled"):
+            self.output_moderation_handler = OutputModerationHandler(
+                tenant_id=self.conversation_message_task.tenant_id,
+                app_id=self.conversation_message_task.app.id,
+                rule=ModerationRule(
+                    type=sensitive_word_avoidance_dict.get("type"),
+                    config=sensitive_word_avoidance_dict.get("config")
+                ),
+                on_message_replace_func=self.conversation_message_task.on_message_replace
+            )
+
     @property
     def always_verbose(self) -> bool:
         """Whether to call verbose callbacks even if verbose is False."""
@@ -59,10 +89,19 @@ class LLMCallbackHandler(BaseCallbackHandler):
         self.llm_message.prompt_tokens = self.model_instance.get_num_tokens([PromptMessage(content=prompts[0])])
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        if not self.conversation_message_task.streaming:
-            self.conversation_message_task.append_message_text(response.generations[0][0].text)
+        if self.output_moderation_handler:
+            self.output_moderation_handler.stop_thread()
+
+            self.llm_message.completion = self.output_moderation_handler.moderation_completion(
+                completion=response.generations[0][0].text,
+                public_event=True if self.conversation_message_task.streaming else False
+            )
+        else:
             self.llm_message.completion = response.generations[0][0].text
 
+        if not self.conversation_message_task.streaming:
+            self.conversation_message_task.append_message_text(self.llm_message.completion)
+
         if response.llm_output and 'token_usage' in response.llm_output:
             if 'prompt_tokens' in response.llm_output['token_usage']:
                 self.llm_message.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
@@ -79,23 +118,161 @@ class LLMCallbackHandler(BaseCallbackHandler):
         self.conversation_message_task.save_message(self.llm_message)
 
     def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        try:
-            self.conversation_message_task.append_message_text(token)
-        except ConversationTaskStoppedException as ex:
+        if self.output_moderation_handler and self.output_moderation_handler.should_direct_output():
+            # stop subscribe new token when output moderation should direct output
+            ex = ConversationTaskInterruptException()
             self.on_llm_error(error=ex)
             raise ex
 
-        self.llm_message.completion += token
+        try:
+            self.conversation_message_task.append_message_text(token)
+            self.llm_message.completion += token
+
+            if self.output_moderation_handler:
+                self.output_moderation_handler.append_new_token(token)
+        except ConversationTaskStoppedException as ex:
+            self.on_llm_error(error=ex)
+            raise ex
 
     def on_llm_error(
             self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
         """Do nothing."""
+        if self.output_moderation_handler:
+            self.output_moderation_handler.stop_thread()
+
         if isinstance(error, ConversationTaskStoppedException):
             if self.conversation_message_task.streaming:
                 self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
                     [PromptMessage(content=self.llm_message.completion)]
                 )
                 self.conversation_message_task.save_message(llm_message=self.llm_message, by_stopped=True)
+        if isinstance(error, ConversationTaskInterruptException):
+            self.llm_message.completion = self.output_moderation_handler.get_final_output()
+            self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                [PromptMessage(content=self.llm_message.completion)]
+            )
+            self.conversation_message_task.save_message(llm_message=self.llm_message)
         else:
             logging.debug("on_llm_error: %s", error)
+
+
+class OutputModerationHandler(BaseModel):
+    DEFAULT_BUFFER_SIZE: int = 300
+
+    tenant_id: str
+    app_id: str
+
+    rule: ModerationRule
+    on_message_replace_func: Any
+
+    thread: Optional[threading.Thread] = None
+    thread_running: bool = True
+    buffer: str = ''
+    is_final_chunk: bool = False
+    final_output: Optional[str] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def should_direct_output(self):
+        return self.final_output is not None
+
+    def get_final_output(self):
+        return self.final_output
+
+    def append_new_token(self, token: str):
+        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.on_message_replace_func(final_output)
+
+        return final_output
+
+    def start_thread(self) -> threading.Thread:
+        buffer_size = int(current_app.config.get('MODERATION_BUFFER_SIZE', self.DEFAULT_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 self.DEFAULT_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):]
+
+                # trigger replace event
+                if self.thread_running:
+                    self.on_message_replace_func(final_output)
+
+                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:
+            logging.error("Moderation Output error: %s", e)
+
+        return None

api/core/chain/sensitive_word_avoidance_chain.py

@@ -1,92 +0,0 @@
-import enum
-import logging
-from typing import List, Dict, Optional, Any
-
-from langchain.callbacks.manager import CallbackManagerForChainRun
-from langchain.chains.base import Chain
-from pydantic import BaseModel
-
-from core.model_providers.error import LLMBadRequestError
-from core.model_providers.model_factory import ModelFactory
-from core.model_providers.models.llm.base import BaseLLM
-from core.model_providers.models.moderation import openai_moderation
-
-
-class SensitiveWordAvoidanceRule(BaseModel):
-    class Type(enum.Enum):
-        MODERATION = "moderation"
-        KEYWORDS = "keywords"
-
-    type: Type
-    canned_response: str = 'Your content violates our usage policy. Please revise and try again.'
-    extra_params: dict = {}
-
-
-class SensitiveWordAvoidanceChain(Chain):
-    input_key: str = "input"  #: :meta private:
-    output_key: str = "output"  #: :meta private:
-
-    model_instance: BaseLLM
-    sensitive_word_avoidance_rule: SensitiveWordAvoidanceRule
-
-    @property
-    def _chain_type(self) -> str:
-        return "sensitive_word_avoidance_chain"
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Expect input key.
-
-        :meta private:
-        """
-        return [self.input_key]
-
-    @property
-    def output_keys(self) -> List[str]:
-        """Return output key.
-
-        :meta private:
-        """
-        return [self.output_key]
-
-    def _check_sensitive_word(self, text: str) -> bool:
-        for word in self.sensitive_word_avoidance_rule.extra_params.get('sensitive_words', []):
-            if word in text:
-                return False
-        return True
-
-    def _check_moderation(self, text: str) -> bool:
-        moderation_model_instance = ModelFactory.get_moderation_model(
-            tenant_id=self.model_instance.model_provider.provider.tenant_id,
-            model_provider_name='openai',
-            model_name=openai_moderation.DEFAULT_MODEL
-        )
-
-        try:
-            return moderation_model_instance.run(text=text)
-        except Exception as ex:
-            logging.exception(ex)
-            raise LLMBadRequestError('Rate limit exceeded, please try again later.')
-
-    def _call(
-            self,
-            inputs: Dict[str, Any],
-            run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        text = inputs[self.input_key]
-
-        if self.sensitive_word_avoidance_rule.type == SensitiveWordAvoidanceRule.Type.KEYWORDS:
-            result = self._check_sensitive_word(text)
-        else:
-            result = self._check_moderation(text)
-
-        if not result:
-            raise SensitiveWordAvoidanceError(self.sensitive_word_avoidance_rule.canned_response)
-
-        return {self.output_key: text}
-
-
-class SensitiveWordAvoidanceError(Exception):
-    def __init__(self, message):
-        super().__init__(message)
-        self.message = message

api/core/completion.py

@@ -1,13 +1,18 @@
+import concurrent
+import json
 import logging
-from typing import Optional, List, Union
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional, List, Union, Tuple
 
+from flask import current_app, Flask
 from requests.exceptions import ChunkedEncodingError
 
 from core.agent.agent_executor import AgentExecuteResult, PlanningStrategy
 from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
 from core.callback_handler.llm_callback_handler import LLMCallbackHandler
-from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceError
-from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
+from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException, \
+    ConversationTaskInterruptException
+from core.external_data_tool.factory import ExternalDataToolFactory
 from core.model_providers.error import LLMBadRequestError
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
     ReadOnlyConversationTokenDBBufferSharedMemory
@@ -18,6 +23,8 @@ from core.orchestrator_rule_parser import OrchestratorRuleParser
 from core.prompt.prompt_template import PromptTemplateParser
 from core.prompt.prompt_transform import PromptTransform
 from models.model import App, AppModelConfig, Account, Conversation, EndUser
+from core.moderation.base import ModerationException, ModerationAction
+from core.moderation.factory import ModerationFactory
 
 
 class Completion:
@@ -76,26 +83,35 @@ class Completion:
         )
 
         try:
-            # parse sensitive_word_avoidance_chain
             chain_callback = MainChainGatherCallbackHandler(conversation_message_task)
-            sensitive_word_avoidance_chain = orchestrator_rule_parser.to_sensitive_word_avoidance_chain(
-                final_model_instance, [chain_callback])
-            if sensitive_word_avoidance_chain:
-                try:
-                    query = sensitive_word_avoidance_chain.run(query)
-                except SensitiveWordAvoidanceError as ex:
-                    cls.run_final_llm(
-                        model_instance=final_model_instance,
-                        mode=app.mode,
-                        app_model_config=app_model_config,
-                        query=query,
-                        inputs=inputs,
-                        agent_execute_result=None,
-                        conversation_message_task=conversation_message_task,
-                        memory=memory,
-                        fake_response=ex.message
-                    )
-                    return
+
+            try:
+                # process sensitive_word_avoidance
+                inputs, query = cls.moderation_for_inputs(app.id, app.tenant_id, app_model_config, inputs, query)
+            except ModerationException as e:
+                cls.run_final_llm(
+                    model_instance=final_model_instance,
+                    mode=app.mode,
+                    app_model_config=app_model_config,
+                    query=query,
+                    inputs=inputs,
+                    agent_execute_result=None,
+                    conversation_message_task=conversation_message_task,
+                    memory=memory,
+                    fake_response=str(e)
+                )
+                return
+
+            # fill in variable inputs from external data tools if exists
+            external_data_tools = app_model_config.external_data_tools_list
+            if external_data_tools:
+                inputs = cls.fill_in_inputs_from_external_data_tools(
+                    tenant_id=app.tenant_id,
+                    app_id=app.id,
+                    external_data_tools=external_data_tools,
+                    inputs=inputs,
+                    query=query
+                )
 
             # get agent executor
             agent_executor = orchestrator_rule_parser.to_agent_executor(
@@ -135,19 +151,110 @@ class Completion:
                 memory=memory,
                 fake_response=fake_response
             )
-        except ConversationTaskStoppedException:
+        except (ConversationTaskInterruptException, ConversationTaskStoppedException):
             return
         except ChunkedEncodingError as e:
             # Interrupt by LLM (like OpenAI), handle it.
             logging.warning(f'ChunkedEncodingError: {e}')
             conversation_message_task.end()
             return
-        
+
+    @classmethod
+    def moderation_for_inputs(cls, app_id: str, tenant_id: str, app_model_config: AppModelConfig, inputs: dict, query: str):
+        if not app_model_config.sensitive_word_avoidance_dict['enabled']:
+            return inputs, query
+
+        type = app_model_config.sensitive_word_avoidance_dict['type']
+
+        moderation = ModerationFactory(type, app_id, tenant_id, app_model_config.sensitive_word_avoidance_dict['config'])
+        moderation_result = moderation.moderation_for_inputs(inputs, query)
+
+        if not moderation_result.flagged:
+            return inputs, query
+
+        if moderation_result.action == ModerationAction.DIRECT_OUTPUT:
+            raise ModerationException(moderation_result.preset_response)
+        elif moderation_result.action == ModerationAction.OVERRIDED:
+            inputs = moderation_result.inputs
+            query = moderation_result.query
+
+        return inputs, query
+
+    @classmethod
+    def fill_in_inputs_from_external_data_tools(cls, tenant_id: str, app_id: str, external_data_tools: list[dict],
+                                                inputs: dict, query: str) -> dict:
+        """
+        Fill in variable inputs from external data tools if exists.
+
+        :param tenant_id: workspace id
+        :param app_id: app id
+        :param external_data_tools: external data tools configs
+        :param inputs: the inputs
+        :param query: the query
+        :return: the filled inputs
+        """
+        # Group tools by type and config
+        grouped_tools = {}
+        for tool in external_data_tools:
+            if not tool.get("enabled"):
+                continue
+
+            tool_key = (tool.get("type"), json.dumps(tool.get("config"), sort_keys=True))
+            grouped_tools.setdefault(tool_key, []).append(tool)
+
+        results = {}
+        with ThreadPoolExecutor() as executor:
+            futures = {}
+            for tools in grouped_tools.values():
+                # Only query the first tool in each group
+                first_tool = tools[0]
+                future = executor.submit(
+                    cls.query_external_data_tool, current_app._get_current_object(), tenant_id, app_id, first_tool,
+                    inputs, query
+                )
+                for tool in tools:
+                    futures[future] = tool
+
+            for future in concurrent.futures.as_completed(futures):
+                tool_key, result = future.result()
+                if tool_key in grouped_tools:
+                    for tool in grouped_tools[tool_key]:
+                        results[tool['variable']] = result
+
+        inputs.update(results)
+        return inputs
+
+    @classmethod
+    def query_external_data_tool(cls, flask_app: Flask, tenant_id: str, app_id: str, external_data_tool: dict,
+                                 inputs: dict, query: str) -> Tuple[Optional[str], Optional[str]]:
+        with flask_app.app_context():
+            tool_variable = external_data_tool.get("variable")
+            tool_type = external_data_tool.get("type")
+            tool_config = external_data_tool.get("config")
+
+            external_data_tool_factory = ExternalDataToolFactory(
+                name=tool_type,
+                tenant_id=tenant_id,
+                app_id=app_id,
+                variable=tool_variable,
+                config=tool_config
+            )
+
+            # query external data tool
+            result = external_data_tool_factory.query(
+                inputs=inputs,
+                query=query
+            )
+
+            tool_key = (external_data_tool.get("type"), json.dumps(external_data_tool.get("config"), sort_keys=True))
+
+            return tool_key, result
+
     @classmethod
     def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str:
         if app.mode != 'completion':
             return query
-        
+
         return inputs.get(app_model_config.dataset_query_variable, "")
 
     @classmethod

api/core/conversation_message_task.py

@@ -290,6 +290,10 @@ class ConversationMessageTask:
                 db.session.commit()
             self.retriever_resource = resource
 
+    def on_message_replace(self, text: str):
+        if text is not None:
+            self._pub_handler.pub_message_replace(text)
+
     def message_end(self):
         self._pub_handler.pub_message_end(self.retriever_resource)
 
@@ -342,6 +346,24 @@ class PubHandler:
             self.pub_end()
             raise ConversationTaskStoppedException()
 
+    def pub_message_replace(self, text: str):
+        content = {
+            'event': 'message_replace',
+            'data': {
+                'task_id': self._task_id,
+                'message_id': str(self._message.id),
+                'text': text,
+                'mode': self._conversation.mode,
+                'conversation_id': str(self._conversation.id)
+            }
+        }
+
+        redis_client.publish(self._channel, json.dumps(content))
+
+        if self._is_stopped():
+            self.pub_end()
+            raise ConversationTaskStoppedException()
+
     def pub_chain(self, message_chain: MessageChain):
         if self._chain_pub:
             content = {
@@ -443,3 +465,7 @@ class PubHandler:
 
 class ConversationTaskStoppedException(Exception):
     pass
+
+
+class ConversationTaskInterruptException(Exception):
+    pass

api/core/extension/api_based_extension_requestor.py

@@ -0,0 +1,62 @@
+import os
+
+import requests
+
+from models.api_based_extension import APIBasedExtensionPoint
+
+
+class APIBasedExtensionRequestor:
+    timeout: (int, int) = (5, 60)
+    """timeout for request connect and read"""
+
+    def __init__(self, api_endpoint: str, api_key: str) -> None:
+        self.api_endpoint = api_endpoint
+        self.api_key = api_key
+
+    def request(self, point: APIBasedExtensionPoint, params: dict) -> dict:
+        """
+        Request the api.
+
+        :param point: the api point
+        :param params: the request params
+        :return: the response json
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer {}".format(self.api_key)
+        }
+
+        url = self.api_endpoint
+
+        try:
+            # proxy support for security
+            proxies = None
+            if os.environ.get("API_BASED_EXTENSION_HTTP_PROXY") and os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"):
+                proxies = {
+                    'http': os.environ.get("API_BASED_EXTENSION_HTTP_PROXY"),
+                    'https': os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"),
+                }
+
+            response = requests.request(
+                method='POST',
+                url=url,
+                json={
+                    'point': point.value,
+                    'params': params
+                },
+                headers=headers,
+                timeout=self.timeout,
+                proxies=proxies
+            )
+        except requests.exceptions.Timeout:
+            raise ValueError("request timeout")
+        except requests.exceptions.ConnectionError:
+            raise ValueError("request connection error")
+
+        if response.status_code != 200:
+            raise ValueError("request error, status_code: {}, content: {}".format(
+                response.status_code,
+                response.text[:100]
+            ))
+
+        return response.json()

api/core/extension/extensible.py

@@ -0,0 +1,111 @@
+import enum
+import importlib.util
+import json
+import logging
+import os
+from collections import OrderedDict
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+class ExtensionModule(enum.Enum):
+    MODERATION = 'moderation'
+    EXTERNAL_DATA_TOOL = 'external_data_tool'
+
+
+class ModuleExtension(BaseModel):
+    extension_class: Any
+    name: str
+    label: Optional[dict] = None
+    form_schema: Optional[list] = None
+    builtin: bool = True
+    position: Optional[int] = None
+
+
+class Extensible:
+    module: ExtensionModule
+
+    name: str
+    tenant_id: str
+    config: Optional[dict] = None
+
+    def __init__(self, tenant_id: str, config: Optional[dict] = None) -> None:
+        self.tenant_id = tenant_id
+        self.config = config
+
+    @classmethod
+    def scan_extensions(cls):
+        extensions = {}
+
+        # get the path of the current class
+        current_path = os.path.abspath(cls.__module__.replace(".", os.path.sep) + '.py')
+        current_dir_path = os.path.dirname(current_path)
+
+        # traverse subdirectories
+        for subdir_name in os.listdir(current_dir_path):
+            if subdir_name.startswith('__'):
+                continue
+
+            subdir_path = os.path.join(current_dir_path, subdir_name)
+            extension_name = subdir_name
+            if os.path.isdir(subdir_path):
+                file_names = os.listdir(subdir_path)
+
+                # is builtin extension, builtin extension
+                # in the front-end page and business logic, there are special treatments.
+                builtin = False
+                position = None
+                if '__builtin__' in file_names:
+                    builtin = True
+
+                    builtin_file_path = os.path.join(subdir_path, '__builtin__')
+                    if os.path.exists(builtin_file_path):
+                        with open(builtin_file_path, 'r') as f:
+                            position = int(f.read().strip())
+
+                if (extension_name + '.py') not in file_names:
+                    logging.warning(f"Missing {extension_name}.py file in {subdir_path}, Skip.")
+                    continue
+
+                # Dynamic loading {subdir_name}.py file and find the subclass of Extensible
+                py_path = os.path.join(subdir_path, extension_name + '.py')
+                spec = importlib.util.spec_from_file_location(extension_name, py_path)
+                mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(mod)
+
+                extension_class = None
+                for name, obj in vars(mod).items():
+                    if isinstance(obj, type) and issubclass(obj, cls) and obj != cls:
+                        extension_class = obj
+                        break
+
+                if not extension_class:
+                    logging.warning(f"Missing subclass of {cls.__name__} in {py_path}, Skip.")
+                    continue
+
+                json_data = {}
+                if not builtin:
+                    if 'schema.json' not in file_names:
+                        logging.warning(f"Missing schema.json file in {subdir_path}, Skip.")
+                        continue
+
+                    json_path = os.path.join(subdir_path, 'schema.json')
+                    json_data = {}
+                    if os.path.exists(json_path):
+                        with open(json_path, 'r') as f:
+                            json_data = json.load(f)
+
+                extensions[extension_name] = ModuleExtension(
+                    extension_class=extension_class,
+                    name=extension_name,
+                    label=json_data.get('label'),
+                    form_schema=json_data.get('form_schema'),
+                    builtin=builtin,
+                    position=position
+                )
+
+        sorted_items = sorted(extensions.items(), key=lambda x: (x[1].position is None, x[1].position))
+        sorted_extensions = OrderedDict(sorted_items)
+
+        return sorted_extensions

api/core/extension/extension.py

@@ -0,0 +1,47 @@
+from core.extension.extensible import ModuleExtension, ExtensionModule
+from core.external_data_tool.base import ExternalDataTool
+from core.moderation.base import Moderation
+
+
+class Extension:
+    __module_extensions: dict[str, dict[str, ModuleExtension]] = {}
+
+    module_classes = {
+        ExtensionModule.MODERATION: Moderation,
+        ExtensionModule.EXTERNAL_DATA_TOOL: ExternalDataTool
+    }
+
+    def init(self):
+        for module, module_class in self.module_classes.items():
+            self.__module_extensions[module.value] = module_class.scan_extensions()
+
+    def module_extensions(self, module: str) -> list[ModuleExtension]:
+        module_extensions = self.__module_extensions.get(module)
+
+        if not module_extensions:
+            raise ValueError(f"Extension Module {module} not found")
+
+        return list(module_extensions.values())
+
+    def module_extension(self, module: ExtensionModule, extension_name: str) -> ModuleExtension:
+        module_extensions = self.__module_extensions.get(module.value)
+
+        if not module_extensions:
+            raise ValueError(f"Extension Module {module} not found")
+
+        module_extension = module_extensions.get(extension_name)
+
+        if not module_extension:
+            raise ValueError(f"Extension {extension_name} not found")
+
+        return module_extension
+
+    def extension_class(self, module: ExtensionModule, extension_name: str) -> type:
+        module_extension = self.module_extension(module, extension_name)
+        return module_extension.extension_class
+
+    def validate_form_schema(self, module: ExtensionModule, extension_name: str, config: dict) -> None:
+        module_extension = self.module_extension(module, extension_name)
+        form_schema = module_extension.form_schema
+
+        # TODO validate form_schema