pref: optimize feature model_load_balancing_enabled value fetch speed… (#4933)

Co-authored-by: Jyong <76649700+johnjyong@users.noreply.github.com>

.github/workflows/api-tests.yml

@@ -4,9 +4,17 @@ on:
   pull_request:
     branches:
       - main
+    paths:
+      - api/**
+      - docker/**
+
+concurrency:
+  group: api-tests-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
 
 jobs:
   test:
+    name: API Tests
     runs-on: ubuntu-latest
     strategy:
       matrix:

.github/workflows/db-migration-test.yml

@@ -0,0 +1,53 @@
+name: DB Migration Test
+
+on:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - api/migrations/**
+
+concurrency:
+  group: db-migration-test-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  db-migration-test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.10"
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+          cache-dependency-path: |
+            ./api/requirements.txt
+
+      - name: Install dependencies
+        run: pip install -r ./api/requirements.txt
+
+      - name: Set up Middleware
+        uses: hoverkraft-tech/compose-action@v2.0.0
+        with:
+          compose-file: |
+            docker/docker-compose.middleware.yaml
+          services: |
+            db
+
+      - name: Prepare configs
+        run: |
+          cd api
+          cp .env.example .env
+
+      - name: Run DB Migration
+        run: |
+          cd api
+          flask db upgrade

.github/workflows/style.yml

@@ -6,7 +6,7 @@ on:
       - main
 
 concurrency:
-  group: dep-${{ github.head_ref || github.run_id }}
+  group: style-${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
 
 jobs:
@@ -18,54 +18,89 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
+      - name: Check changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v44
+        with:
+          files: api/**
+
       - name: Set up Python
         uses: actions/setup-python@v5
+        if: steps.changed-files.outputs.any_changed == 'true'
         with:
           python-version: '3.10'
 
       - name: Python dependencies
+        if: steps.changed-files.outputs.any_changed == 'true'
         run: pip install ruff dotenv-linter
 
       - name: Ruff check
+        if: steps.changed-files.outputs.any_changed == 'true'
         run: ruff check ./api
 
       - name: Dotenv check
+        if: steps.changed-files.outputs.any_changed == 'true'
         run: dotenv-linter ./api/.env.example ./web/.env.example
 
       - name: Lint hints
         if: failure()
         run: echo "Please run 'dev/reformat' to fix the fixable linting errors."
 
-  test:
-    name: ESLint and SuperLinter
+  web-style:
+    name: Web Style
     runs-on: ubuntu-latest
-    needs: python-style
+    defaults:
+      run:
+        working-directory: ./web
 
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+
+      - name: Check changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v44
         with:
-          fetch-depth: 0
+          files: web/**
 
       - name: Setup NodeJS
         uses: actions/setup-node@v4
+        if: steps.changed-files.outputs.any_changed == 'true'
         with:
           node-version: 20
           cache: yarn
           cache-dependency-path: ./web/package.json
 
       - name: Web dependencies
-        run: |
-          cd ./web
-          yarn install --frozen-lockfile
+        if: steps.changed-files.outputs.any_changed == 'true'
+        run: yarn install --frozen-lockfile
 
       - name: Web style check
-        run: |
-          cd ./web
-          yarn run lint
+        if: steps.changed-files.outputs.any_changed == 'true'
+        run: yarn run lint
+
+
+  superlinter:
+    name: SuperLinter
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Check changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v44
+        with:
+          files: |
+            **.sh
+            **.yaml
+            **.yml
+            Dockerfile
 
       - name: Super-linter
         uses: super-linter/super-linter/slim@v6
+        if: steps.changed-files.outputs.any_changed == 'true'
         env:
           BASH_SEVERITY: warning
           DEFAULT_BRANCH: main
@@ -76,4 +111,5 @@ jobs:
           VALIDATE_BASH_EXEC: true
           VALIDATE_GITHUB_ACTIONS: true
           VALIDATE_DOCKERFILE_HADOLINT: true
+          VALIDATE_XML: true
           VALIDATE_YAML: true

.github/workflows/tool-test-sdks.yaml

@@ -4,6 +4,13 @@ on:
   pull_request:
     branches:
       - main
+    paths:
+      - sdks/**
+
+concurrency:
+  group: sdk-tests-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
 jobs:
   build:
     name: unit test for Node.js SDK

.gitignore

@@ -134,7 +134,8 @@ dmypy.json
 web/.vscode/settings.json
 
 # Intellij IDEA Files
-.idea/
+.idea/*
+!.idea/vcs.xml
 .ideaDataSources/
 
 api/.env

.idea/vcs.xml

@@ -0,0 +1,16 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="IssueNavigationConfiguration">
+    <option name="links">
+      <list>
+        <IssueNavigationLink>
+          <option name="issueRegexp" value="#(\d+)" />
+          <option name="linkRegexp" value="https://github.com/langgenius/dify/issues/$1" />
+        </IssueNavigationLink>
+      </list>
+    </option>
+  </component>
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

api/.env.example

@@ -127,10 +127,11 @@ RESEND_API_KEY=
 RESEND_API_URL=https://api.resend.com
 # smtp configuration
 SMTP_SERVER=smtp.gmail.com
-SMTP_PORT=587
+SMTP_PORT=465
 SMTP_USERNAME=123
 SMTP_PASSWORD=abc
-SMTP_USE_TLS=false
+SMTP_USE_TLS=true
+SMTP_OPPORTUNISTIC_TLS=false
 
 # Sentry configuration
 SENTRY_DSN=
@@ -182,3 +183,8 @@ LOG_FILE=
 
 # Indexing configuration
 INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH=1000
+
+# Workflow runtime configuration
+WORKFLOW_MAX_EXECUTION_STEPS=500
+WORKFLOW_MAX_EXECUTION_TIME=1200
+WORKFLOW_CALL_MAX_DEPTH=5

api/config.py

@@ -70,6 +70,7 @@ DEFAULTS = {
     'INVITE_EXPIRY_HOURS': 72,
     'BILLING_ENABLED': 'False',
     'CAN_REPLACE_LOGO': 'False',
+    'MODEL_LB_ENABLED': 'False',
     'ETL_TYPE': 'dify',
     'KEYWORD_STORE': 'jieba',
     'BATCH_UPLOAD_LIMIT': 20,
@@ -81,8 +82,9 @@ DEFAULTS = {
     'INNER_API': 'False',
     'ENTERPRISE_ENABLED': 'False',
     'INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH': 1000,
-    'WORKFLOW_MAX_EXECUTION_STEPS': 50,
-    'WORKFLOW_MAX_EXECUTION_TIME': 600,
+    'WORKFLOW_MAX_EXECUTION_STEPS': 500,
+    'WORKFLOW_MAX_EXECUTION_TIME': 1200,
+    'WORKFLOW_CALL_MAX_DEPTH': 5,
 }
 
 
@@ -113,7 +115,7 @@ class Config:
         # ------------------------
         # General Configurations.
         # ------------------------
-        self.CURRENT_VERSION = "0.6.9"
+        self.CURRENT_VERSION = "0.6.10"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = get_env('EDITION')
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')
@@ -122,6 +124,7 @@ class Config:
         self.LOG_FILE = get_env('LOG_FILE')
         self.LOG_FORMAT = get_env('LOG_FORMAT')
         self.LOG_DATEFORMAT = get_env('LOG_DATEFORMAT')
+        self.API_COMPRESSION_ENABLED = get_bool_env('API_COMPRESSION_ENABLED')
 
         # The backend URL prefix of the console API.
         # used to concatenate the login authorization callback or notion integration callback.
@@ -209,27 +212,41 @@ class Config:
             if self.CELERY_BACKEND == 'database' else self.CELERY_BROKER_URL
         self.BROKER_USE_SSL = self.CELERY_BROKER_URL.startswith('rediss://')
 
+        # ------------------------
+        # Code Execution Sandbox Configurations.
+        # ------------------------
+        self.CODE_EXECUTION_ENDPOINT = get_env('CODE_EXECUTION_ENDPOINT')
+        self.CODE_EXECUTION_API_KEY = get_env('CODE_EXECUTION_API_KEY')
+
         # ------------------------
         # File Storage Configurations.
         # ------------------------
         self.STORAGE_TYPE = get_env('STORAGE_TYPE')
         self.STORAGE_LOCAL_PATH = get_env('STORAGE_LOCAL_PATH')
+
+        # S3 Storage settings
         self.S3_ENDPOINT = get_env('S3_ENDPOINT')
         self.S3_BUCKET_NAME = get_env('S3_BUCKET_NAME')
         self.S3_ACCESS_KEY = get_env('S3_ACCESS_KEY')
         self.S3_SECRET_KEY = get_env('S3_SECRET_KEY')
         self.S3_REGION = get_env('S3_REGION')
         self.S3_ADDRESS_STYLE = get_env('S3_ADDRESS_STYLE')
+
+        # Azure Blob Storage settings
         self.AZURE_BLOB_ACCOUNT_NAME = get_env('AZURE_BLOB_ACCOUNT_NAME')
         self.AZURE_BLOB_ACCOUNT_KEY = get_env('AZURE_BLOB_ACCOUNT_KEY')
         self.AZURE_BLOB_CONTAINER_NAME = get_env('AZURE_BLOB_CONTAINER_NAME')
         self.AZURE_BLOB_ACCOUNT_URL = get_env('AZURE_BLOB_ACCOUNT_URL')
+
+        # Aliyun Storage settings
         self.ALIYUN_OSS_BUCKET_NAME = get_env('ALIYUN_OSS_BUCKET_NAME')
         self.ALIYUN_OSS_ACCESS_KEY = get_env('ALIYUN_OSS_ACCESS_KEY')
         self.ALIYUN_OSS_SECRET_KEY = get_env('ALIYUN_OSS_SECRET_KEY')
         self.ALIYUN_OSS_ENDPOINT = get_env('ALIYUN_OSS_ENDPOINT')
         self.ALIYUN_OSS_REGION = get_env('ALIYUN_OSS_REGION')
         self.ALIYUN_OSS_AUTH_VERSION = get_env('ALIYUN_OSS_AUTH_VERSION')
+
+        # Google Cloud Storage settings
         self.GOOGLE_STORAGE_BUCKET_NAME = get_env('GOOGLE_STORAGE_BUCKET_NAME')
         self.GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64 = get_env('GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64')
 
@@ -239,6 +256,7 @@ class Config:
         # ------------------------
         self.VECTOR_STORE = get_env('VECTOR_STORE')
         self.KEYWORD_STORE = get_env('KEYWORD_STORE')
+
         # qdrant settings
         self.QDRANT_URL = get_env('QDRANT_URL')
         self.QDRANT_API_KEY = get_env('QDRANT_API_KEY')
@@ -294,6 +312,7 @@ class Config:
         self.SMTP_USERNAME = get_env('SMTP_USERNAME')
         self.SMTP_PASSWORD = get_env('SMTP_PASSWORD')
         self.SMTP_USE_TLS = get_bool_env('SMTP_USE_TLS')
+        self.SMTP_OPPORTUNISTIC_TLS = get_bool_env('SMTP_OPPORTUNISTIC_TLS')
 
         # ------------------------
         # Workspace Configurations.
@@ -321,9 +340,23 @@ 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'))
         self.UPLOAD_IMAGE_FILE_SIZE_LIMIT = int(get_env('UPLOAD_IMAGE_FILE_SIZE_LIMIT'))
+        self.BATCH_UPLOAD_LIMIT = get_env('BATCH_UPLOAD_LIMIT')
+
+        # RAG ETL Configurations.
+        self.ETL_TYPE = get_env('ETL_TYPE')
+        self.UNSTRUCTURED_API_URL = get_env('UNSTRUCTURED_API_URL')
+        self.UNSTRUCTURED_API_KEY = get_env('UNSTRUCTURED_API_KEY')
+        self.KEYWORD_DATA_SOURCE_TYPE = get_env('KEYWORD_DATA_SOURCE_TYPE')
+
+        # Indexing Configurations.
+        self.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH = get_env('INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH')
+
+        # Tool Configurations.
+        self.TOOL_ICON_CACHE_MAX_AGE = get_env('TOOL_ICON_CACHE_MAX_AGE')
 
         self.WORKFLOW_MAX_EXECUTION_STEPS = int(get_env('WORKFLOW_MAX_EXECUTION_STEPS'))
         self.WORKFLOW_MAX_EXECUTION_TIME = int(get_env('WORKFLOW_MAX_EXECUTION_TIME'))
+        self.WORKFLOW_CALL_MAX_DEPTH = int(get_env('WORKFLOW_CALL_MAX_DEPTH'))
 
         # Moderation in app Configurations.
         self.OUTPUT_MODERATION_BUFFER_SIZE = int(get_env('OUTPUT_MODERATION_BUFFER_SIZE'))
@@ -375,24 +408,15 @@ class Config:
         self.HOSTED_FETCH_APP_TEMPLATES_MODE = get_env('HOSTED_FETCH_APP_TEMPLATES_MODE')
         self.HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN = get_env('HOSTED_FETCH_APP_TEMPLATES_REMOTE_DOMAIN')
 
-        self.ETL_TYPE = get_env('ETL_TYPE')
-        self.UNSTRUCTURED_API_URL = get_env('UNSTRUCTURED_API_URL')
-        self.UNSTRUCTURED_API_KEY = get_env('UNSTRUCTURED_API_KEY')
+        # Model Load Balancing Configurations.
+        self.MODEL_LB_ENABLED = get_bool_env('MODEL_LB_ENABLED')
+
+        # Platform Billing Configurations.
         self.BILLING_ENABLED = get_bool_env('BILLING_ENABLED')
-        self.CAN_REPLACE_LOGO = get_bool_env('CAN_REPLACE_LOGO')
 
-        self.BATCH_UPLOAD_LIMIT = get_env('BATCH_UPLOAD_LIMIT')
-
-        self.CODE_EXECUTION_ENDPOINT = get_env('CODE_EXECUTION_ENDPOINT')
-        self.CODE_EXECUTION_API_KEY = get_env('CODE_EXECUTION_API_KEY')
-
-        self.API_COMPRESSION_ENABLED = get_bool_env('API_COMPRESSION_ENABLED')
-        self.TOOL_ICON_CACHE_MAX_AGE = get_env('TOOL_ICON_CACHE_MAX_AGE')
-
-        self.KEYWORD_DATA_SOURCE_TYPE = get_env('KEYWORD_DATA_SOURCE_TYPE')
+        # ------------------------
+        # Enterprise feature Configurations.
+        # **Before using, please contact business@dify.ai by email to inquire about licensing matters.**
+        # ------------------------
         self.ENTERPRISE_ENABLED = get_bool_env('ENTERPRISE_ENABLED')
-
-        # ------------------------
-        # Indexing Configurations.
-        # ------------------------
-        self.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH = get_env('INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH')
+        self.CAN_REPLACE_LOGO = get_bool_env('CAN_REPLACE_LOGO')

api/controllers/console/__init__.py

@@ -54,4 +54,4 @@ from .explore import (
 from .tag import tags
 
 # Import workspace controllers
-from .workspace import account, members, model_providers, models, tool_providers, workspace
+from .workspace import account, load_balancing_config, members, model_providers, models, tool_providers, workspace

api/controllers/console/app/audio.py

@@ -85,7 +85,7 @@ class ChatMessageTextApi(Resource):
             response = AudioService.transcript_tts(
                 app_model=app_model,
                 text=request.form['text'],
-                voice=request.form['voice'] if request.form.get('voice') else app_model.app_model_config.text_to_speech_dict.get('voice'),
+                voice=request.form['voice'],
                 streaming=False
             )
 

api/controllers/console/datasets/datasets_document.py

@@ -1,10 +1,12 @@
 import logging
+from argparse import ArgumentTypeError
 from datetime import datetime, timezone
 
 from flask import request
 from flask_login import current_user
 from flask_restful import Resource, fields, marshal, marshal_with, reqparse
 from sqlalchemy import asc, desc
+from transformers.hf_argparser import string_to_bool
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -141,7 +143,11 @@ class DatasetDocumentListApi(Resource):
         limit = request.args.get('limit', default=20, type=int)
         search = request.args.get('keyword', default=None, type=str)
         sort = request.args.get('sort', default='-created_at', type=str)
-        fetch = request.args.get('fetch', default=False, type=bool)
+        # "yes", "true", "t", "y", "1" convert to True, while others convert to False.
+        try:
+            fetch = string_to_bool(request.args.get('fetch', default='false'))
+        except (ArgumentTypeError, ValueError, Exception) as e:
+            fetch = False
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
@@ -924,6 +930,28 @@ class DocumentRetryApi(DocumentResource):
         return {'result': 'success'}, 204
 
 
+class DocumentRenameApi(DocumentResource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(document_fields)
+    def post(self, dataset_id, document_id):
+        # The role of the current user in the ta table must be admin or owner
+        if not current_user.is_admin_or_owner:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        try:
+            document = DocumentService.rename_document(dataset_id, document_id, args['name'])
+        except services.errors.document.DocumentIndexingError:
+            raise DocumentIndexingError('Cannot delete document during indexing.')
+
+        return document
+
+
 api.add_resource(GetProcessRuleApi, '/datasets/process-rule')
 api.add_resource(DatasetDocumentListApi,
                  '/datasets/<uuid:dataset_id>/documents')
@@ -950,3 +978,5 @@ api.add_resource(DocumentStatusApi,
 api.add_resource(DocumentPauseApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause')
 api.add_resource(DocumentRecoverApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume')
 api.add_resource(DocumentRetryApi, '/datasets/<uuid:dataset_id>/retry')
+api.add_resource(DocumentRenameApi,
+                 '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/rename')

api/controllers/console/feature.py

@@ -1,14 +1,19 @@
 from flask_login import current_user
 from flask_restful import Resource
 
+from libs.login import login_required
 from services.feature_service import FeatureService
 
 from . import api
-from .wraps import cloud_utm_record
+from .setup import setup_required
+from .wraps import account_initialization_required, cloud_utm_record
 
 
 class FeatureApi(Resource):
 
+    @setup_required
+    @login_required
+    @account_initialization_required
     @cloud_utm_record
     def get(self):
         return FeatureService.get_features(current_user.current_tenant_id).dict()

api/controllers/console/version.py

@@ -17,13 +17,19 @@ class VersionApi(Resource):
         args = parser.parse_args()
         check_update_url = current_app.config['CHECK_UPDATE_URL']
 
-        if not check_update_url:
-            return {
-                'version': '0.0.0',
-                'release_date': '',
-                'release_notes': '',
-                'can_auto_update': False
+        result = {
+            'version': current_app.config['CURRENT_VERSION'],
+            'release_date': '',
+            'release_notes': '',
+            'can_auto_update': False,
+            'features': {
+                'can_replace_logo': current_app.config['CAN_REPLACE_LOGO'],
+                'model_load_balancing_enabled': current_app.config['MODEL_LB_ENABLED']
             }
+        }
+
+        if not check_update_url:
+            return result
 
         try:
             response = requests.get(check_update_url, {
@@ -31,20 +37,15 @@ class VersionApi(Resource):
             })
         except Exception as error:
             logging.warning("Check update version error: {}.".format(str(error)))
-            return {
-                'version': args.get('current_version'),
-                'release_date': '',
-                'release_notes': '',
-                'can_auto_update': False
-            }
+            result['version'] = args.get('current_version')
+            return result
 
         content = json.loads(response.content)
-        return {
-            'version': content['version'],
-            'release_date': content['releaseDate'],
-            'release_notes': content['releaseNotes'],
-            'can_auto_update': content['canAutoUpdate']
-        }
+        result['version'] = content['version']
+        result['release_date'] = content['releaseDate']
+        result['release_notes'] = content['releaseNotes']
+        result['can_auto_update'] = content['canAutoUpdate']
+        return result
 
 
 api.add_resource(VersionApi, '/version')

api/controllers/console/workspace/load_balancing_config.py

@@ -0,0 +1,106 @@
+from flask_restful import Resource, reqparse
+from werkzeug.exceptions import Forbidden
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.errors.validate import CredentialsValidateFailedError
+from libs.login import current_user, login_required
+from models.account import TenantAccountRole
+from services.model_load_balancing_service import ModelLoadBalancingService
+
+
+class LoadBalancingCredentialsValidateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider: str):
+        if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):
+            raise Forbidden()
+
+        tenant_id = current_user.current_tenant_id
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=[mt.value for mt in ModelType], location='json')
+        parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        # validate model load balancing credentials
+        model_load_balancing_service = ModelLoadBalancingService()
+
+        result = True
+        error = None
+
+        try:
+            model_load_balancing_service.validate_load_balancing_credentials(
+                tenant_id=tenant_id,
+                provider=provider,
+                model=args['model'],
+                model_type=args['model_type'],
+                credentials=args['credentials']
+            )
+        except CredentialsValidateFailedError as ex:
+            result = False
+            error = str(ex)
+
+        response = {'result': 'success' if result else 'error'}
+
+        if not result:
+            response['error'] = error
+
+        return response
+
+
+class LoadBalancingConfigCredentialsValidateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider: str, config_id: str):
+        if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):
+            raise Forbidden()
+
+        tenant_id = current_user.current_tenant_id
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=[mt.value for mt in ModelType], location='json')
+        parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        # validate model load balancing config credentials
+        model_load_balancing_service = ModelLoadBalancingService()
+
+        result = True
+        error = None
+
+        try:
+            model_load_balancing_service.validate_load_balancing_credentials(
+                tenant_id=tenant_id,
+                provider=provider,
+                model=args['model'],
+                model_type=args['model_type'],
+                credentials=args['credentials'],
+                config_id=config_id,
+            )
+        except CredentialsValidateFailedError as ex:
+            result = False
+            error = str(ex)
+
+        response = {'result': 'success' if result else 'error'}
+
+        if not result:
+            response['error'] = error
+
+        return response
+
+
+# Load Balancing Config
+api.add_resource(LoadBalancingCredentialsValidateApi,
+                 '/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/credentials-validate')
+
+api.add_resource(LoadBalancingConfigCredentialsValidateApi,
+                 '/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/<string:config_id>/credentials-validate')

api/controllers/console/workspace/models.py

@@ -12,6 +12,7 @@ from core.model_runtime.errors.validate import CredentialsValidateFailedError
 from core.model_runtime.utils.encoders import jsonable_encoder
 from libs.login import login_required
 from models.account import TenantAccountRole
+from services.model_load_balancing_service import ModelLoadBalancingService
 from services.model_provider_service import ModelProviderService
 
 
@@ -104,21 +105,56 @@ class ModelProviderModelApi(Resource):
         parser.add_argument('model', type=str, required=True, nullable=False, location='json')
         parser.add_argument('model_type', type=str, required=True, nullable=False,
                             choices=[mt.value for mt in ModelType], location='json')
-        parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+        parser.add_argument('credentials', type=dict, required=False, nullable=True, location='json')
+        parser.add_argument('load_balancing', type=dict, required=False, nullable=True, location='json')
+        parser.add_argument('config_from', type=str, required=False, nullable=True, location='json')
         args = parser.parse_args()
 
-        model_provider_service = ModelProviderService()
+        model_load_balancing_service = ModelLoadBalancingService()
 
-        try:
-            model_provider_service.save_model_credentials(
+        if ('load_balancing' in args and args['load_balancing'] and
+                'enabled' in args['load_balancing'] and args['load_balancing']['enabled']):
+            if 'configs' not in args['load_balancing']:
+                raise ValueError('invalid load balancing configs')
+
+            # save load balancing configs
+            model_load_balancing_service.update_load_balancing_configs(
                 tenant_id=tenant_id,
                 provider=provider,
                 model=args['model'],
                 model_type=args['model_type'],
-                credentials=args['credentials']
+                configs=args['load_balancing']['configs']
             )
-        except CredentialsValidateFailedError as ex:
-            raise ValueError(str(ex))
+
+            # enable load balancing
+            model_load_balancing_service.enable_model_load_balancing(
+                tenant_id=tenant_id,
+                provider=provider,
+                model=args['model'],
+                model_type=args['model_type']
+            )
+        else:
+            # disable load balancing
+            model_load_balancing_service.disable_model_load_balancing(
+                tenant_id=tenant_id,
+                provider=provider,
+                model=args['model'],
+                model_type=args['model_type']
+            )
+
+            if args.get('config_from', '') != 'predefined-model':
+                model_provider_service = ModelProviderService()
+
+                try:
+                    model_provider_service.save_model_credentials(
+                        tenant_id=tenant_id,
+                        provider=provider,
+                        model=args['model'],
+                        model_type=args['model_type'],
+                        credentials=args['credentials']
+                    )
+                except CredentialsValidateFailedError as ex:
+                    raise ValueError(str(ex))
 
         return {'result': 'success'}, 200
 
@@ -170,11 +206,73 @@ class ModelProviderModelCredentialApi(Resource):
             model=args['model']
         )
 
+        model_load_balancing_service = ModelLoadBalancingService()
+        is_load_balancing_enabled, load_balancing_configs = model_load_balancing_service.get_load_balancing_configs(
+            tenant_id=tenant_id,
+            provider=provider,
+            model=args['model'],
+            model_type=args['model_type']
+        )
+
         return {
-            "credentials": credentials
+            "credentials": credentials,
+            "load_balancing": {
+                "enabled": is_load_balancing_enabled,
+                "configs": load_balancing_configs
+            }
         }
 
 
+class ModelProviderModelEnableApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(self, provider: str):
+        tenant_id = current_user.current_tenant_id
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=[mt.value for mt in ModelType], location='json')
+        args = parser.parse_args()
+
+        model_provider_service = ModelProviderService()
+        model_provider_service.enable_model(
+            tenant_id=tenant_id,
+            provider=provider,
+            model=args['model'],
+            model_type=args['model_type']
+        )
+
+        return {'result': 'success'}
+
+
+class ModelProviderModelDisableApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(self, provider: str):
+        tenant_id = current_user.current_tenant_id
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=[mt.value for mt in ModelType], location='json')
+        args = parser.parse_args()
+
+        model_provider_service = ModelProviderService()
+        model_provider_service.disable_model(
+            tenant_id=tenant_id,
+            provider=provider,
+            model=args['model'],
+            model_type=args['model_type']
+        )
+
+        return {'result': 'success'}
+
+
 class ModelProviderModelValidateApi(Resource):
 
     @setup_required
@@ -259,6 +357,10 @@ class ModelProviderAvailableModelApi(Resource):
 
 
 api.add_resource(ModelProviderModelApi, '/workspaces/current/model-providers/<string:provider>/models')
+api.add_resource(ModelProviderModelEnableApi, '/workspaces/current/model-providers/<string:provider>/models/enable',
+                 endpoint='model-provider-model-enable')
+api.add_resource(ModelProviderModelDisableApi, '/workspaces/current/model-providers/<string:provider>/models/disable',
+                 endpoint='model-provider-model-disable')
 api.add_resource(ModelProviderModelCredentialApi,
                  '/workspaces/current/model-providers/<string:provider>/models/credentials')
 api.add_resource(ModelProviderModelValidateApi,

api/controllers/service_api/wraps.py

@@ -8,7 +8,7 @@ from flask import current_app, request
 from flask_login import user_logged_in
 from flask_restful import Resource
 from pydantic import BaseModel
-from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
+from werkzeug.exceptions import Forbidden, Unauthorized
 
 from extensions.ext_database import db
 from libs.login import _get_user
@@ -39,17 +39,17 @@ def validate_app_token(view: Optional[Callable] = None, *, fetch_user_arg: Optio
 
             app_model = db.session.query(App).filter(App.id == api_token.app_id).first()
             if not app_model:
-                raise NotFound()
+                raise Forbidden("The app no longer exists.")
 
             if app_model.status != 'normal':
-                raise NotFound()
+                raise Forbidden("The app's status is abnormal.")
 
             if not app_model.enable_api:
-                raise NotFound()
+                raise Forbidden("The app's API service has been disabled.")
 
             tenant = db.session.query(Tenant).filter(Tenant.id == app_model.tenant_id).first()
             if tenant.status == TenantStatus.ARCHIVE:
-                raise NotFound()
+                raise Forbidden("The workspace's status is archived.")
 
             kwargs['app_model'] = app_model
 

api/core/agent/base_agent_runner.py

@@ -39,6 +39,7 @@ from core.tools.entities.tool_entities import (
 from core.tools.tool.dataset_retriever_tool import DatasetRetrieverTool
 from core.tools.tool.tool import Tool
 from core.tools.tool_manager import ToolManager
+from core.tools.utils.tool_parameter_converter import ToolParameterConverter
 from extensions.ext_database import db
 from models.model import Conversation, Message, MessageAgentThought
 from models.tools import ToolConversationVariables
@@ -128,6 +129,8 @@ class BaseAgentRunner(AppRunner):
             self.files = application_generate_entity.files
         else:
             self.files = []
+        self.query = None
+        self._current_thoughts: list[PromptMessage] = []
 
     def _repack_app_generate_entity(self, app_generate_entity: AgentChatAppGenerateEntity) \
             -> AgentChatAppGenerateEntity:
@@ -184,21 +187,11 @@ class BaseAgentRunner(AppRunner):
             if parameter.form != ToolParameter.ToolParameterForm.LLM:
                 continue
 
-            parameter_type = 'string'
+            parameter_type = ToolParameterConverter.get_parameter_type(parameter.type)
             enum = []
-            if parameter.type == ToolParameter.ToolParameterType.STRING:
-                parameter_type = 'string'
-            elif parameter.type == ToolParameter.ToolParameterType.BOOLEAN:
-                parameter_type = 'boolean'
-            elif parameter.type == ToolParameter.ToolParameterType.NUMBER:
-                parameter_type = 'number'
-            elif parameter.type == ToolParameter.ToolParameterType.SELECT:
-                for option in parameter.options:
-                    enum.append(option.value)
-                parameter_type = 'string'
-            else:
-                raise ValueError(f"parameter type {parameter.type} is not supported")
-            
+            if parameter.type == ToolParameter.ToolParameterType.SELECT:
+                enum = [option.value for option in parameter.options]
+
             message_tool.parameters['properties'][parameter.name] = {
                 "type": parameter_type,
                 "description": parameter.llm_description or '',
@@ -279,20 +272,10 @@ class BaseAgentRunner(AppRunner):
             if parameter.form != ToolParameter.ToolParameterForm.LLM:
                 continue
 
-            parameter_type = 'string'
+            parameter_type = ToolParameterConverter.get_parameter_type(parameter.type)
             enum = []
-            if parameter.type == ToolParameter.ToolParameterType.STRING:
-                parameter_type = 'string'
-            elif parameter.type == ToolParameter.ToolParameterType.BOOLEAN:
-                parameter_type = 'boolean'
-            elif parameter.type == ToolParameter.ToolParameterType.NUMBER:
-                parameter_type = 'number'
-            elif parameter.type == ToolParameter.ToolParameterType.SELECT:
-                for option in parameter.options:
-                    enum.append(option.value)
-                parameter_type = 'string'
-            else:
-                raise ValueError(f"parameter type {parameter.type} is not supported")
+            if parameter.type == ToolParameter.ToolParameterType.SELECT:
+                enum = [option.value for option in parameter.options]
         
             prompt_tool.parameters['properties'][parameter.name] = {
                 "type": parameter_type,
@@ -464,7 +447,7 @@ class BaseAgentRunner(AppRunner):
         for message in messages:
             if message.id == self.message.id:
                 continue
-            
+
             result.append(self.organize_agent_user_prompt(message))
             agent_thoughts: list[MessageAgentThought] = message.agent_thoughts
             if agent_thoughts:
@@ -545,3 +528,4 @@ class BaseAgentRunner(AppRunner):
                 return UserPromptMessage(content=prompt_message_contents)
         else:
             return UserPromptMessage(content=message.query)
+         

api/core/agent/cot_agent_runner.py

@@ -15,6 +15,7 @@ from core.model_runtime.entities.message_entities import (
     ToolPromptMessage,
     UserPromptMessage,
 )
+from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform
 from core.tools.entities.tool_entities import ToolInvokeMeta
 from core.tools.tool.tool import Tool
 from core.tools.tool_engine import ToolEngine
@@ -373,7 +374,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
 
         return message
 
-    def _organize_historic_prompt_messages(self) -> list[PromptMessage]:
+    def _organize_historic_prompt_messages(self, current_session_messages: list[PromptMessage] = None) -> list[PromptMessage]:
         """
             organize historic prompt messages
         """
@@ -381,6 +382,13 @@ class CotAgentRunner(BaseAgentRunner, ABC):
         scratchpad: list[AgentScratchpadUnit] = []
         current_scratchpad: AgentScratchpadUnit = None
 
+        self.history_prompt_messages = AgentHistoryPromptTransform(
+            model_config=self.model_config,
+            prompt_messages=current_session_messages or [],
+            history_messages=self.history_prompt_messages,
+            memory=self.memory
+        ).get_prompt()
+
         for message in self.history_prompt_messages:
             if isinstance(message, AssistantPromptMessage):
                 current_scratchpad = AgentScratchpadUnit(

api/core/agent/cot_chat_agent_runner.py

@@ -32,9 +32,6 @@ class CotChatAgentRunner(CotAgentRunner):
         # organize system prompt
         system_message = self._organize_system_prompt()
 
-        # organize historic prompt messages
-        historic_messages = self._historic_prompt_messages
-
         # organize current assistant messages
         agent_scratchpad = self._agent_scratchpad
         if not agent_scratchpad:
@@ -57,6 +54,13 @@ class CotChatAgentRunner(CotAgentRunner):
         query_messages = UserPromptMessage(content=self._query)
 
         if assistant_messages:
+            # organize historic prompt messages
+            historic_messages = self._organize_historic_prompt_messages([
+                system_message,
+                query_messages,
+                *assistant_messages,
+                UserPromptMessage(content='continue')
+            ])            
             messages = [
                 system_message,
                 *historic_messages,
@@ -65,6 +69,8 @@ class CotChatAgentRunner(CotAgentRunner):
                 UserPromptMessage(content='continue')
             ]
         else:
+            # organize historic prompt messages
+            historic_messages = self._organize_historic_prompt_messages([system_message, query_messages])
             messages = [system_message, *historic_messages, query_messages]
 
         # join all messages

api/core/agent/cot_completion_agent_runner.py

@@ -19,11 +19,11 @@ class CotCompletionAgentRunner(CotAgentRunner):
         
         return system_prompt
 
-    def _organize_historic_prompt(self) -> str:
+    def _organize_historic_prompt(self, current_session_messages: list[PromptMessage] = None) -> str:
         """
         Organize historic prompt
         """
-        historic_prompt_messages = self._historic_prompt_messages
+        historic_prompt_messages = self._organize_historic_prompt_messages(current_session_messages)
         historic_prompt = ""
 
         for message in historic_prompt_messages:

api/core/agent/fc_agent_runner.py

@@ -17,6 +17,7 @@ from core.model_runtime.entities.message_entities import (
     ToolPromptMessage,
     UserPromptMessage,
 )
+from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform
 from core.tools.entities.tool_entities import ToolInvokeMeta
 from core.tools.tool_engine import ToolEngine
 from models.model import Message
@@ -24,21 +25,18 @@ from models.model import Message
 logger = logging.getLogger(__name__)
 
 class FunctionCallAgentRunner(BaseAgentRunner):
+
     def run(self, 
             message: Message, query: str, **kwargs: Any
     ) -> Generator[LLMResultChunk, None, None]:
         """
         Run FunctionCall agent application
         """
+        self.query = query
         app_generate_entity = self.application_generate_entity
 
         app_config = self.app_config
 
-        prompt_template = app_config.prompt_template.simple_prompt_template or ''
-        prompt_messages = self.history_prompt_messages
-        prompt_messages = self._init_system_message(prompt_template, prompt_messages)
-        prompt_messages = self._organize_user_query(query, prompt_messages)
-
         # convert tools into ModelRuntime Tool format
         tool_instances, prompt_messages_tools = self._init_prompt_tools()
 
@@ -81,6 +79,7 @@ class FunctionCallAgentRunner(BaseAgentRunner):
             )
 
             # recalc llm max tokens
+            prompt_messages = self._organize_prompt_messages()
             self.recalc_llm_max_tokens(self.model_config, prompt_messages)
             # invoke model
             chunks: Union[Generator[LLMResultChunk, None, None], LLMResult] = model_instance.invoke_llm(
@@ -203,7 +202,7 @@ class FunctionCallAgentRunner(BaseAgentRunner):
             else:
                 assistant_message.content = response
             
-            prompt_messages.append(assistant_message)
+            self._current_thoughts.append(assistant_message)
 
             # save thought
             self.save_agent_thought(
@@ -265,12 +264,14 @@ class FunctionCallAgentRunner(BaseAgentRunner):
                     }
                 
                 tool_responses.append(tool_response)
-                prompt_messages = self._organize_assistant_message(
-                    tool_call_id=tool_call_id,
-                    tool_call_name=tool_call_name,
-                    tool_response=tool_response['tool_response'],
-                    prompt_messages=prompt_messages,
-                )
+                if tool_response['tool_response'] is not None:
+                    self._current_thoughts.append(
+                        ToolPromptMessage(
+                            content=tool_response['tool_response'],
+                            tool_call_id=tool_call_id,
+                            name=tool_call_name,
+                        )
+                    ) 
 
             if len(tool_responses) > 0:
                 # save agent thought
@@ -300,8 +301,6 @@ class FunctionCallAgentRunner(BaseAgentRunner):
 
             iteration_step += 1
 
-            prompt_messages = self._clear_user_prompt_image_messages(prompt_messages)
-
         self.update_db_variables(self.variables_pool, self.db_variables_pool)
         # publish end event
         self.queue_manager.publish(QueueMessageEndEvent(llm_result=LLMResult(
@@ -393,24 +392,6 @@ class FunctionCallAgentRunner(BaseAgentRunner):
 
         return prompt_messages
     
-    def _organize_assistant_message(self, tool_call_id: str = None, tool_call_name: str = None, tool_response: str = None, 
-                                    prompt_messages: list[PromptMessage] = None) -> list[PromptMessage]:
-        """
-        Organize assistant message
-        """
-        prompt_messages = deepcopy(prompt_messages)
-
-        if tool_response is not None:
-            prompt_messages.append(
-                ToolPromptMessage(
-                    content=tool_response,
-                    tool_call_id=tool_call_id,
-                    name=tool_call_name,
-                )
-            )
-
-        return prompt_messages
-    
     def _clear_user_prompt_image_messages(self, prompt_messages: list[PromptMessage]) -> list[PromptMessage]:
         """
         As for now, gpt supports both fc and vision at the first iteration.
@@ -428,4 +409,26 @@ class FunctionCallAgentRunner(BaseAgentRunner):
                         for content in prompt_message.content 
                     ])
 
-        return prompt_messages
+        return prompt_messages
+
+    def _organize_prompt_messages(self):
+        prompt_template = self.app_config.prompt_template.simple_prompt_template or ''
+        self.history_prompt_messages = self._init_system_message(prompt_template, self.history_prompt_messages)
+        query_prompt_messages = self._organize_user_query(self.query, [])
+
+        self.history_prompt_messages = AgentHistoryPromptTransform(
+            model_config=self.model_config,
+            prompt_messages=[*query_prompt_messages, *self._current_thoughts],
+            history_messages=self.history_prompt_messages,
+            memory=self.memory
+        ).get_prompt()
+
+        prompt_messages = [
+            *self.history_prompt_messages,
+            *query_prompt_messages,
+            *self._current_thoughts
+        ]
+        if len(self._current_thoughts) != 0:
+            # clear messages after the first iteration
+            prompt_messages = self._clear_user_prompt_image_messages(prompt_messages)
+        return prompt_messages

api/core/app/apps/base_app_runner.py

@@ -1,6 +1,6 @@
 import time
 from collections.abc import Generator
-from typing import Optional, Union, cast
+from typing import Optional, Union
 
 from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@@ -16,11 +16,11 @@ from core.app.features.hosting_moderation.hosting_moderation import HostingModer
 from core.external_data_tool.external_data_fetch import ExternalDataFetch
 from core.file.file_obj import FileVar
 from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_manager import ModelInstance
 from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage
 from core.model_runtime.entities.model_entities import ModelPropertyKey
 from core.model_runtime.errors.invoke import InvokeBadRequestError
-from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.moderation.input_moderation import InputModeration
 from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
 from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
@@ -45,8 +45,11 @@ class AppRunner:
         :param query: query
         :return:
         """
-        model_type_instance = model_config.provider_model_bundle.model_type_instance
-        model_type_instance = cast(LargeLanguageModel, model_type_instance)
+        # Invoke model
+        model_instance = ModelInstance(
+            provider_model_bundle=model_config.provider_model_bundle,
+            model=model_config.model
+        )
 
         model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
 
@@ -73,9 +76,7 @@ class AppRunner:
             query=query
         )
 
-        prompt_tokens = model_type_instance.get_num_tokens(
-            model_config.model,
-            model_config.credentials,
+        prompt_tokens = model_instance.get_llm_num_tokens(
             prompt_messages
         )
 
@@ -89,8 +90,10 @@ class AppRunner:
     def recalc_llm_max_tokens(self, model_config: ModelConfigWithCredentialsEntity,
                               prompt_messages: list[PromptMessage]):
         # recalc max_tokens if sum(prompt_token +  max_tokens) over model token limit
-        model_type_instance = model_config.provider_model_bundle.model_type_instance
-        model_type_instance = cast(LargeLanguageModel, model_type_instance)
+        model_instance = ModelInstance(
+            provider_model_bundle=model_config.provider_model_bundle,
+            model=model_config.model
+        )
 
         model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
 
@@ -107,9 +110,7 @@ class AppRunner:
         if max_tokens is None:
             max_tokens = 0
 
-        prompt_tokens = model_type_instance.get_num_tokens(
-            model_config.model,
-            model_config.credentials,
+        prompt_tokens = model_instance.get_llm_num_tokens(
             prompt_messages
         )
 

api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -37,6 +37,7 @@ from core.app.entities.task_entities import (
 )
 from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
 from core.app.task_pipeline.message_cycle_manage import MessageCycleManage
+from core.model_manager import ModelInstance
 from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from core.model_runtime.entities.message_entities import (
     AssistantPromptMessage,
@@ -317,29 +318,30 @@ class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline, MessageCycleMan
         """
         model_config = self._model_config
         model = model_config.model
-        model_type_instance = model_config.provider_model_bundle.model_type_instance
-        model_type_instance = cast(LargeLanguageModel, model_type_instance)
+
+        model_instance = ModelInstance(
+            provider_model_bundle=model_config.provider_model_bundle,
+            model=model_config.model
+        )
 
         # calculate num tokens
         prompt_tokens = 0
         if event.stopped_by != QueueStopEvent.StopBy.ANNOTATION_REPLY:
-            prompt_tokens = model_type_instance.get_num_tokens(
-                model,
-                model_config.credentials,
+            prompt_tokens = model_instance.get_llm_num_tokens(
                 self._task_state.llm_result.prompt_messages
             )
 
         completion_tokens = 0
         if event.stopped_by == QueueStopEvent.StopBy.USER_MANUAL:
-            completion_tokens = model_type_instance.get_num_tokens(
-                model,
-                model_config.credentials,
+            completion_tokens = model_instance.get_llm_num_tokens(
                 [self._task_state.llm_result.message]
             )
 
         credentials = model_config.credentials
 
         # transform usage
+        model_type_instance = model_config.provider_model_bundle.model_type_instance
+        model_type_instance = cast(LargeLanguageModel, model_type_instance)
         self._task_state.llm_result.usage = model_type_instance._calc_response_usage(
             model,
             credentials,

api/core/entities/model_entities.py

@@ -16,6 +16,7 @@ class ModelStatus(Enum):
     NO_CONFIGURE = "no-configure"
     QUOTA_EXCEEDED = "quota-exceeded"
     NO_PERMISSION = "no-permission"
+    DISABLED = "disabled"
 
 
 class SimpleModelProviderEntity(BaseModel):
@@ -43,12 +44,19 @@ class SimpleModelProviderEntity(BaseModel):
         )
 
 
-class ModelWithProviderEntity(ProviderModel):
+class ProviderModelWithStatusEntity(ProviderModel):
+    """
+    Model class for model response.
+    """
+    status: ModelStatus
+    load_balancing_enabled: bool = False
+
+
+class ModelWithProviderEntity(ProviderModelWithStatusEntity):
     """
     Model with provider entity.
     """
     provider: SimpleModelProviderEntity
-    status: ModelStatus
 
 
 class DefaultModelProviderEntity(BaseModel):

api/core/entities/provider_configuration.py

@@ -1,6 +1,7 @@
 import datetime
 import json
 import logging
+from collections import defaultdict
 from collections.abc import Iterator
 from json import JSONDecodeError
 from typing import Optional
@@ -8,7 +9,12 @@ from typing import Optional
 from pydantic import BaseModel
 
 from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, SimpleModelProviderEntity
-from core.entities.provider_entities import CustomConfiguration, SystemConfiguration, SystemConfigurationStatus
+from core.entities.provider_entities import (
+    CustomConfiguration,
+    ModelSettings,
+    SystemConfiguration,
+    SystemConfigurationStatus,
+)
 from core.helper import encrypter
 from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
 from core.model_runtime.entities.model_entities import FetchFrom, ModelType
@@ -22,7 +28,14 @@ from core.model_runtime.model_providers import model_provider_factory
 from core.model_runtime.model_providers.__base.ai_model import AIModel
 from core.model_runtime.model_providers.__base.model_provider import ModelProvider
 from extensions.ext_database import db
-from models.provider import Provider, ProviderModel, ProviderType, TenantPreferredModelProvider
+from models.provider import (
+    LoadBalancingModelConfig,
+    Provider,
+    ProviderModel,
+    ProviderModelSetting,
+    ProviderType,
+    TenantPreferredModelProvider,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -39,6 +52,7 @@ class ProviderConfiguration(BaseModel):
     using_provider_type: ProviderType
     system_configuration: SystemConfiguration
     custom_configuration: CustomConfiguration
+    model_settings: list[ModelSettings]
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -62,6 +76,14 @@ class ProviderConfiguration(BaseModel):
         :param model: model name
         :return:
         """
+        if self.model_settings:
+            # check if model is disabled by admin
+            for model_setting in self.model_settings:
+                if (model_setting.model_type == model_type
+                        and model_setting.model == model):
+                    if not model_setting.enabled:
+                        raise ValueError(f'Model {model} is disabled.')
+
         if self.using_provider_type == ProviderType.SYSTEM:
             restrict_models = []
             for quota_configuration in self.system_configuration.quota_configurations:
@@ -80,15 +102,17 @@ class ProviderConfiguration(BaseModel):
 
             return copy_credentials
         else:
+            credentials = None
             if self.custom_configuration.models:
                 for model_configuration in self.custom_configuration.models:
                     if model_configuration.model_type == model_type and model_configuration.model == model:
-                        return model_configuration.credentials
+                        credentials = model_configuration.credentials
+                        break
 
             if self.custom_configuration.provider:
-                return self.custom_configuration.provider.credentials
-            else:
-                return None
+                credentials = self.custom_configuration.provider.credentials
+
+            return credentials
 
     def get_system_configuration_status(self) -> SystemConfigurationStatus:
         """
@@ -130,7 +154,7 @@ class ProviderConfiguration(BaseModel):
             return credentials
 
         # Obfuscate credentials
-        return self._obfuscated_credentials(
+        return self.obfuscated_credentials(
             credentials=credentials,
             credential_form_schemas=self.provider.provider_credential_schema.credential_form_schemas
             if self.provider.provider_credential_schema else []
@@ -151,7 +175,7 @@ class ProviderConfiguration(BaseModel):
         ).first()
 
         # Get provider credential secret variables
-        provider_credential_secret_variables = self._extract_secret_variables(
+        provider_credential_secret_variables = self.extract_secret_variables(
             self.provider.provider_credential_schema.credential_form_schemas
             if self.provider.provider_credential_schema else []
         )
@@ -274,7 +298,7 @@ class ProviderConfiguration(BaseModel):
                     return credentials
 
                 # Obfuscate credentials
-                return self._obfuscated_credentials(
+                return self.obfuscated_credentials(
                     credentials=credentials,
                     credential_form_schemas=self.provider.model_credential_schema.credential_form_schemas
                     if self.provider.model_credential_schema else []
@@ -302,7 +326,7 @@ class ProviderConfiguration(BaseModel):
         ).first()
 
         # Get provider credential secret variables
-        provider_credential_secret_variables = self._extract_secret_variables(
+        provider_credential_secret_variables = self.extract_secret_variables(
             self.provider.model_credential_schema.credential_form_schemas
             if self.provider.model_credential_schema else []
         )
@@ -402,6 +426,160 @@ class ProviderConfiguration(BaseModel):
 
             provider_model_credentials_cache.delete()
 
+    def enable_model(self, model_type: ModelType, model: str) -> ProviderModelSetting:
+        """
+        Enable model.
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        model_setting = db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == self.tenant_id,
+            ProviderModelSetting.provider_name == self.provider.provider,
+            ProviderModelSetting.model_type == model_type.to_origin_model_type(),
+            ProviderModelSetting.model_name == model
+        ).first()
+
+        if model_setting:
+            model_setting.enabled = True
+            model_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+            db.session.commit()
+        else:
+            model_setting = ProviderModelSetting(
+                tenant_id=self.tenant_id,
+                provider_name=self.provider.provider,
+                model_type=model_type.to_origin_model_type(),
+                model_name=model,
+                enabled=True
+            )
+            db.session.add(model_setting)
+            db.session.commit()
+
+        return model_setting
+
+    def disable_model(self, model_type: ModelType, model: str) -> ProviderModelSetting:
+        """
+        Disable model.
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        model_setting = db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == self.tenant_id,
+            ProviderModelSetting.provider_name == self.provider.provider,
+            ProviderModelSetting.model_type == model_type.to_origin_model_type(),
+            ProviderModelSetting.model_name == model
+        ).first()
+
+        if model_setting:
+            model_setting.enabled = False
+            model_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+            db.session.commit()
+        else:
+            model_setting = ProviderModelSetting(
+                tenant_id=self.tenant_id,
+                provider_name=self.provider.provider,
+                model_type=model_type.to_origin_model_type(),
+                model_name=model,
+                enabled=False
+            )
+            db.session.add(model_setting)
+            db.session.commit()
+
+        return model_setting
+
+    def get_provider_model_setting(self, model_type: ModelType, model: str) -> Optional[ProviderModelSetting]:
+        """
+        Get provider model setting.
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        return db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == self.tenant_id,
+            ProviderModelSetting.provider_name == self.provider.provider,
+            ProviderModelSetting.model_type == model_type.to_origin_model_type(),
+            ProviderModelSetting.model_name == model
+        ).first()
+
+    def enable_model_load_balancing(self, model_type: ModelType, model: str) -> ProviderModelSetting:
+        """
+        Enable model load balancing.
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        load_balancing_config_count = db.session.query(LoadBalancingModelConfig) \
+            .filter(
+            LoadBalancingModelConfig.tenant_id == self.tenant_id,
+            LoadBalancingModelConfig.provider_name == self.provider.provider,
+            LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(),
+            LoadBalancingModelConfig.model_name == model
+        ).count()
+
+        if load_balancing_config_count <= 1:
+            raise ValueError('Model load balancing configuration must be more than 1.')
+
+        model_setting = db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == self.tenant_id,
+            ProviderModelSetting.provider_name == self.provider.provider,
+            ProviderModelSetting.model_type == model_type.to_origin_model_type(),
+            ProviderModelSetting.model_name == model
+        ).first()
+
+        if model_setting:
+            model_setting.load_balancing_enabled = True
+            model_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+            db.session.commit()
+        else:
+            model_setting = ProviderModelSetting(
+                tenant_id=self.tenant_id,
+                provider_name=self.provider.provider,
+                model_type=model_type.to_origin_model_type(),
+                model_name=model,
+                load_balancing_enabled=True
+            )
+            db.session.add(model_setting)
+            db.session.commit()
+
+        return model_setting
+
+    def disable_model_load_balancing(self, model_type: ModelType, model: str) -> ProviderModelSetting:
+        """
+        Disable model load balancing.
+        :param model_type: model type
+        :param model: model name
+        :return:
+        """
+        model_setting = db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == self.tenant_id,
+            ProviderModelSetting.provider_name == self.provider.provider,
+            ProviderModelSetting.model_type == model_type.to_origin_model_type(),
+            ProviderModelSetting.model_name == model
+        ).first()
+
+        if model_setting:
+            model_setting.load_balancing_enabled = False
+            model_setting.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
+            db.session.commit()
+        else:
+            model_setting = ProviderModelSetting(
+                tenant_id=self.tenant_id,
+                provider_name=self.provider.provider,
+                model_type=model_type.to_origin_model_type(),
+                model_name=model,
+                load_balancing_enabled=False
+            )
+            db.session.add(model_setting)
+            db.session.commit()
+
+        return model_setting
+
     def get_provider_instance(self) -> ModelProvider:
         """
         Get provider instance.
@@ -453,7 +631,7 @@ class ProviderConfiguration(BaseModel):
 
         db.session.commit()
 
-    def _extract_secret_variables(self, credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
+    def extract_secret_variables(self, credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
         """
         Extract secret input form variables.
 
@@ -467,7 +645,7 @@ class ProviderConfiguration(BaseModel):
 
         return secret_input_form_variables
 
-    def _obfuscated_credentials(self, credentials: dict, credential_form_schemas: list[CredentialFormSchema]) -> dict:
+    def obfuscated_credentials(self, credentials: dict, credential_form_schemas: list[CredentialFormSchema]) -> dict:
         """
         Obfuscated credentials.
 
@@ -476,7 +654,7 @@ class ProviderConfiguration(BaseModel):
         :return:
         """
         # Get provider credential secret variables
-        credential_secret_variables = self._extract_secret_variables(
+        credential_secret_variables = self.extract_secret_variables(
             credential_form_schemas
         )
 
@@ -522,15 +700,22 @@ class ProviderConfiguration(BaseModel):
         else:
             model_types = provider_instance.get_provider_schema().supported_model_types
 
+        # Group model settings by model type and model
+        model_setting_map = defaultdict(dict)
+        for model_setting in self.model_settings:
+            model_setting_map[model_setting.model_type][model_setting.model] = model_setting
+
         if self.using_provider_type == ProviderType.SYSTEM:
             provider_models = self._get_system_provider_models(
                 model_types=model_types,
-                provider_instance=provider_instance
+                provider_instance=provider_instance,
+                model_setting_map=model_setting_map
             )
         else:
             provider_models = self._get_custom_provider_models(
                 model_types=model_types,
-                provider_instance=provider_instance
+                provider_instance=provider_instance,
+                model_setting_map=model_setting_map
             )
 
         if only_active:
@@ -541,18 +726,27 @@ class ProviderConfiguration(BaseModel):
 
     def _get_system_provider_models(self,
                                     model_types: list[ModelType],
-                                    provider_instance: ModelProvider) -> list[ModelWithProviderEntity]:
+                                    provider_instance: ModelProvider,
+                                    model_setting_map: dict[ModelType, dict[str, ModelSettings]]) \
+            -> list[ModelWithProviderEntity]:
         """
         Get system provider models.
 
         :param model_types: model types
         :param provider_instance: provider instance
+        :param model_setting_map: model setting map
         :return:
         """
         provider_models = []
         for model_type in model_types:
-            provider_models.extend(
-                [
+            for m in provider_instance.models(model_type):
+                status = ModelStatus.ACTIVE
+                if m.model_type in model_setting_map and m.model in model_setting_map[m.model_type]:
+                    model_setting = model_setting_map[m.model_type][m.model]
+                    if model_setting.enabled is False:
+                        status = ModelStatus.DISABLED
+
+                provider_models.append(
                     ModelWithProviderEntity(
                         model=m.model,
                         label=m.label,
@@ -562,11 +756,9 @@ class ProviderConfiguration(BaseModel):
                         model_properties=m.model_properties,
                         deprecated=m.deprecated,
                         provider=SimpleModelProviderEntity(self.provider),
-                        status=ModelStatus.ACTIVE
+                        status=status
                     )
-                    for m in provider_instance.models(model_type)
-                ]
-            )
+                )
 
         if self.provider.provider not in original_provider_configurate_methods:
             original_provider_configurate_methods[self.provider.provider] = []
@@ -586,7 +778,8 @@ class ProviderConfiguration(BaseModel):
                 break
 
             if should_use_custom_model:
-                if original_provider_configurate_methods[self.provider.provider] == [ConfigurateMethod.CUSTOMIZABLE_MODEL]:
+                if original_provider_configurate_methods[self.provider.provider] == [
+                    ConfigurateMethod.CUSTOMIZABLE_MODEL]:
                     # only customizable model
                     for restrict_model in restrict_models:
                         copy_credentials = self.system_configuration.credentials.copy()
@@ -611,6 +804,13 @@ class ProviderConfiguration(BaseModel):
                         if custom_model_schema.model_type not in model_types:
                             continue
 
+                        status = ModelStatus.ACTIVE
+                        if (custom_model_schema.model_type in model_setting_map
+                                and custom_model_schema.model in model_setting_map[custom_model_schema.model_type]):
+                            model_setting = model_setting_map[custom_model_schema.model_type][custom_model_schema.model]
+                            if model_setting.enabled is False:
+                                status = ModelStatus.DISABLED
+
                         provider_models.append(
                             ModelWithProviderEntity(
                                 model=custom_model_schema.model,
@@ -621,7 +821,7 @@ class ProviderConfiguration(BaseModel):
                                 model_properties=custom_model_schema.model_properties,
                                 deprecated=custom_model_schema.deprecated,
                                 provider=SimpleModelProviderEntity(self.provider),
-                                status=ModelStatus.ACTIVE
+                                status=status
                             )
                         )
 
@@ -632,16 +832,20 @@ class ProviderConfiguration(BaseModel):
                     m.status = ModelStatus.NO_PERMISSION
                 elif not quota_configuration.is_valid:
                     m.status = ModelStatus.QUOTA_EXCEEDED
+
         return provider_models
 
     def _get_custom_provider_models(self,
                                     model_types: list[ModelType],
-                                    provider_instance: ModelProvider) -> list[ModelWithProviderEntity]:
+                                    provider_instance: ModelProvider,
+                                    model_setting_map: dict[ModelType, dict[str, ModelSettings]]) \
+            -> list[ModelWithProviderEntity]:
         """
         Get custom provider models.
 
         :param model_types: model types
         :param provider_instance: provider instance
+        :param model_setting_map: model setting map
         :return:
         """
         provider_models = []
@@ -656,6 +860,16 @@ class ProviderConfiguration(BaseModel):
 
             models = provider_instance.models(model_type)
             for m in models:
+                status = ModelStatus.ACTIVE if credentials else ModelStatus.NO_CONFIGURE
+                load_balancing_enabled = False
+                if m.model_type in model_setting_map and m.model in model_setting_map[m.model_type]:
+                    model_setting = model_setting_map[m.model_type][m.model]
+                    if model_setting.enabled is False:
+                        status = ModelStatus.DISABLED
+
+                    if len(model_setting.load_balancing_configs) > 1:
+                        load_balancing_enabled = True
+
                 provider_models.append(
                     ModelWithProviderEntity(
                         model=m.model,
@@ -666,7 +880,8 @@ class ProviderConfiguration(BaseModel):
                         model_properties=m.model_properties,
                         deprecated=m.deprecated,
                         provider=SimpleModelProviderEntity(self.provider),
-                        status=ModelStatus.ACTIVE if credentials else ModelStatus.NO_CONFIGURE
+                        status=status,
+                        load_balancing_enabled=load_balancing_enabled
                     )
                 )
 
@@ -690,6 +905,17 @@ class ProviderConfiguration(BaseModel):
             if not custom_model_schema:
                 continue
 
+            status = ModelStatus.ACTIVE
+            load_balancing_enabled = False
+            if (custom_model_schema.model_type in model_setting_map
+                    and custom_model_schema.model in model_setting_map[custom_model_schema.model_type]):
+                model_setting = model_setting_map[custom_model_schema.model_type][custom_model_schema.model]
+                if model_setting.enabled is False:
+                    status = ModelStatus.DISABLED
+
+                if len(model_setting.load_balancing_configs) > 1:
+                    load_balancing_enabled = True
+
             provider_models.append(
                 ModelWithProviderEntity(
                     model=custom_model_schema.model,
@@ -700,7 +926,8 @@ class ProviderConfiguration(BaseModel):
                     model_properties=custom_model_schema.model_properties,
                     deprecated=custom_model_schema.deprecated,
                     provider=SimpleModelProviderEntity(self.provider),
-                    status=ModelStatus.ACTIVE
+                    status=status,
+                    load_balancing_enabled=load_balancing_enabled
                 )
             )
 

api/core/entities/provider_entities.py

@@ -72,3 +72,22 @@ class CustomConfiguration(BaseModel):
     """
     provider: Optional[CustomProviderConfiguration] = None
     models: list[CustomModelConfiguration] = []
+
+
+class ModelLoadBalancingConfiguration(BaseModel):
+    """
+    Class for model load balancing configuration.
+    """
+    id: str
+    name: str
+    credentials: dict
+
+
+class ModelSettings(BaseModel):
+    """
+    Model class for model settings.
+    """
+    model: str
+    model_type: ModelType
+    enabled: bool = True
+    load_balancing_configs: list[ModelLoadBalancingConfiguration] = []

api/core/extension/extensible.py

@@ -7,7 +7,7 @@ from typing import Any, Optional
 
 from pydantic import BaseModel
 
-from core.utils.position_helper import sort_to_dict_by_position_map
+from core.helper.position_helper import sort_to_dict_by_position_map
 
 
 class ExtensionModule(enum.Enum):

api/core/helper/model_provider_cache.py

@@ -9,6 +9,7 @@ from extensions.ext_redis import redis_client
 class ProviderCredentialsCacheType(Enum):
     PROVIDER = "provider"
     MODEL = "provider_model"
+    LOAD_BALANCING_MODEL = "load_balancing_provider_model"
 
 
 class ProviderCredentialsCache:

api/core/indexing_runner.py

@@ -12,7 +12,6 @@ from flask import Flask, current_app
 from flask_login import current_user
 from sqlalchemy.orm.exc import ObjectDeletedError
 
-from core.docstore.dataset_docstore import DatasetDocumentStore
 from core.errors.error import ProviderTokenNotInitError
 from core.llm_generator.llm_generator import LLMGenerator
 from core.model_manager import ModelInstance, ModelManager
@@ -20,12 +19,16 @@ from core.model_runtime.entities.model_entities import ModelType, PriceType
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from core.rag.datasource.keyword.keyword_factory import Keyword
+from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from core.rag.models.document import Document
-from core.splitter.fixed_text_splitter import EnhanceRecursiveCharacterTextSplitter, FixedRecursiveCharacterTextSplitter
-from core.splitter.text_splitter import TextSplitter
+from core.rag.splitter.fixed_text_splitter import (
+    EnhanceRecursiveCharacterTextSplitter,
+    FixedRecursiveCharacterTextSplitter,
+)
+from core.rag.splitter.text_splitter import TextSplitter
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from extensions.ext_storage import storage
@@ -283,11 +286,7 @@ class IndexingRunner:
                 if len(preview_texts) < 5:
                     preview_texts.append(document.page_content)
                 if indexing_technique == 'high_quality' or embedding_model_instance:
-                    embedding_model_type_instance = embedding_model_instance.model_type_instance
-                    embedding_model_type_instance = cast(TextEmbeddingModel, embedding_model_type_instance)
-                    tokens += embedding_model_type_instance.get_num_tokens(
-                        model=embedding_model_instance.model,
-                        credentials=embedding_model_instance.credentials,
+                    tokens += embedding_model_instance.get_text_embedding_num_tokens(
                         texts=[self.filter_string(document.page_content)]
                     )
 
@@ -655,10 +654,6 @@ class IndexingRunner:
         tokens = 0
         chunk_size = 10
 
-        embedding_model_type_instance = None
-        if embedding_model_instance:
-            embedding_model_type_instance = embedding_model_instance.model_type_instance
-            embedding_model_type_instance = cast(TextEmbeddingModel, embedding_model_type_instance)
         # create keyword index
         create_keyword_thread = threading.Thread(target=self._process_keyword_index,
                                                  args=(current_app._get_current_object(),
@@ -671,8 +666,7 @@ class IndexingRunner:
                     chunk_documents = documents[i:i + chunk_size]
                     futures.append(executor.submit(self._process_chunk, current_app._get_current_object(), index_processor,
                                                    chunk_documents, dataset,
-                                                   dataset_document, embedding_model_instance,
-                                                   embedding_model_type_instance))
+                                                   dataset_document, embedding_model_instance))
 
                 for future in futures:
                     tokens += future.result()
@@ -713,7 +707,7 @@ class IndexingRunner:
                 db.session.commit()
 
     def _process_chunk(self, flask_app, index_processor, chunk_documents, dataset, dataset_document,
-                       embedding_model_instance, embedding_model_type_instance):
+                       embedding_model_instance):
         with flask_app.app_context():
             # check document is paused
             self._check_document_paused_status(dataset_document.id)
@@ -721,9 +715,7 @@ class IndexingRunner:
             tokens = 0
             if dataset.indexing_technique == 'high_quality' or embedding_model_type_instance:
                 tokens += sum(
-                    embedding_model_type_instance.get_num_tokens(
-                        embedding_model_instance.model,
-                        embedding_model_instance.credentials,
+                    embedding_model_instance.get_text_embedding_num_tokens(
                         [document.page_content]
                     )
                     for document in chunk_documents

api/core/memory/token_buffer_memory.py

@@ -9,8 +9,6 @@ from core.model_runtime.entities.message_entities import (
     TextPromptMessageContent,
     UserPromptMessage,
 )
-from core.model_runtime.entities.model_entities import ModelType
-from core.model_runtime.model_providers import model_provider_factory
 from extensions.ext_database import db
 from models.model import AppMode, Conversation, Message
 
@@ -78,12 +76,7 @@ class TokenBufferMemory:
             return []
 
         # prune the chat message if it exceeds the max token limit
-        provider_instance = model_provider_factory.get_provider_instance(self.model_instance.provider)
-        model_type_instance = provider_instance.get_model_instance(ModelType.LLM)
-
-        curr_message_tokens = model_type_instance.get_num_tokens(
-            self.model_instance.model,
-            self.model_instance.credentials,
+        curr_message_tokens = self.model_instance.get_llm_num_tokens(
             prompt_messages
         )
 
@@ -91,9 +84,7 @@ class TokenBufferMemory:
             pruned_memory = []
             while curr_message_tokens > max_token_limit and prompt_messages:
                 pruned_memory.append(prompt_messages.pop(0))
-                curr_message_tokens = model_type_instance.get_num_tokens(
-                    self.model_instance.model,
-                    self.model_instance.credentials,
+                curr_message_tokens = self.model_instance.get_llm_num_tokens(
                     prompt_messages
                 )
 

api/core/model_manager.py

@@ -1,7 +1,10 @@
+import logging
+import os
 from collections.abc import Generator
 from typing import IO, Optional, Union, cast
 
-from core.entities.provider_configuration import ProviderModelBundle
+from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
+from core.entities.provider_entities import ModelLoadBalancingConfiguration
 from core.errors.error import ProviderTokenNotInitError
 from core.model_runtime.callbacks.base_callback import Callback
 from core.model_runtime.entities.llm_entities import LLMResult
@@ -9,6 +12,7 @@ from core.model_runtime.entities.message_entities import PromptMessage, PromptMe
 from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.entities.rerank_entities import RerankResult
 from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
+from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeConnectionError, InvokeRateLimitError
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.model_runtime.model_providers.__base.moderation_model import ModerationModel
 from core.model_runtime.model_providers.__base.rerank_model import RerankModel
@@ -16,6 +20,10 @@ from core.model_runtime.model_providers.__base.speech2text_model import Speech2T
 from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from core.model_runtime.model_providers.__base.tts_model import TTSModel
 from core.provider_manager import ProviderManager
+from extensions.ext_redis import redis_client
+from models.provider import ProviderType
+
+logger = logging.getLogger(__name__)
 
 
 class ModelInstance:
@@ -29,6 +37,12 @@ class ModelInstance:
         self.provider = provider_model_bundle.configuration.provider.provider
         self.credentials = self._fetch_credentials_from_bundle(provider_model_bundle, model)
         self.model_type_instance = self.provider_model_bundle.model_type_instance
+        self.load_balancing_manager = self._get_load_balancing_manager(
+            configuration=provider_model_bundle.configuration,
+            model_type=provider_model_bundle.model_type_instance.model_type,
+            model=model,
+            credentials=self.credentials
+        )
 
     def _fetch_credentials_from_bundle(self, provider_model_bundle: ProviderModelBundle, model: str) -> dict:
         """
@@ -37,8 +51,10 @@ class ModelInstance:
         :param model: model name
         :return:
         """
-        credentials = provider_model_bundle.configuration.get_current_credentials(
-            model_type=provider_model_bundle.model_type_instance.model_type,
+        configuration = provider_model_bundle.configuration
+        model_type = provider_model_bundle.model_type_instance.model_type
+        credentials = configuration.get_current_credentials(
+            model_type=model_type,
             model=model
         )
 
@@ -47,6 +63,43 @@ class ModelInstance:
 
         return credentials
 
+    def _get_load_balancing_manager(self, configuration: ProviderConfiguration,
+                                    model_type: ModelType,
+                                    model: str,
+                                    credentials: dict) -> Optional["LBModelManager"]:
+        """
+        Get load balancing model credentials
+        :param configuration: provider configuration
+        :param model_type: model type
+        :param model: model name
+        :param credentials: model credentials
+        :return:
+        """
+        if configuration.model_settings and configuration.using_provider_type == ProviderType.CUSTOM:
+            current_model_setting = None
+            # check if model is disabled by admin
+            for model_setting in configuration.model_settings:
+                if (model_setting.model_type == model_type
+                        and model_setting.model == model):
+                    current_model_setting = model_setting
+                    break
+
+            # check if load balancing is enabled
+            if current_model_setting and current_model_setting.load_balancing_configs:
+                # use load balancing proxy to choose credentials
+                lb_model_manager = LBModelManager(
+                    tenant_id=configuration.tenant_id,
+                    provider=configuration.provider.provider,
+                    model_type=model_type,
+                    model=model,
+                    load_balancing_configs=current_model_setting.load_balancing_configs,
+                    managed_credentials=credentials if configuration.custom_configuration.provider else None
+                )
+
+                return lb_model_manager
+
+        return None
+
     def invoke_llm(self, prompt_messages: list[PromptMessage], model_parameters: Optional[dict] = None,
                    tools: Optional[list[PromptMessageTool]] = None, stop: Optional[list[str]] = None,
                    stream: bool = True, user: Optional[str] = None, callbacks: list[Callback] = None) \
@@ -67,7 +120,8 @@ class ModelInstance:
             raise Exception("Model type instance is not LargeLanguageModel")
 
         self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             prompt_messages=prompt_messages,
@@ -79,6 +133,27 @@ class ModelInstance:
             callbacks=callbacks
         )
 
+    def get_llm_num_tokens(self, prompt_messages: list[PromptMessage],
+                           tools: Optional[list[PromptMessageTool]] = None) -> int:
+        """
+        Get number of tokens for llm
+
+        :param prompt_messages: prompt messages
+        :param tools: tools for tool calling
+        :return:
+        """
+        if not isinstance(self.model_type_instance, LargeLanguageModel):
+            raise Exception("Model type instance is not LargeLanguageModel")
+
+        self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance)
+        return self._round_robin_invoke(
+            function=self.model_type_instance.get_num_tokens,
+            model=self.model,
+            credentials=self.credentials,
+            prompt_messages=prompt_messages,
+            tools=tools
+        )
+
     def invoke_text_embedding(self, texts: list[str], user: Optional[str] = None) \
             -> TextEmbeddingResult:
         """
@@ -92,13 +167,32 @@ class ModelInstance:
             raise Exception("Model type instance is not TextEmbeddingModel")
 
         self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             texts=texts,
             user=user
         )
 
+    def get_text_embedding_num_tokens(self, texts: list[str]) -> int:
+        """
+        Get number of tokens for text embedding
+
+        :param texts: texts to embed
+        :return:
+        """
+        if not isinstance(self.model_type_instance, TextEmbeddingModel):
+            raise Exception("Model type instance is not TextEmbeddingModel")
+
+        self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance)
+        return self._round_robin_invoke(
+            function=self.model_type_instance.get_num_tokens,
+            model=self.model,
+            credentials=self.credentials,
+            texts=texts
+        )
+
     def invoke_rerank(self, query: str, docs: list[str], score_threshold: Optional[float] = None,
                       top_n: Optional[int] = None,
                       user: Optional[str] = None) \
@@ -117,7 +211,8 @@ class ModelInstance:
             raise Exception("Model type instance is not RerankModel")
 
         self.model_type_instance = cast(RerankModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             query=query,
@@ -140,7 +235,8 @@ class ModelInstance:
             raise Exception("Model type instance is not ModerationModel")
 
         self.model_type_instance = cast(ModerationModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             text=text,
@@ -160,7 +256,8 @@ class ModelInstance:
             raise Exception("Model type instance is not Speech2TextModel")
 
         self.model_type_instance = cast(Speech2TextModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             file=file,
@@ -183,7 +280,8 @@ class ModelInstance:
             raise Exception("Model type instance is not TTSModel")
 
         self.model_type_instance = cast(TTSModel, self.model_type_instance)
-        return self.model_type_instance.invoke(
+        return self._round_robin_invoke(
+            function=self.model_type_instance.invoke,
             model=self.model,
             credentials=self.credentials,
             content_text=content_text,
@@ -193,6 +291,43 @@ class ModelInstance:
             streaming=streaming
         )
 
+    def _round_robin_invoke(self, function: callable, *args, **kwargs):
+        """
+        Round-robin invoke
+        :param function: function to invoke
+        :param args: function args
+        :param kwargs: function kwargs
+        :return:
+        """
+        if not self.load_balancing_manager:
+            return function(*args, **kwargs)
+
+        last_exception = None
+        while True:
+            lb_config = self.load_balancing_manager.fetch_next()
+            if not lb_config:
+                if not last_exception:
+                    raise ProviderTokenNotInitError("Model credentials is not initialized.")
+                else:
+                    raise last_exception
+
+            try:
+                if 'credentials' in kwargs:
+                    del kwargs['credentials']
+                return function(*args, **kwargs, credentials=lb_config.credentials)
+            except InvokeRateLimitError as e:
+                # expire in 60 seconds
+                self.load_balancing_manager.cooldown(lb_config, expire=60)
+                last_exception = e
+                continue
+            except (InvokeAuthorizationError, InvokeConnectionError) as e:
+                # expire in 10 seconds
+                self.load_balancing_manager.cooldown(lb_config, expire=10)
+                last_exception = e
+                continue
+            except Exception as e:
+                raise e
+
     def get_tts_voices(self, language: str) -> list:
         """
         Invoke large language tts model voices
@@ -226,6 +361,7 @@ class ModelManager:
         """
         if not provider:
             return self.get_default_model_instance(tenant_id, model_type)
+
         provider_model_bundle = self._provider_manager.get_provider_model_bundle(
             tenant_id=tenant_id,
             provider=provider,
@@ -255,3 +391,141 @@ class ModelManager:
             model_type=model_type,
             model=default_model_entity.model
         )
+
+
+class LBModelManager:
+    def __init__(self, tenant_id: str,
+                 provider: str,
+                 model_type: ModelType,
+                 model: str,
+                 load_balancing_configs: list[ModelLoadBalancingConfiguration],
+                 managed_credentials: Optional[dict] = None) -> None:
+        """
+        Load balancing model manager
+        :param load_balancing_configs: all load balancing configurations
+        :param managed_credentials: credentials if load balancing configuration name is __inherit__
+        """
+        self._tenant_id = tenant_id
+        self._provider = provider
+        self._model_type = model_type
+        self._model = model
+        self._load_balancing_configs = load_balancing_configs
+
+        for load_balancing_config in self._load_balancing_configs:
+            if load_balancing_config.name == "__inherit__":
+                if not managed_credentials:
+                    # remove __inherit__ if managed credentials is not provided
+                    self._load_balancing_configs.remove(load_balancing_config)
+                else:
+                    load_balancing_config.credentials = managed_credentials
+
+    def fetch_next(self) -> Optional[ModelLoadBalancingConfiguration]:
+        """
+        Get next model load balancing config
+        Strategy: Round Robin
+        :return:
+        """
+        cache_key = "model_lb_index:{}:{}:{}:{}".format(
+            self._tenant_id,
+            self._provider,
+            self._model_type.value,
+            self._model
+        )
+
+        cooldown_load_balancing_configs = []
+        max_index = len(self._load_balancing_configs)
+
+        while True:
+            current_index = redis_client.incr(cache_key)
+            if current_index >= 10000000:
+                current_index = 1
+                redis_client.set(cache_key, current_index)
+
+            redis_client.expire(cache_key, 3600)
+            if current_index > max_index:
+                current_index = current_index % max_index
+
+            real_index = current_index - 1
+            if real_index > max_index:
+                real_index = 0
+
+            config = self._load_balancing_configs[real_index]
+
+            if self.in_cooldown(config):
+                cooldown_load_balancing_configs.append(config)
+                if len(cooldown_load_balancing_configs) >= len(self._load_balancing_configs):
+                    # all configs are in cooldown
+                    return None
+
+                continue
+
+            if bool(os.environ.get("DEBUG", 'False').lower() == 'true'):
+                logger.info(f"Model LB\nid: {config.id}\nname:{config.name}\n"
+                            f"tenant_id: {self._tenant_id}\nprovider: {self._provider}\n"
+                            f"model_type: {self._model_type.value}\nmodel: {self._model}")
+
+            return config
+
+        return None
+
+    def cooldown(self, config: ModelLoadBalancingConfiguration, expire: int = 60) -> None:
+        """
+        Cooldown model load balancing config
+        :param config: model load balancing config
+        :param expire: cooldown time
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            self._tenant_id,
+            self._provider,
+            self._model_type.value,
+            self._model,
+            config.id
+        )
+
+        redis_client.setex(cooldown_cache_key, expire, 'true')
+
+    def in_cooldown(self, config: ModelLoadBalancingConfiguration) -> bool:
+        """
+        Check if model load balancing config is in cooldown
+        :param config: model load balancing config
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            self._tenant_id,
+            self._provider,
+            self._model_type.value,
+            self._model,
+            config.id
+        )
+
+        return redis_client.exists(cooldown_cache_key)
+
+    @classmethod
+    def get_config_in_cooldown_and_ttl(cls, tenant_id: str,
+                                       provider: str,
+                                       model_type: ModelType,
+                                       model: str,
+                                       config_id: str) -> tuple[bool, int]:
+        """
+        Get model load balancing config is in cooldown and ttl
+        :param tenant_id: workspace id
+        :param provider: provider name
+        :param model_type: model type
+        :param model: model name
+        :param config_id: model load balancing config id
+        :return:
+        """
+        cooldown_cache_key = "model_lb_index:cooldown:{}:{}:{}:{}:{}".format(
+            tenant_id,
+            provider,
+            model_type.value,
+            model,
+            config_id
+        )
+
+        ttl = redis_client.ttl(cooldown_cache_key)
+        if ttl == -2:
+            return False, 0
+
+        return True, ttl

api/core/model_runtime/README.md

@@ -20,7 +20,7 @@ This module provides the interface for invoking and authenticating various model
 
   ![image-20231210143654461](./docs/en_US/images/index/image-20231210143654461.png)
 
-  Displays a list of all supported providers, including provider names, icons, supported model types list, predefined model list, configuration method, and credentials form rules, etc. For detailed rule design, see: [Schema](./schema.md).
+  Displays a list of all supported providers, including provider names, icons, supported model types list, predefined model list, configuration method, and credentials form rules, etc. For detailed rule design, see: [Schema](./docs/en_US/schema.md).
 
 - Selectable model list display
 

api/core/model_runtime/model_providers/__base/ai_model.py

@@ -3,6 +3,7 @@ import os
 from abc import ABC, abstractmethod
 from typing import Optional
 
+from core.helper.position_helper import get_position_map, sort_by_position_map
 from core.model_runtime.entities.common_entities import I18nObject
 from core.model_runtime.entities.defaults import PARAMETER_RULE_TEMPLATE
 from core.model_runtime.entities.model_entities import (
@@ -17,7 +18,6 @@ from core.model_runtime.entities.model_entities import (
 from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
 from core.model_runtime.model_providers.__base.tokenizers.gpt2_tokenzier import GPT2Tokenizer
 from core.tools.utils.yaml_utils import load_yaml_file
-from core.utils.position_helper import get_position_map, sort_by_position_map
 
 
 class AIModel(ABC):

api/core/model_runtime/model_providers/__base/model_provider.py

@@ -1,11 +1,11 @@
 import os
 from abc import ABC, abstractmethod
 
+from core.helper.module_import_helper import get_subclasses_from_module, import_module_from_source
 from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from core.model_runtime.entities.provider_entities import ProviderEntity
 from core.model_runtime.model_providers.__base.ai_model import AIModel
 from core.tools.utils.yaml_utils import load_yaml_file
-from core.utils.module_import_helper import get_subclasses_from_module, import_module_from_source
 
 
 class ModelProvider(ABC):

api/core/model_runtime/model_providers/_position.yaml

@@ -4,6 +4,7 @@
 - google
 - vertex_ai
 - nvidia
+- nvidia_nim
 - cohere
 - bedrock
 - togetherai

api/core/model_runtime/model_providers/baichuan/llm/baichuan2-53b.yaml

@@ -6,7 +6,7 @@ features:
   - agent-thought
 model_properties:
   mode: chat
-  context_size: 4000
+  context_size: 32000
 parameter_rules:
   - name: temperature
     use_template: temperature

api/core/model_runtime/model_providers/baichuan/llm/baichuan2-turbo.yaml

@@ -6,7 +6,7 @@ features:
   - agent-thought
 model_properties:
   mode: chat
-  context_size: 192000
+  context_size: 32000
 parameter_rules:
   - name: temperature
     use_template: temperature

api/core/model_runtime/model_providers/baichuan/llm/baichuan3-turbo-128k.yaml

@@ -0,0 +1,45 @@
+model: baichuan3-turbo-128k
+label:
+  en_US: Baichuan3-Turbo-128k
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 128000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+  - name: top_p
+    use_template: top_p
+  - name: top_k
+    label:
+      zh_Hans: 取样数量
+      en_US: Top k
+    type: int
+    help:
+      zh_Hans: 仅从每个后续标记的前 K 个选项中采样。
+      en_US: Only sample from the top K options for each subsequent token.
+    required: false
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    default: 8000
+    min: 1
+    max: 128000
+  - name: presence_penalty
+    use_template: presence_penalty
+  - name: frequency_penalty
+    use_template: frequency_penalty
+    default: 1
+    min: 1
+    max: 2
+  - name: with_search_enhance
+    label:
+      zh_Hans: 搜索增强
+      en_US: Search Enhance
+    type: boolean
+    help:
+      zh_Hans: 允许模型自行进行外部搜索，以增强生成结果。
+      en_US: Allow the model to perform external search to enhance the generation results.
+    required: false

api/core/model_runtime/model_providers/baichuan/llm/baichuan3-turbo.yaml

@@ -0,0 +1,45 @@
+model: baichuan3-turbo
+label:
+  en_US: Baichuan3-Turbo
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 32000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+  - name: top_p
+    use_template: top_p
+  - name: top_k
+    label:
+      zh_Hans: 取样数量
+      en_US: Top k
+    type: int
+    help:
+      zh_Hans: 仅从每个后续标记的前 K 个选项中采样。
+      en_US: Only sample from the top K options for each subsequent token.
+    required: false
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    default: 8000
+    min: 1
+    max: 32000
+  - name: presence_penalty
+    use_template: presence_penalty
+  - name: frequency_penalty
+    use_template: frequency_penalty
+    default: 1
+    min: 1
+    max: 2
+  - name: with_search_enhance
+    label:
+      zh_Hans: 搜索增强
+      en_US: Search Enhance
+    type: boolean
+    help:
+      zh_Hans: 允许模型自行进行外部搜索，以增强生成结果。
+      en_US: Allow the model to perform external search to enhance the generation results.
+    required: false

api/core/model_runtime/model_providers/baichuan/llm/baichuan4.yaml

@@ -0,0 +1,45 @@
+model: baichuan4
+label:
+  en_US: Baichuan4
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 32000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+  - name: top_p
+    use_template: top_p
+  - name: top_k
+    label:
+      zh_Hans: 取样数量
+      en_US: Top k
+    type: int
+    help:
+      zh_Hans: 仅从每个后续标记的前 K 个选项中采样。
+      en_US: Only sample from the top K options for each subsequent token.
+    required: false
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    default: 8000
+    min: 1
+    max: 32000
+  - name: presence_penalty
+    use_template: presence_penalty
+  - name: frequency_penalty
+    use_template: frequency_penalty
+    default: 1
+    min: 1
+    max: 2
+  - name: with_search_enhance
+    label:
+      zh_Hans: 搜索增强
+      en_US: Search Enhance
+    type: boolean
+    help:
+      zh_Hans: 允许模型自行进行外部搜索，以增强生成结果。
+      en_US: Allow the model to perform external search to enhance the generation results.
+    required: false

api/core/model_runtime/model_providers/baichuan/llm/baichuan_turbo.py

@@ -51,6 +51,9 @@ class BaichuanModel:
             'baichuan2-turbo': 'Baichuan2-Turbo',
             'baichuan2-turbo-192k': 'Baichuan2-Turbo-192k',
             'baichuan2-53b': 'Baichuan2-53B',
+            'baichuan3-turbo': 'Baichuan3-Turbo',
+            'baichuan3-turbo-128k': 'Baichuan3-Turbo-128k',
+            'baichuan4': 'Baichuan4',
         }[model]
 
     def _handle_chat_generate_response(self, response) -> BaichuanMessage:
@@ -110,7 +113,8 @@ class BaichuanModel:
     def _build_parameters(self, model: str, stream: bool, messages: list[BaichuanMessage],
                                parameters: dict[str, Any]) \
         -> dict[str, Any]:
-        if model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b':
+        if (model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b'
+                or model == 'baichuan3-turbo' or model == 'baichuan3-turbo-128k'  or model == 'baichuan4'):
             prompt_messages = []
             for message in messages:
                 if message.role == BaichuanMessage.Role.USER.value or message.role == BaichuanMessage.Role._SYSTEM.value:
@@ -143,7 +147,8 @@ class BaichuanModel:
             raise BadRequestError(f"Unknown model: {model}")
         
     def _build_headers(self, model: str, data: dict[str, Any]) -> dict[str, Any]:
-        if model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b':
+        if (model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b'
+                or model == 'baichuan3-turbo' or model == 'baichuan3-turbo-128k'  or model == 'baichuan4'):
             # there is no secret key for turbo api
             return {
                 'Content-Type': 'application/json',
@@ -160,7 +165,8 @@ class BaichuanModel:
                  parameters: dict[str, Any], timeout: int) \
         -> Union[Generator, BaichuanMessage]:
         
-        if model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b':
+        if (model == 'baichuan2-turbo' or model == 'baichuan2-turbo-192k' or model == 'baichuan2-53b'
+                or model == 'baichuan3-turbo' or model == 'baichuan3-turbo-128k'  or model == 'baichuan4'):
             api_base = 'https://api.baichuan-ai.com/v1/chat/completions'
         else:
             raise BadRequestError(f"Unknown model: {model}")

api/core/model_runtime/model_providers/bedrock/llm/llm.py

@@ -358,26 +358,25 @@ class BedrockLargeLanguageModel(LargeLanguageModel):
 
         return message_dict
 
-    def get_num_tokens(self, model: str, credentials: dict, messages: list[PromptMessage] | str,
+    def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list[PromptMessage] | str,
                        tools: Optional[list[PromptMessageTool]] = None) -> int:
         """
         Get number of tokens for given prompt messages
 
         :param model: model name
         :param credentials: model credentials
-        :param messages: prompt messages or message string
+        :param prompt_messages: prompt messages or message string
         :param tools: tools for tool calling
         :return:md = genai.GenerativeModel(model)
         """
         prefix = model.split('.')[0]
         model_name = model.split('.')[1]
-        if isinstance(messages, str):
-            prompt = messages
+        if isinstance(prompt_messages, str):
+            prompt = prompt_messages
         else:
-            prompt = self._convert_messages_to_prompt(messages, prefix, model_name)
+            prompt = self._convert_messages_to_prompt(prompt_messages, prefix, model_name)
 
         return self._get_num_tokens_by_gpt2(prompt)
-    
 
     def validate_credentials(self, model: str, credentials: dict) -> None:
         """

api/core/model_runtime/model_providers/model_provider_factory.py

@@ -4,13 +4,13 @@ from typing import Optional
 
 from pydantic import BaseModel
 
+from core.helper.module_import_helper import load_single_subclass_from_source
+from core.helper.position_helper import get_position_map, sort_to_dict_by_position_map
 from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.entities.provider_entities import ProviderConfig, ProviderEntity, SimpleProviderEntity
 from core.model_runtime.model_providers.__base.model_provider import ModelProvider
 from core.model_runtime.schema_validators.model_credential_schema_validator import ModelCredentialSchemaValidator
 from core.model_runtime.schema_validators.provider_credential_schema_validator import ProviderCredentialSchemaValidator
-from core.utils.module_import_helper import load_single_subclass_from_source
-from core.utils.position_helper import get_position_map, sort_to_dict_by_position_map
 
 logger = logging.getLogger(__name__)
 

api/core/model_runtime/model_providers/nvidia_nim/_assets/icon_s_en.svg

@@ -0,0 +1,3 @@
+<svg width="567" height="376" viewBox="0 0 567 376" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M58.0366 161.868C58.0366 161.868 109.261 86.2912 211.538 78.4724V51.053C98.2528 60.1511 0.152344 156.098 0.152344 156.098C0.152344 156.098 55.7148 316.717 211.538 331.426V302.282C97.1876 287.896 58.0366 161.868 58.0366 161.868ZM211.538 244.32V271.013C125.114 255.603 101.125 165.768 101.125 165.768C101.125 165.768 142.621 119.799 211.538 112.345V141.633C211.486 141.633 211.449 141.617 211.406 141.617C175.235 137.276 146.978 171.067 146.978 171.067C146.978 171.067 162.816 227.949 211.538 244.32ZM211.538 0.47998V51.053C214.864 50.7981 218.189 50.5818 221.533 50.468C350.326 46.1273 434.243 156.098 434.243 156.098C434.243 156.098 337.861 273.296 237.448 273.296C228.245 273.296 219.63 272.443 211.538 271.009V302.282C218.695 303.201 225.903 303.667 233.119 303.675C326.56 303.675 394.134 255.954 459.566 199.474C470.415 208.162 514.828 229.299 523.958 238.55C461.745 290.639 316.752 332.626 234.551 332.626C226.627 332.626 219.018 332.148 211.538 331.426V375.369H566.701V0.47998H211.538ZM211.538 112.345V78.4724C214.829 78.2425 218.146 78.0672 221.533 77.9602C314.148 75.0512 374.909 157.548 374.909 157.548C374.909 157.548 309.281 248.693 238.914 248.693C228.787 248.693 219.707 247.065 211.536 244.318V141.631C247.591 145.987 254.848 161.914 276.524 198.049L324.737 157.398C324.737 157.398 289.544 111.243 230.219 111.243C223.768 111.241 217.597 111.696 211.538 112.345Z" fill="#77B900"/>
+</svg>

api/core/model_runtime/model_providers/nvidia_nim/llm/llm.py

@@ -0,0 +1,12 @@
+import logging
+
+from core.model_runtime.model_providers.openai_api_compatible.llm.llm import OAIAPICompatLargeLanguageModel
+
+logger = logging.getLogger(__name__)
+
+
+class NVIDIANIMProvider(OAIAPICompatLargeLanguageModel):
+    """
+    Model class for NVIDIA NIM large language model.
+    """
+    pass

api/core/model_runtime/model_providers/nvidia_nim/nvidia_nim.py

@@ -0,0 +1,11 @@
+import logging
+
+from core.model_runtime.model_providers.__base.model_provider import ModelProvider
+
+logger = logging.getLogger(__name__)
+
+
+class NVIDIANIMProvider(ModelProvider):
+
+    def validate_provider_credentials(self, credentials: dict) -> None:
+        pass

api/core/model_runtime/model_providers/nvidia_nim/nvidia_nim.yaml

@@ -0,0 +1,79 @@
+provider: nvidia_nim
+label:
+  en_US: NVIDIA NIM
+description:
+  en_US: NVIDIA NIM, a set of easy-to-use inference microservices.
+  zh_Hans: NVIDIA NIM，一组易于使用的模型推理微服务。
+icon_small:
+  en_US: icon_s_en.svg
+icon_large:
+  en_US: icon_l_en.png
+background: "#EFFDFD"
+help:
+  title:
+    en_US: Learn more about NVIDIA NIM
+    zh_Hans: 了解 NVIDIA NIM 更多信息
+  url:
+    en_US: https://www.nvidia.com/en-us/ai/
+supported_model_types:
+  - llm
+configurate_methods:
+  - customizable-model
+model_credential_schema:
+  model:
+    label:
+      en_US: Model Name
+      zh_Hans: 模型名称
+    placeholder:
+      en_US: Enter full model name
+      zh_Hans: 输入模型全称
+  credential_form_schemas:
+    - variable: endpoint_url
+      label:
+        zh_Hans: API endpoint URL
+        en_US: API endpoint URL
+      type: text-input
+      required: true
+      placeholder:
+        zh_Hans: Base URL, e.g. http://192.168.1.100:8000/v1
+        en_US: Base URL, e.g. http://192.168.1.100:8000/v1
+    - variable: mode
+      show_on:
+        - variable: __model_type
+          value: llm
+      label:
+        en_US: Completion mode
+      type: select
+      required: false
+      default: chat
+      placeholder:
+        zh_Hans: 选择对话类型
+        en_US: Select completion mode
+      options:
+        - value: completion
+          label:
+            en_US: Completion
+            zh_Hans: 补全
+        - value: chat
+          label:
+            en_US: Chat
+            zh_Hans: 对话
+    - variable: context_size
+      label:
+        zh_Hans: 模型上下文长度
+        en_US: Model context size
+      required: true
+      type: text-input
+      default: '4096'
+      placeholder:
+        zh_Hans: 在此输入您的模型上下文长度
+        en_US: Enter your Model context size
+    - variable: max_tokens_to_sample
+      label:
+        zh_Hans: 最大 token 上限
+        en_US: Upper bound for max tokens
+      show_on:
+        - variable: __model_type
+          value: llm
+      default: '4096'
+      type: text-input

api/core/model_runtime/model_providers/ollama/llm/llm.py

@@ -8,7 +8,12 @@ from urllib.parse import urljoin
 
 import requests
 
-from core.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMResultChunk, LLMResultChunkDelta
+from core.model_runtime.entities.llm_entities import (
+    LLMMode,
+    LLMResult,
+    LLMResultChunk,
+    LLMResultChunkDelta,
+)
 from core.model_runtime.entities.message_entities import (
     AssistantPromptMessage,
     ImagePromptMessageContent,
@@ -40,7 +45,9 @@ from core.model_runtime.errors.invoke import (
     InvokeServerUnavailableError,
 )
 from core.model_runtime.errors.validate import CredentialsValidateFailedError
-from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.model_runtime.model_providers.__base.large_language_model import (
+    LargeLanguageModel,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -50,11 +57,17 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
     Model class for Ollama large language model.
     """
 
-    def _invoke(self, model: str, credentials: dict,
-                prompt_messages: list[PromptMessage], model_parameters: dict,
-                tools: Optional[list[PromptMessageTool]] = None, stop: Optional[list[str]] = None,
-                stream: bool = True, user: Optional[str] = None) \
-            -> Union[LLMResult, Generator]:
+    def _invoke(
+        self,
+        model: str,
+        credentials: dict,
+        prompt_messages: list[PromptMessage],
+        model_parameters: dict,
+        tools: Optional[list[PromptMessageTool]] = None,
+        stop: Optional[list[str]] = None,
+        stream: bool = True,
+        user: Optional[str] = None,
+    ) -> Union[LLMResult, Generator]:
         """
         Invoke large language model
 
@@ -75,11 +88,16 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
             model_parameters=model_parameters,
             stop=stop,
             stream=stream,
-            user=user
+            user=user,
         )
 
-    def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list[PromptMessage],
-                       tools: Optional[list[PromptMessageTool]] = None) -> int:
+    def get_num_tokens(
+        self,
+        model: str,
+        credentials: dict,
+        prompt_messages: list[PromptMessage],
+        tools: Optional[list[PromptMessageTool]] = None,
+    ) -> int:
         """
         Get number of tokens for given prompt messages
 
@@ -100,10 +118,12 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
             if isinstance(first_prompt_message.content, str):
                 text = first_prompt_message.content
             else:
-                text = ''
+                text = ""
                 for message_content in first_prompt_message.content:
                     if message_content.type == PromptMessageContentType.TEXT:
-                        message_content = cast(TextPromptMessageContent, message_content)
+                        message_content = cast(
+                            TextPromptMessageContent, message_content
+                        )
                         text = message_content.data
                         break
             return self._get_num_tokens_by_gpt2(text)
@@ -121,19 +141,28 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
                 model=model,
                 credentials=credentials,
                 prompt_messages=[UserPromptMessage(content="ping")],
-                model_parameters={
-                    'num_predict': 5
-                },
-                stream=False
+                model_parameters={"num_predict": 5},
+                stream=False,
             )
         except InvokeError as ex:
-            raise CredentialsValidateFailedError(f'An error occurred during credentials validation: {ex.description}')
+            raise CredentialsValidateFailedError(
+                f"An error occurred during credentials validation: {ex.description}"
+            )
         except Exception as ex:
-            raise CredentialsValidateFailedError(f'An error occurred during credentials validation: {str(ex)}')
+            raise CredentialsValidateFailedError(
+                f"An error occurred during credentials validation: {str(ex)}"
+            )
 
-    def _generate(self, model: str, credentials: dict,
-                  prompt_messages: list[PromptMessage], model_parameters: dict, stop: Optional[list[str]] = None,
-                  stream: bool = True, user: Optional[str] = None) -> Union[LLMResult, Generator]:
+    def _generate(
+        self,
+        model: str,
+        credentials: dict,
+        prompt_messages: list[PromptMessage],
+        model_parameters: dict,
+        stop: Optional[list[str]] = None,
+        stream: bool = True,
+        user: Optional[str] = None,
+    ) -> Union[LLMResult, Generator]:
         """
         Invoke llm completion model
 
@@ -146,76 +175,93 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
         :param user: unique user id
         :return: full response or stream response chunk generator result
         """
-        headers = {
-            'Content-Type': 'application/json'
-        }
+        headers = {"Content-Type": "application/json"}
 
-        endpoint_url = credentials['base_url']
-        if not endpoint_url.endswith('/'):
-            endpoint_url += '/'
+        endpoint_url = credentials["base_url"]
+        if not endpoint_url.endswith("/"):
+            endpoint_url += "/"
 
         # prepare the payload for a simple ping to the model
-        data = {
-            'model': model,
-            'stream': stream
-        }
+        data = {"model": model, "stream": stream}
 
-        if 'format' in model_parameters:
-            data['format'] = model_parameters['format']
-            del model_parameters['format']
+        if "format" in model_parameters:
+            data["format"] = model_parameters["format"]
+            del model_parameters["format"]
 
-        data['options'] = model_parameters or {}
+        if "keep_alive" in model_parameters:
+            data["keep_alive"] = model_parameters["keep_alive"]
+            del model_parameters["keep_alive"]
+
+        data["options"] = model_parameters or {}
 
         if stop:
-            data['stop'] = "\n".join(stop)
+            data["stop"] = "\n".join(stop)
 
-        completion_type = LLMMode.value_of(credentials['mode'])
+        completion_type = LLMMode.value_of(credentials["mode"])
 
         if completion_type is LLMMode.CHAT:
-            endpoint_url = urljoin(endpoint_url, 'api/chat')
-            data['messages'] = [self._convert_prompt_message_to_dict(m) for m in prompt_messages]
+            endpoint_url = urljoin(endpoint_url, "api/chat")
+            data["messages"] = [
+                self._convert_prompt_message_to_dict(m) for m in prompt_messages
+            ]
         else:
-            endpoint_url = urljoin(endpoint_url, 'api/generate')
+            endpoint_url = urljoin(endpoint_url, "api/generate")
             first_prompt_message = prompt_messages[0]
             if isinstance(first_prompt_message, UserPromptMessage):
                 first_prompt_message = cast(UserPromptMessage, first_prompt_message)
                 if isinstance(first_prompt_message.content, str):
-                    data['prompt'] = first_prompt_message.content
+                    data["prompt"] = first_prompt_message.content
                 else:
-                    text = ''
+                    text = ""
                     images = []
                     for message_content in first_prompt_message.content:
                         if message_content.type == PromptMessageContentType.TEXT:
-                            message_content = cast(TextPromptMessageContent, message_content)
+                            message_content = cast(
+                                TextPromptMessageContent, message_content
+                            )
                             text = message_content.data
                         elif message_content.type == PromptMessageContentType.IMAGE:
-                            message_content = cast(ImagePromptMessageContent, message_content)
-                            image_data = re.sub(r'^data:image\/[a-zA-Z]+;base64,', '', message_content.data)
+                            message_content = cast(
+                                ImagePromptMessageContent, message_content
+                            )
+                            image_data = re.sub(
+                                r"^data:image\/[a-zA-Z]+;base64,",
+                                "",
+                                message_content.data,
+                            )
                             images.append(image_data)
 
-                    data['prompt'] = text
-                    data['images'] = images
+                    data["prompt"] = text
+                    data["images"] = images
 
         # send a post request to validate the credentials
         response = requests.post(
-            endpoint_url,
-            headers=headers,
-            json=data,
-            timeout=(10, 300),
-            stream=stream
+            endpoint_url, headers=headers, json=data, timeout=(10, 300), stream=stream
         )
 
         response.encoding = "utf-8"
         if response.status_code != 200:
-            raise InvokeError(f"API request failed with status code {response.status_code}: {response.text}")
+            raise InvokeError(
+                f"API request failed with status code {response.status_code}: {response.text}"
+            )
 
         if stream:
-            return self._handle_generate_stream_response(model, credentials, completion_type, response, prompt_messages)
+            return self._handle_generate_stream_response(
+                model, credentials, completion_type, response, prompt_messages
+            )
 
-        return self._handle_generate_response(model, credentials, completion_type, response, prompt_messages)
+        return self._handle_generate_response(
+            model, credentials, completion_type, response, prompt_messages
+        )
 
-    def _handle_generate_response(self, model: str, credentials: dict, completion_type: LLMMode,
-                                  response: requests.Response, prompt_messages: list[PromptMessage]) -> LLMResult:
+    def _handle_generate_response(
+        self,
+        model: str,
+        credentials: dict,
+        completion_type: LLMMode,
+        response: requests.Response,
+        prompt_messages: list[PromptMessage],
+    ) -> LLMResult:
         """
         Handle llm completion response
 
@@ -229,14 +275,14 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
         response_json = response.json()
 
         if completion_type is LLMMode.CHAT:
-            message = response_json.get('message', {})
-            response_content = message.get('content', '')
+            message = response_json.get("message", {})
+            response_content = message.get("content", "")
         else:
-            response_content = response_json['response']
+            response_content = response_json["response"]
 
         assistant_message = AssistantPromptMessage(content=response_content)
 
-        if 'prompt_eval_count' in response_json and 'eval_count' in response_json:
+        if "prompt_eval_count" in response_json and "eval_count" in response_json:
             # transform usage
             prompt_tokens = response_json["prompt_eval_count"]
             completion_tokens = response_json["eval_count"]
@@ -246,7 +292,9 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
             completion_tokens = self._get_num_tokens_by_gpt2(assistant_message.content)
 
         # transform usage
-        usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)
+        usage = self._calc_response_usage(
+            model, credentials, prompt_tokens, completion_tokens
+        )
 
         # transform response
         result = LLMResult(
@@ -258,8 +306,14 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
 
         return result
 
-    def _handle_generate_stream_response(self, model: str, credentials: dict, completion_type: LLMMode,
-                                         response: requests.Response, prompt_messages: list[PromptMessage]) -> Generator:
+    def _handle_generate_stream_response(
+        self,
+        model: str,
+        credentials: dict,
+        completion_type: LLMMode,
+        response: requests.Response,
+        prompt_messages: list[PromptMessage],
+    ) -> Generator:
         """
         Handle llm completion stream response
 
@@ -270,17 +324,20 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
         :param prompt_messages: prompt messages
         :return: llm response chunk generator result
         """
-        full_text = ''
+        full_text = ""
         chunk_index = 0
 
-        def create_final_llm_result_chunk(index: int, message: AssistantPromptMessage, finish_reason: str) \
-                -> LLMResultChunk:
+        def create_final_llm_result_chunk(
+            index: int, message: AssistantPromptMessage, finish_reason: str
+        ) -> LLMResultChunk:
             # calculate num tokens
             prompt_tokens = self._get_num_tokens_by_gpt2(prompt_messages[0].content)
             completion_tokens = self._get_num_tokens_by_gpt2(full_text)
 
             # transform usage
-            usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)
+            usage = self._calc_response_usage(
+                model, credentials, prompt_tokens, completion_tokens
+            )
 
             return LLMResultChunk(
                 model=model,
@@ -289,11 +346,11 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
                     index=index,
                     message=message,
                     finish_reason=finish_reason,
-                    usage=usage
-                )
+                    usage=usage,
+                ),
             )
 
-        for chunk in response.iter_lines(decode_unicode=True, delimiter='\n'):
+        for chunk in response.iter_lines(decode_unicode=True, delimiter="\n"):
             if not chunk:
                 continue
 
@@ -304,7 +361,7 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
                 yield create_final_llm_result_chunk(
                     index=chunk_index,
                     message=AssistantPromptMessage(content=""),
-                    finish_reason="Non-JSON encountered."
+                    finish_reason="Non-JSON encountered.",
                 )
 
                 chunk_index += 1
@@ -314,55 +371,57 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
                 if not chunk_json:
                     continue
 
-                if 'message' not in chunk_json:
-                    text = ''
+                if "message" not in chunk_json:
+                    text = ""
                 else:
-                    text = chunk_json.get('message').get('content', '')
+                    text = chunk_json.get("message").get("content", "")
             else:
                 if not chunk_json:
                     continue
 
                 # transform assistant message to prompt message
-                text = chunk_json['response']
+                text = chunk_json["response"]
 
-            assistant_prompt_message = AssistantPromptMessage(
-                content=text
-            )
+            assistant_prompt_message = AssistantPromptMessage(content=text)
 
             full_text += text
 
-            if chunk_json['done']:
+            if chunk_json["done"]:
                 # calculate num tokens
-                if 'prompt_eval_count' in chunk_json and 'eval_count' in chunk_json:
+                if "prompt_eval_count" in chunk_json and "eval_count" in chunk_json:
                     # transform usage
                     prompt_tokens = chunk_json["prompt_eval_count"]
                     completion_tokens = chunk_json["eval_count"]
                 else:
                     # calculate num tokens
-                    prompt_tokens = self._get_num_tokens_by_gpt2(prompt_messages[0].content)
+                    prompt_tokens = self._get_num_tokens_by_gpt2(
+                        prompt_messages[0].content
+                    )
                     completion_tokens = self._get_num_tokens_by_gpt2(full_text)
 
                 # transform usage
-                usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)
+                usage = self._calc_response_usage(
+                    model, credentials, prompt_tokens, completion_tokens
+                )
 
                 yield LLMResultChunk(
-                    model=chunk_json['model'],
+                    model=chunk_json["model"],
                     prompt_messages=prompt_messages,
                     delta=LLMResultChunkDelta(
                         index=chunk_index,
                         message=assistant_prompt_message,
-                        finish_reason='stop',
-                        usage=usage
-                    )
+                        finish_reason="stop",
+                        usage=usage,
+                    ),
                 )
             else:
                 yield LLMResultChunk(
-                    model=chunk_json['model'],
+                    model=chunk_json["model"],
                     prompt_messages=prompt_messages,
                     delta=LLMResultChunkDelta(
                         index=chunk_index,
                         message=assistant_prompt_message,
-                    )
+                    ),
                 )
 
             chunk_index += 1
@@ -376,15 +435,21 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
             if isinstance(message.content, str):
                 message_dict = {"role": "user", "content": message.content}
             else:
-                text = ''
+                text = ""
                 images = []
                 for message_content in message.content:
                     if message_content.type == PromptMessageContentType.TEXT:
-                        message_content = cast(TextPromptMessageContent, message_content)
+                        message_content = cast(
+                            TextPromptMessageContent, message_content
+                        )
                         text = message_content.data
                     elif message_content.type == PromptMessageContentType.IMAGE:
-                        message_content = cast(ImagePromptMessageContent, message_content)
-                        image_data = re.sub(r'^data:image\/[a-zA-Z]+;base64,', '', message_content.data)
+                        message_content = cast(
+                            ImagePromptMessageContent, message_content
+                        )
+                        image_data = re.sub(
+                            r"^data:image\/[a-zA-Z]+;base64,", "", message_content.data
+                        )
                         images.append(image_data)
 
                 message_dict = {"role": "user", "content": text, "images": images}
@@ -414,7 +479,9 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
 
         return num_tokens
 
-    def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:
+    def get_customizable_model_schema(
+        self, model: str, credentials: dict
+    ) -> AIModelEntity:
         """
         Get customizable model schema.
 
@@ -425,20 +492,19 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
         """
         extras = {}
 
-        if 'vision_support' in credentials and credentials['vision_support'] == 'true':
-            extras['features'] = [ModelFeature.VISION]
+        if "vision_support" in credentials and credentials["vision_support"] == "true":
+            extras["features"] = [ModelFeature.VISION]
 
         entity = AIModelEntity(
             model=model,
-            label=I18nObject(
-                zh_Hans=model,
-                en_US=model
-            ),
+            label=I18nObject(zh_Hans=model, en_US=model),
             model_type=ModelType.LLM,
             fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
             model_properties={
-                ModelPropertyKey.MODE: credentials.get('mode'),
-                ModelPropertyKey.CONTEXT_SIZE: int(credentials.get('context_size', 4096)),
+                ModelPropertyKey.MODE: credentials.get("mode"),
+                ModelPropertyKey.CONTEXT_SIZE: int(
+                    credentials.get("context_size", 4096)
+                ),
             },
             parameter_rules=[
                 ParameterRule(
@@ -446,152 +512,195 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
                     use_template=DefaultParameterName.TEMPERATURE.value,
                     label=I18nObject(en_US="Temperature"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="The temperature of the model. "
-                                          "Increasing the temperature will make the model answer "
-                                          "more creatively. (Default: 0.8)"),
+                    help=I18nObject(
+                        en_US="The temperature of the model. "
+                        "Increasing the temperature will make the model answer "
+                        "more creatively. (Default: 0.8)"
+                    ),
                     default=0.1,
                     min=0,
-                    max=1
+                    max=1,
                 ),
                 ParameterRule(
                     name=DefaultParameterName.TOP_P.value,
                     use_template=DefaultParameterName.TOP_P.value,
                     label=I18nObject(en_US="Top P"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="Works together with top-k. A higher value (e.g., 0.95) will lead to "
-                                          "more diverse text, while a lower value (e.g., 0.5) will generate more "
-                                          "focused and conservative text. (Default: 0.9)"),
+                    help=I18nObject(
+                        en_US="Works together with top-k. A higher value (e.g., 0.95) will lead to "
+                        "more diverse text, while a lower value (e.g., 0.5) will generate more "
+                        "focused and conservative text. (Default: 0.9)"
+                    ),
                     default=0.9,
                     min=0,
-                    max=1
+                    max=1,
                 ),
                 ParameterRule(
                     name="top_k",
                     label=I18nObject(en_US="Top K"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Reduces the probability of generating nonsense. "
-                                          "A higher value (e.g. 100) will give more diverse answers, "
-                                          "while a lower value (e.g. 10) will be more conservative. (Default: 40)"),
+                    help=I18nObject(
+                        en_US="Reduces the probability of generating nonsense. "
+                        "A higher value (e.g. 100) will give more diverse answers, "
+                        "while a lower value (e.g. 10) will be more conservative. (Default: 40)"
+                    ),
                     min=1,
-                    max=100
+                    max=100,
                 ),
                 ParameterRule(
-                    name='repeat_penalty',
+                    name="repeat_penalty",
                     label=I18nObject(en_US="Repeat Penalty"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="Sets how strongly to penalize repetitions. "
-                                          "A higher value (e.g., 1.5) will penalize repetitions more strongly, "
-                                          "while a lower value (e.g., 0.9) will be more lenient. (Default: 1.1)"),
+                    help=I18nObject(
+                        en_US="Sets how strongly to penalize repetitions. "
+                        "A higher value (e.g., 1.5) will penalize repetitions more strongly, "
+                        "while a lower value (e.g., 0.9) will be more lenient. (Default: 1.1)"
+                    ),
                     min=-2,
-                    max=2
+                    max=2,
                 ),
                 ParameterRule(
-                    name='num_predict',
-                    use_template='max_tokens',
+                    name="num_predict",
+                    use_template="max_tokens",
                     label=I18nObject(en_US="Num Predict"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Maximum number of tokens to predict when generating text. "
-                                          "(Default: 128, -1 = infinite generation, -2 = fill context)"),
-                    default=512 if int(credentials.get('max_tokens', 4096)) >= 768 else 128,
+                    help=I18nObject(
+                        en_US="Maximum number of tokens to predict when generating text. "
+                        "(Default: 128, -1 = infinite generation, -2 = fill context)"
+                    ),
+                    default=(
+                        512 if int(credentials.get("max_tokens", 4096)) >= 768 else 128
+                    ),
                     min=-2,
-                    max=int(credentials.get('max_tokens', 4096)),
+                    max=int(credentials.get("max_tokens", 4096)),
                 ),
                 ParameterRule(
-                    name='mirostat',
+                    name="mirostat",
                     label=I18nObject(en_US="Mirostat sampling"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Enable Mirostat sampling for controlling perplexity. "
-                                          "(default: 0, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)"),
+                    help=I18nObject(
+                        en_US="Enable Mirostat sampling for controlling perplexity. "
+                        "(default: 0, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)"
+                    ),
                     min=0,
-                    max=2
+                    max=2,
                 ),
                 ParameterRule(
-                    name='mirostat_eta',
+                    name="mirostat_eta",
                     label=I18nObject(en_US="Mirostat Eta"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="Influences how quickly the algorithm responds to feedback from "
-                                          "the generated text. A lower learning rate will result in slower adjustments, "
-                                          "while a higher learning rate will make the algorithm more responsive. "
-                                          "(Default: 0.1)"),
-                    precision=1
+                    help=I18nObject(
+                        en_US="Influences how quickly the algorithm responds to feedback from "
+                        "the generated text. A lower learning rate will result in slower adjustments, "
+                        "while a higher learning rate will make the algorithm more responsive. "
+                        "(Default: 0.1)"
+                    ),
+                    precision=1,
                 ),
                 ParameterRule(
-                    name='mirostat_tau',
+                    name="mirostat_tau",
                     label=I18nObject(en_US="Mirostat Tau"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="Controls the balance between coherence and diversity of the output. "
-                                          "A lower value will result in more focused and coherent text. (Default: 5.0)"),
-                    precision=1
+                    help=I18nObject(
+                        en_US="Controls the balance between coherence and diversity of the output. "
+                        "A lower value will result in more focused and coherent text. (Default: 5.0)"
+                    ),
+                    precision=1,
                 ),
                 ParameterRule(
-                    name='num_ctx',
+                    name="num_ctx",
                     label=I18nObject(en_US="Size of context window"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Sets the size of the context window used to generate the next token. "
-                                          "(Default: 2048)"),
+                    help=I18nObject(
+                        en_US="Sets the size of the context window used to generate the next token. "
+                        "(Default: 2048)"
+                    ),
                     default=2048,
-                    min=1
+                    min=1,
                 ),
                 ParameterRule(
                     name='num_gpu',
-                    label=I18nObject(en_US="Num GPU"),
+                    label=I18nObject(en_US="GPU Layers"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="The number of layers to send to the GPU(s). "
-                                          "On macOS it defaults to 1 to enable metal support, 0 to disable."),
-                    min=0,
-                    max=1
+                    help=I18nObject(en_US="The number of layers to offload to the GPU(s). "
+                                          "On macOS it defaults to 1 to enable metal support, 0 to disable."
+                                          "As long as a model fits into one gpu it stays in one. "
+                                          "It does not set the number of GPU(s). "),
+                    min=-1,
+                    default=1
                 ),
                 ParameterRule(
-                    name='num_thread',
+                    name="num_thread",
                     label=I18nObject(en_US="Num Thread"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Sets the number of threads to use during computation. "
-                                          "By default, Ollama will detect this for optimal performance. "
-                                          "It is recommended to set this value to the number of physical CPU cores "
-                                          "your system has (as opposed to the logical number of cores)."),
+                    help=I18nObject(
+                        en_US="Sets the number of threads to use during computation. "
+                        "By default, Ollama will detect this for optimal performance. "
+                        "It is recommended to set this value to the number of physical CPU cores "
+                        "your system has (as opposed to the logical number of cores)."
+                    ),
                     min=1,
                 ),
                 ParameterRule(
-                    name='repeat_last_n',
+                    name="repeat_last_n",
                     label=I18nObject(en_US="Repeat last N"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Sets how far back for the model to look back to prevent repetition. "
-                                          "(Default: 64, 0 = disabled, -1 = num_ctx)"),
-                    min=-1
+                    help=I18nObject(
+                        en_US="Sets how far back for the model to look back to prevent repetition. "
+                        "(Default: 64, 0 = disabled, -1 = num_ctx)"
+                    ),
+                    min=-1,
                 ),
                 ParameterRule(
-                    name='tfs_z',
+                    name="tfs_z",
                     label=I18nObject(en_US="TFS Z"),
                     type=ParameterType.FLOAT,
-                    help=I18nObject(en_US="Tail free sampling is used to reduce the impact of less probable tokens "
-                                          "from the output. A higher value (e.g., 2.0) will reduce the impact more, "
-                                          "while a value of 1.0 disables this setting. (default: 1)"),
-                    precision=1
+                    help=I18nObject(
+                        en_US="Tail free sampling is used to reduce the impact of less probable tokens "
+                        "from the output. A higher value (e.g., 2.0) will reduce the impact more, "
+                        "while a value of 1.0 disables this setting. (default: 1)"
+                    ),
+                    precision=1,
                 ),
                 ParameterRule(
-                    name='seed',
+                    name="seed",
                     label=I18nObject(en_US="Seed"),
                     type=ParameterType.INT,
-                    help=I18nObject(en_US="Sets the random number seed to use for generation. Setting this to "
-                                          "a specific number will make the model generate the same text for "
-                                          "the same prompt. (Default: 0)"),
+                    help=I18nObject(
+                        en_US="Sets the random number seed to use for generation. Setting this to "
+                        "a specific number will make the model generate the same text for "
+                        "the same prompt. (Default: 0)"
+                    ),
                 ),
                 ParameterRule(
-                    name='format',
+                    name="keep_alive",
+                    label=I18nObject(en_US="Keep Alive"),
+                    type=ParameterType.STRING,
+                    help=I18nObject(
+                        en_US="Sets how long the model is kept in memory after generating a response. "
+                        "This must be a duration string with a unit (e.g., '10m' for 10 minutes or '24h' for 24 hours). "
+                        "A negative number keeps the model loaded indefinitely, and '0' unloads the model immediately after generating a response. "
+                        "Valid time units are 's','m','h'. (Default: 5m)"
+                    ),
+                ),
+                ParameterRule(
+                    name="format",
                     label=I18nObject(en_US="Format"),
                     type=ParameterType.STRING,
-                    help=I18nObject(en_US="the format to return a response in."
-                                          " Currently the only accepted value is json."),
-                    options=['json'],
-                )
+                    help=I18nObject(
+                        en_US="the format to return a response in."
+                        " Currently the only accepted value is json."
+                    ),
+                    options=["json"],
+                ),
             ],
             pricing=PriceConfig(
-                input=Decimal(credentials.get('input_price', 0)),
-                output=Decimal(credentials.get('output_price', 0)),
-                unit=Decimal(credentials.get('unit', 0)),
-                currency=credentials.get('currency', "USD")
+                input=Decimal(credentials.get("input_price", 0)),
+                output=Decimal(credentials.get("output_price", 0)),
+                unit=Decimal(credentials.get("unit", 0)),
+                currency=credentials.get("currency", "USD"),
             ),
-            **extras
+            **extras,
         )
 
         return entity
@@ -619,10 +728,10 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
             ],
             InvokeServerUnavailableError: [
                 requests.exceptions.ConnectionError,  # Engine Overloaded
-                requests.exceptions.HTTPError  # Server Error
+                requests.exceptions.HTTPError,  # Server Error
             ],
             InvokeConnectionError: [
                 requests.exceptions.ConnectTimeout,  # Timeout
-                requests.exceptions.ReadTimeout  # Timeout
-            ]
+                requests.exceptions.ReadTimeout,  # Timeout
+            ],
         }

api/core/model_runtime/model_providers/vertex_ai/llm/anthropic.claude-3-haiku.yaml

@@ -0,0 +1,56 @@
+model: claude-3-haiku@20240307
+label:
+  en_US: Claude 3 Haiku
+model_type: llm
+features:
+  - agent-thought
+  - vision
+model_properties:
+  mode: chat
+  context_size: 200000
+parameter_rules:
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    type: int
+    default: 4096
+    min: 1
+    max: 4096
+    help:
+      zh_Hans: 停止前生成的最大令牌数。请注意，Anthropic Claude 模型可能会在达到 max_tokens 的值之前停止生成令牌。不同的 Anthropic Claude 模型对此参数具有不同的最大值。
+      en_US: The maximum number of tokens to generate before stopping. Note that Anthropic Claude models might stop generating tokens before reaching the value of max_tokens. Different Anthropic Claude models have different maximum values for this parameter.
+  # docs: https://docs.anthropic.com/claude/docs/system-prompts
+  - name: temperature
+    use_template: temperature
+    required: false
+    type: float
+    default: 1
+    min: 0.0
+    max: 1.0
+    help:
+      zh_Hans: 生成内容的随机性。
+      en_US: The amount of randomness injected into the response.
+  - name: top_p
+    required: false
+    type: float
+    default: 0.999
+    min: 0.000
+    max: 1.000
+    help:
+      zh_Hans: 在核采样中，Anthropic Claude 按概率递减顺序计算每个后续标记的所有选项的累积分布，并在达到 top_p 指定的特定概率时将其切断。您应该更改温度或top_p，但不能同时更改两者。
+      en_US: In nucleus sampling, Anthropic Claude computes the cumulative distribution over all the options for each subsequent token in decreasing probability order and cuts it off once it reaches a particular probability specified by top_p. You should alter either temperature or top_p, but not both.
+  - name: top_k
+    required: false
+    type: int
+    default: 0
+    min: 0
+    # tip docs from aws has error, max value is 500
+    max: 500
+    help:
+      zh_Hans: 对于每个后续标记，仅从前 K 个选项中进行采样。使用 top_k 删除长尾低概率响应。
+      en_US: Only sample from the top K options for each subsequent token. Use top_k to remove long tail low probability responses.
+pricing:
+  input: '0.00025'
+  output: '0.00125'
+  unit: '0.001'
+  currency: USD

api/core/model_runtime/model_providers/vertex_ai/llm/anthropic.claude-3-opus.yaml

@@ -0,0 +1,56 @@
+model: claude-3-opus@20240229
+label:
+  en_US: Claude 3 Opus
+model_type: llm
+features:
+  - agent-thought
+  - vision
+model_properties:
+  mode: chat
+  context_size: 200000
+parameter_rules:
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    type: int
+    default: 4096
+    min: 1
+    max: 4096
+    help:
+      zh_Hans: 停止前生成的最大令牌数。请注意，Anthropic Claude 模型可能会在达到 max_tokens 的值之前停止生成令牌。不同的 Anthropic Claude 模型对此参数具有不同的最大值。
+      en_US: The maximum number of tokens to generate before stopping. Note that Anthropic Claude models might stop generating tokens before reaching the value of max_tokens. Different Anthropic Claude models have different maximum values for this parameter.
+  # docs: https://docs.anthropic.com/claude/docs/system-prompts
+  - name: temperature
+    use_template: temperature
+    required: false
+    type: float
+    default: 1
+    min: 0.0
+    max: 1.0
+    help:
+      zh_Hans: 生成内容的随机性。
+      en_US: The amount of randomness injected into the response.
+  - name: top_p
+    required: false
+    type: float
+    default: 0.999
+    min: 0.000
+    max: 1.000
+    help:
+      zh_Hans: 在核采样中，Anthropic Claude 按概率递减顺序计算每个后续标记的所有选项的累积分布，并在达到 top_p 指定的特定概率时将其切断。您应该更改温度或top_p，但不能同时更改两者。
+      en_US: In nucleus sampling, Anthropic Claude computes the cumulative distribution over all the options for each subsequent token in decreasing probability order and cuts it off once it reaches a particular probability specified by top_p. You should alter either temperature or top_p, but not both.
+  - name: top_k
+    required: false
+    type: int
+    default: 0
+    min: 0
+    # tip docs from aws has error, max value is 500
+    max: 500
+    help:
+      zh_Hans: 对于每个后续标记，仅从前 K 个选项中进行采样。使用 top_k 删除长尾低概率响应。
+      en_US: Only sample from the top K options for each subsequent token. Use top_k to remove long tail low probability responses.
+pricing:
+  input: '0.015'
+  output: '0.075'
+  unit: '0.001'
+  currency: USD

api/core/model_runtime/model_providers/vertex_ai/llm/anthropic.claude-3-sonnet.yaml

@@ -0,0 +1,55 @@
+model: claude-3-sonnet@20240229
+label:
+  en_US: Claude 3 Sonnet
+model_type: llm
+features:
+  - agent-thought
+  - vision
+model_properties:
+  mode: chat
+  context_size: 200000
+parameter_rules:
+  - name: max_tokens
+    use_template: max_tokens
+    required: true
+    type: int
+    default: 4096
+    min: 1
+    max: 4096
+    help:
+      zh_Hans: 停止前生成的最大令牌数。请注意，Anthropic Claude 模型可能会在达到 max_tokens 的值之前停止生成令牌。不同的 Anthropic Claude 模型对此参数具有不同的最大值。
+      en_US: The maximum number of tokens to generate before stopping. Note that Anthropic Claude models might stop generating tokens before reaching the value of max_tokens. Different Anthropic Claude models have different maximum values for this parameter.
+  - name: temperature
+    use_template: temperature
+    required: false
+    type: float
+    default: 1
+    min: 0.0
+    max: 1.0
+    help:
+      zh_Hans: 生成内容的随机性。
+      en_US: The amount of randomness injected into the response.
+  - name: top_p
+    required: false
+    type: float
+    default: 0.999
+    min: 0.000
+    max: 1.000
+    help:
+      zh_Hans: 在核采样中，Anthropic Claude 按概率递减顺序计算每个后续标记的所有选项的累积分布，并在达到 top_p 指定的特定概率时将其切断。您应该更改温度或top_p，但不能同时更改两者。
+      en_US: In nucleus sampling, Anthropic Claude computes the cumulative distribution over all the options for each subsequent token in decreasing probability order and cuts it off once it reaches a particular probability specified by top_p. You should alter either temperature or top_p, but not both.
+  - name: top_k
+    required: false
+    type: int
+    default: 0
+    min: 0
+    # tip docs from aws has error, max value is 500
+    max: 500
+    help:
+      zh_Hans: 对于每个后续标记，仅从前 K 个选项中进行采样。使用 top_k 删除长尾低概率响应。
+      en_US: Only sample from the top K options for each subsequent token. Use top_k to remove long tail low probability responses.
+pricing:
+  input: '0.003'
+  output: '0.015'
+  unit: '0.001'
+  currency: USD

api/core/model_runtime/model_providers/vertex_ai/llm/llm.py

@@ -2,21 +2,32 @@ import base64
 import json
 import logging
 from collections.abc import Generator
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 import google.api_core.exceptions as exceptions
 import vertexai.generative_models as glm
+from anthropic import AnthropicVertex, Stream
+from anthropic.types import (
+    ContentBlockDeltaEvent,
+    Message,
+    MessageDeltaEvent,
+    MessageStartEvent,
+    MessageStopEvent,
+    MessageStreamEvent,
+)
 from google.cloud import aiplatform
 from google.oauth2 import service_account
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory
 
-from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
+from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from core.model_runtime.entities.message_entities import (
     AssistantPromptMessage,
+    ImagePromptMessageContent,
     PromptMessage,
     PromptMessageContentType,
     PromptMessageTool,
     SystemPromptMessage,
+    TextPromptMessageContent,
     ToolPromptMessage,
     UserPromptMessage,
 )
@@ -63,9 +74,287 @@ class VertexAiLargeLanguageModel(LargeLanguageModel):
         :param user: unique user id
         :return: full response or stream response chunk generator result
         """
-        # invoke model
+        # invoke anthropic models via anthropic official SDK
+        if "claude" in model:
+            return self._generate_anthropic(model, credentials, prompt_messages, model_parameters, stop, stream, user)
+        # invoke Gemini model
         return self._generate(model, credentials, prompt_messages, model_parameters, tools, stop, stream, user)
-    
+
+    def _generate_anthropic(self, model: str, credentials: dict, prompt_messages: list[PromptMessage], model_parameters: dict,
+                stop: Optional[list[str]] = None, stream: bool = True, user: Optional[str] = None) -> Union[LLMResult, Generator]:
+        """
+        Invoke Anthropic large language model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_messages: prompt messages
+        :param model_parameters: model parameters
+        :param stop: stop words
+        :param stream: is stream response
+        :return: full response or stream response chunk generator result
+        """
+        # use Anthropic official SDK references
+        # - https://github.com/anthropics/anthropic-sdk-python
+        project_id = credentials["vertex_project_id"]
+
+        if 'opus' in model:
+            location = 'us-east5'
+        else:
+            location = 'us-central1'
+            
+        client = AnthropicVertex(
+            region=location, 
+            project_id=project_id
+        )
+
+        extra_model_kwargs = {}
+        if stop:
+            extra_model_kwargs['stop_sequences'] = stop
+
+        system, prompt_message_dicts = self._convert_claude_prompt_messages(prompt_messages)
+
+        if system:
+            extra_model_kwargs['system'] = system
+
+        response = client.messages.create(
+            model=model,
+            messages=prompt_message_dicts,
+            stream=stream,
+            **model_parameters,
+            **extra_model_kwargs
+        )
+
+        if stream:
+            return self._handle_claude_stream_response(model, credentials, response, prompt_messages)
+
+        return self._handle_claude_response(model, credentials, response, prompt_messages)
+
+    def _handle_claude_response(self, model: str, credentials: dict, response: Message,
+                                prompt_messages: list[PromptMessage]) -> LLMResult:
+        """
+        Handle llm chat response
+
+        :param model: model name
+        :param credentials: credentials
+        :param response: response
+        :param prompt_messages: prompt messages
+        :return: full response chunk generator result
+        """
+
+        # transform assistant message to prompt message
+        assistant_prompt_message = AssistantPromptMessage(
+            content=response.content[0].text
+        )
+
+        # calculate num tokens
+        if response.usage:
+            # transform usage
+            prompt_tokens = response.usage.input_tokens
+            completion_tokens = response.usage.output_tokens
+        else:
+            # calculate num tokens
+            prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
+            completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])
+
+        # transform usage
+        usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)
+
+        # transform response
+        response = LLMResult(
+            model=response.model,
+            prompt_messages=prompt_messages,
+            message=assistant_prompt_message,
+            usage=usage
+        )
+
+        return response
+
+    def _handle_claude_stream_response(self, model: str, credentials: dict, response: Stream[MessageStreamEvent],
+                                        prompt_messages: list[PromptMessage], ) -> Generator:
+        """
+        Handle llm chat stream response
+
+        :param model: model name
+        :param credentials: credentials
+        :param response: response
+        :param prompt_messages: prompt messages
+        :return: full response or stream response chunk generator result
+        """
+
+        try:
+            full_assistant_content = ''
+            return_model = None
+            input_tokens = 0
+            output_tokens = 0
+            finish_reason = None
+            index = 0
+
+            for chunk in response:
+                if isinstance(chunk, MessageStartEvent):
+                    return_model = chunk.message.model
+                    input_tokens = chunk.message.usage.input_tokens
+                elif isinstance(chunk, MessageDeltaEvent):
+                    output_tokens = chunk.usage.output_tokens
+                    finish_reason = chunk.delta.stop_reason
+                elif isinstance(chunk, MessageStopEvent):
+                    usage = self._calc_response_usage(model, credentials, input_tokens, output_tokens)
+                    yield LLMResultChunk(
+                        model=return_model,
+                        prompt_messages=prompt_messages,
+                        delta=LLMResultChunkDelta(
+                            index=index + 1,
+                            message=AssistantPromptMessage(
+                                content=''
+                            ),
+                            finish_reason=finish_reason,
+                            usage=usage
+                        )
+                    )
+                elif isinstance(chunk, ContentBlockDeltaEvent):
+                    chunk_text = chunk.delta.text if chunk.delta.text else ''
+                    full_assistant_content += chunk_text
+                    assistant_prompt_message = AssistantPromptMessage(
+                        content=chunk_text if chunk_text else '',
+                    )
+                    index = chunk.index
+                    yield LLMResultChunk(
+                        model=model,
+                        prompt_messages=prompt_messages,
+                        delta=LLMResultChunkDelta(
+                            index=index,
+                            message=assistant_prompt_message,
+                        )
+                    )
+        except Exception as ex:
+            raise InvokeError(str(ex))
+
+    def _calc_claude_response_usage(self, model: str, credentials: dict, prompt_tokens: int, completion_tokens: int) -> LLMUsage:
+        """
+        Calculate response usage
+
+        :param model: model name
+        :param credentials: model credentials
+        :param prompt_tokens: prompt tokens
+        :param completion_tokens: completion tokens
+        :return: usage
+        """
+        # get prompt price info
+        prompt_price_info = self.get_price(
+            model=model,
+            credentials=credentials,
+            price_type=PriceType.INPUT,
+            tokens=prompt_tokens,
+        )
+
+        # get completion price info
+        completion_price_info = self.get_price(
+            model=model,
+            credentials=credentials,
+            price_type=PriceType.OUTPUT,
+            tokens=completion_tokens
+        )
+
+        # transform usage
+        usage = LLMUsage(
+            prompt_tokens=prompt_tokens,
+            prompt_unit_price=prompt_price_info.unit_price,
+            prompt_price_unit=prompt_price_info.unit,
+            prompt_price=prompt_price_info.total_amount,
+            completion_tokens=completion_tokens,
+            completion_unit_price=completion_price_info.unit_price,
+            completion_price_unit=completion_price_info.unit,
+            completion_price=completion_price_info.total_amount,
+            total_tokens=prompt_tokens + completion_tokens,
+            total_price=prompt_price_info.total_amount + completion_price_info.total_amount,
+            currency=prompt_price_info.currency,
+            latency=time.perf_counter() - self.started_at
+        )
+
+        return usage
+
+    def _convert_claude_prompt_messages(self, prompt_messages: list[PromptMessage]) -> tuple[str, list[dict]]:
+        """
+        Convert prompt messages to dict list and system
+        """
+
+        system = ""
+        first_loop = True
+        for message in prompt_messages:
+            if isinstance(message, SystemPromptMessage):
+                message.content=message.content.strip()
+                if first_loop:
+                    system=message.content
+                    first_loop=False
+                else:
+                    system+="\n"
+                    system+=message.content
+
+        prompt_message_dicts = []
+        for message in prompt_messages:
+            if not isinstance(message, SystemPromptMessage):
+                prompt_message_dicts.append(self._convert_claude_prompt_message_to_dict(message))
+
+        return system, prompt_message_dicts
+
+    def _convert_claude_prompt_message_to_dict(self, message: PromptMessage) -> dict:
+        """
+        Convert PromptMessage to dict
+        """
+        if isinstance(message, UserPromptMessage):
+            message = cast(UserPromptMessage, message)
+            if isinstance(message.content, str):
+                message_dict = {"role": "user", "content": message.content}
+            else:
+                sub_messages = []
+                for message_content in message.content:
+                    if message_content.type == PromptMessageContentType.TEXT:
+                        message_content = cast(TextPromptMessageContent, message_content)
+                        sub_message_dict = {
+                            "type": "text",
+                            "text": message_content.data
+                        }
+                        sub_messages.append(sub_message_dict)
+                    elif message_content.type == PromptMessageContentType.IMAGE:
+                        message_content = cast(ImagePromptMessageContent, message_content)
+                        if not message_content.data.startswith("data:"):
+                            # fetch image data from url
+                            try:
+                                image_content = requests.get(message_content.data).content
+                                mime_type, _ = mimetypes.guess_type(message_content.data)
+                                base64_data = base64.b64encode(image_content).decode('utf-8')
+                            except Exception as ex:
+                                raise ValueError(f"Failed to fetch image data from url {message_content.data}, {ex}")
+                        else:
+                            data_split = message_content.data.split(";base64,")
+                            mime_type = data_split[0].replace("data:", "")
+                            base64_data = data_split[1]
+
+                        if mime_type not in ["image/jpeg", "image/png", "image/gif", "image/webp"]:
+                            raise ValueError(f"Unsupported image type {mime_type}, "
+                                             f"only support image/jpeg, image/png, image/gif, and image/webp")
+
+                        sub_message_dict = {
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": mime_type,
+                                "data": base64_data
+                            }
+                        }
+                        sub_messages.append(sub_message_dict)
+
+                message_dict = {"role": "user", "content": sub_messages}
+        elif isinstance(message, AssistantPromptMessage):
+            message = cast(AssistantPromptMessage, message)
+            message_dict = {"role": "assistant", "content": message.content}
+        elif isinstance(message, SystemPromptMessage):
+            message = cast(SystemPromptMessage, message)
+            message_dict = {"role": "system", "content": message.content}
+        else:
+            raise ValueError(f"Got unknown type {message}")
+
+        return message_dict
+
     def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list[PromptMessage],
                        tools: Optional[list[PromptMessageTool]] = None) -> int:
         """
@@ -367,10 +656,8 @@ class VertexAiLargeLanguageModel(LargeLanguageModel):
                     else:
                         metadata, data = c.data.split(',', 1)
                         mime_type = metadata.split(';', 1)[0].split(':')[1]
-                        blob = {"inline_data":{"mime_type":mime_type,"data":data}}
-                        parts.append(blob)
-                
-                glm_content = glm.Content(role="user", parts=[parts])
+                        parts.append(glm.Part.from_data(mime_type=mime_type, data=data))
+                glm_content = glm.Content(role="user", parts=parts)
             return glm_content
         elif isinstance(message, AssistantPromptMessage):
             if message.content:
@@ -438,4 +725,4 @@ class VertexAiLargeLanguageModel(LargeLanguageModel):
                 exceptions.RequestRangeNotSatisfiable,
                 exceptions.Cancelled,
             ]
-        }
+        }

api/core/model_runtime/model_providers/wenxin/llm/ernie-3.5-128k.yaml

@@ -0,0 +1,37 @@
+model: ernie-3.5-128k
+label:
+  en_US: Ernie-3.5-128K
+model_type: llm
+features:
+  - agent-thought
+model_properties:
+  mode: chat
+  context_size: 128000
+parameter_rules:
+  - name: temperature
+    use_template: temperature
+    min: 0.1
+    max: 1.0
+    default: 0.8
+  - name: top_p
+    use_template: top_p
+  - name: max_tokens
+    use_template: max_tokens
+    default: 4096
+    min: 2
+    max: 4096
+  - name: presence_penalty
+    use_template: presence_penalty
+  - name: frequency_penalty
+    use_template: frequency_penalty
+  - name: response_format
+    use_template: response_format
+  - name: disable_search
+    label:
+      zh_Hans: 禁用搜索
+      en_US: Disable Search
+    type: boolean
+    help:
+      zh_Hans: 禁用模型自行进行外部搜索。
+      en_US: Disable the model to perform external search.
+    required: false

api/core/model_runtime/model_providers/wenxin/llm/ernie-3.5-4k-0205.yaml

@@ -35,3 +35,4 @@ parameter_rules:
       zh_Hans: 禁用模型自行进行外部搜索。
       en_US: Disable the model to perform external search.
     required: false
+deprecated: true

api/core/model_runtime/model_providers/wenxin/llm/ernie-3.5-8k-1222.yaml

@@ -35,3 +35,4 @@ parameter_rules:
       zh_Hans: 禁用模型自行进行外部搜索。
       en_US: Disable the model to perform external search.
     required: false
+deprecated: true

api/core/model_runtime/model_providers/wenxin/llm/ernie_bot.py

@@ -129,6 +129,7 @@ class ErnieBotModel:
         'ernie-3.5-8k-0205': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-3.5-8k-0205',
         'ernie-3.5-8k-1222': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-3.5-8k-1222',
         'ernie-3.5-4k-0205': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-3.5-4k-0205',
+        'ernie-3.5-128k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-3.5-128k',
         'ernie-4.0-8k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro',
         'ernie-speed-8k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_speed',
         'ernie-speed-128k': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k',
@@ -144,7 +145,9 @@ class ErnieBotModel:
         'ernie-3.5-8k',
         'ernie-3.5-8k-0205',
         'ernie-3.5-8k-1222',
-        'ernie-3.5-4k-0205'
+        'ernie-3.5-4k-0205',
+        'ernie-3.5-128k',
+        'ernie-4.0-8k'
     ]
 
     api_key: str = ''

api/core/prompt/agent_history_prompt_transform.py

@@ -0,0 +1,82 @@
+from typing import Optional, cast
+
+from core.app.entities.app_invoke_entities import (
+    ModelConfigWithCredentialsEntity,
+)
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_runtime.entities.message_entities import (
+    PromptMessage,
+    SystemPromptMessage,
+    UserPromptMessage,
+)
+from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.prompt.prompt_transform import PromptTransform
+
+
+class AgentHistoryPromptTransform(PromptTransform):
+    """
+    History Prompt Transform for Agent App
+    """
+    def __init__(self,
+                 model_config: ModelConfigWithCredentialsEntity,
+                 prompt_messages: list[PromptMessage],
+                 history_messages: list[PromptMessage],
+                 memory: Optional[TokenBufferMemory] = None,
+                 ):
+        self.model_config = model_config
+        self.prompt_messages = prompt_messages
+        self.history_messages = history_messages
+        self.memory = memory
+
+    def get_prompt(self) -> list[PromptMessage]:
+        prompt_messages = []
+        num_system = 0
+        for prompt_message in self.history_messages:
+            if isinstance(prompt_message, SystemPromptMessage):
+                prompt_messages.append(prompt_message)
+                num_system += 1
+
+        if not self.memory:
+            return prompt_messages
+
+        max_token_limit = self._calculate_rest_token(self.prompt_messages, self.model_config)
+
+        model_type_instance = self.model_config.provider_model_bundle.model_type_instance
+        model_type_instance = cast(LargeLanguageModel, model_type_instance)
+
+        curr_message_tokens = model_type_instance.get_num_tokens(
+            self.memory.model_instance.model,
+            self.memory.model_instance.credentials,
+            self.history_messages
+        )
+        if curr_message_tokens <= max_token_limit:
+            return self.history_messages
+
+        # number of prompt has been appended in current message
+        num_prompt = 0
+        # append prompt messages in desc order
+        for prompt_message in self.history_messages[::-1]:
+            if isinstance(prompt_message, SystemPromptMessage):
+                continue
+            prompt_messages.append(prompt_message)
+            num_prompt += 1
+            # a message is start with UserPromptMessage
+            if isinstance(prompt_message, UserPromptMessage):
+                curr_message_tokens = model_type_instance.get_num_tokens(
+                    self.memory.model_instance.model,
+                    self.memory.model_instance.credentials,
+                    prompt_messages
+                )
+                # if current message token is overflow, drop all the prompts in current message and break
+                if curr_message_tokens > max_token_limit:
+                    prompt_messages = prompt_messages[:-num_prompt]
+                    break
+                num_prompt = 0
+        # return prompt messages in asc order
+        message_prompts = prompt_messages[num_system:]
+        message_prompts.reverse()
+
+        # merge system and message prompt
+        prompt_messages = prompt_messages[:num_system]
+        prompt_messages.extend(message_prompts)
+        return prompt_messages

api/core/prompt/prompt_transform.py

@@ -1,10 +1,10 @@
-from typing import Optional, cast
+from typing import Optional
 
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_manager import ModelInstance
 from core.model_runtime.entities.message_entities import PromptMessage
 from core.model_runtime.entities.model_entities import ModelPropertyKey
-from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.prompt.entities.advanced_prompt_entities import MemoryConfig
 
 
@@ -25,12 +25,12 @@ class PromptTransform:
 
         model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
         if model_context_tokens:
-            model_type_instance = model_config.provider_model_bundle.model_type_instance
-            model_type_instance = cast(LargeLanguageModel, model_type_instance)
+            model_instance = ModelInstance(
+                provider_model_bundle=model_config.provider_model_bundle,
+                model=model_config.model
+            )
 
-            curr_message_tokens = model_type_instance.get_num_tokens(
-                model_config.model,
-                model_config.credentials,
+            curr_message_tokens = model_instance.get_llm_num_tokens(
                 prompt_messages
             )
 

api/core/provider_manager.py

@@ -11,6 +11,8 @@ from core.entities.provider_entities import (
     CustomConfiguration,
     CustomModelConfiguration,
     CustomProviderConfiguration,
+    ModelLoadBalancingConfiguration,
+    ModelSettings,
     QuotaConfiguration,
     SystemConfiguration,
 )
@@ -25,14 +27,18 @@ from core.model_runtime.entities.provider_entities import (
 from core.model_runtime.model_providers import model_provider_factory
 from extensions import ext_hosting_provider
 from extensions.ext_database import db
+from extensions.ext_redis import redis_client
 from models.provider import (
+    LoadBalancingModelConfig,
     Provider,
     ProviderModel,
+    ProviderModelSetting,
     ProviderQuotaType,
     ProviderType,
     TenantDefaultModel,
     TenantPreferredModelProvider,
 )
+from services.feature_service import FeatureService
 
 
 class ProviderManager:
@@ -98,6 +104,13 @@ class ProviderManager:
         # Get All preferred provider types of the workspace
         provider_name_to_preferred_model_provider_records_dict = self._get_all_preferred_model_providers(tenant_id)
 
+        # Get All provider model settings
+        provider_name_to_provider_model_settings_dict = self._get_all_provider_model_settings(tenant_id)
+
+        # Get All load balancing configs
+        provider_name_to_provider_load_balancing_model_configs_dict \
+            = self._get_all_provider_load_balancing_configs(tenant_id)
+
         provider_configurations = ProviderConfigurations(
             tenant_id=tenant_id
         )
@@ -147,13 +160,28 @@ class ProviderManager:
                     if system_configuration.enabled and has_valid_quota:
                         using_provider_type = ProviderType.SYSTEM
 
+            # Get provider load balancing configs
+            provider_model_settings = provider_name_to_provider_model_settings_dict.get(provider_name)
+
+            # Get provider load balancing configs
+            provider_load_balancing_configs \
+                = provider_name_to_provider_load_balancing_model_configs_dict.get(provider_name)
+
+            # Convert to model settings
+            model_settings = self._to_model_settings(
+                provider_entity=provider_entity,
+                provider_model_settings=provider_model_settings,
+                load_balancing_model_configs=provider_load_balancing_configs
+            )
+
             provider_configuration = ProviderConfiguration(
                 tenant_id=tenant_id,
                 provider=provider_entity,
                 preferred_provider_type=preferred_provider_type,
                 using_provider_type=using_provider_type,
                 system_configuration=system_configuration,
-                custom_configuration=custom_configuration
+                custom_configuration=custom_configuration,
+                model_settings=model_settings
             )
 
             provider_configurations[provider_name] = provider_configuration
@@ -338,7 +366,7 @@ class ProviderManager:
         """
         Get All preferred provider types of the workspace.
 
-        :param tenant_id:
+        :param tenant_id: workspace id
         :return:
         """
         preferred_provider_types = db.session.query(TenantPreferredModelProvider) \
@@ -353,6 +381,56 @@ class ProviderManager:
 
         return provider_name_to_preferred_provider_type_records_dict
 
+    def _get_all_provider_model_settings(self, tenant_id: str) -> dict[str, list[ProviderModelSetting]]:
+        """
+        Get All provider model settings of the workspace.
+
+        :param tenant_id: workspace id
+        :return:
+        """
+        provider_model_settings = db.session.query(ProviderModelSetting) \
+            .filter(
+            ProviderModelSetting.tenant_id == tenant_id
+        ).all()
+
+        provider_name_to_provider_model_settings_dict = defaultdict(list)
+        for provider_model_setting in provider_model_settings:
+            (provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name]
+             .append(provider_model_setting))
+
+        return provider_name_to_provider_model_settings_dict
+
+    def _get_all_provider_load_balancing_configs(self, tenant_id: str) -> dict[str, list[LoadBalancingModelConfig]]:
+        """
+        Get All provider load balancing configs of the workspace.
+
+        :param tenant_id: workspace id
+        :return:
+        """
+        cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled"
+        cache_result = redis_client.get(cache_key)
+        if cache_result is None:
+            model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled
+            redis_client.setex(cache_key, 120, str(model_load_balancing_enabled))
+        else:
+            cache_result = cache_result.decode('utf-8')
+            model_load_balancing_enabled = cache_result == 'True'
+
+        if not model_load_balancing_enabled:
+            return dict()
+
+        provider_load_balancing_configs = db.session.query(LoadBalancingModelConfig) \
+            .filter(
+            LoadBalancingModelConfig.tenant_id == tenant_id
+        ).all()
+
+        provider_name_to_provider_load_balancing_model_configs_dict = defaultdict(list)
+        for provider_load_balancing_config in provider_load_balancing_configs:
+            (provider_name_to_provider_load_balancing_model_configs_dict[provider_load_balancing_config.provider_name]
+             .append(provider_load_balancing_config))
+
+        return provider_name_to_provider_load_balancing_model_configs_dict
+
     def _init_trial_provider_records(self, tenant_id: str,
                                      provider_name_to_provider_records_dict: dict[str, list]) -> dict[str, list]:
         """
@@ -726,3 +804,97 @@ class ProviderManager:
                 secret_input_form_variables.append(credential_form_schema.variable)
 
         return secret_input_form_variables
+
+    def _to_model_settings(self, provider_entity: ProviderEntity,
+                           provider_model_settings: Optional[list[ProviderModelSetting]] = None,
+                           load_balancing_model_configs: Optional[list[LoadBalancingModelConfig]] = None) \
+            -> list[ModelSettings]:
+        """
+        Convert to model settings.
+
+        :param provider_model_settings: provider model settings include enabled, load balancing enabled
+        :param load_balancing_model_configs: load balancing model configs
+        :return:
+        """
+        # Get provider model credential secret variables
+        model_credential_secret_variables = self._extract_secret_variables(
+            provider_entity.model_credential_schema.credential_form_schemas
+            if provider_entity.model_credential_schema else []
+        )
+
+        model_settings = []
+        if not provider_model_settings:
+            return model_settings
+
+        for provider_model_setting in provider_model_settings:
+            load_balancing_configs = []
+            if provider_model_setting.load_balancing_enabled and load_balancing_model_configs:
+                for load_balancing_model_config in load_balancing_model_configs:
+                    if (load_balancing_model_config.model_name == provider_model_setting.model_name
+                            and load_balancing_model_config.model_type == provider_model_setting.model_type):
+                        if not load_balancing_model_config.enabled:
+                            continue
+
+                        if not load_balancing_model_config.encrypted_config:
+                            if load_balancing_model_config.name == "__inherit__":
+                                load_balancing_configs.append(ModelLoadBalancingConfiguration(
+                                    id=load_balancing_model_config.id,
+                                    name=load_balancing_model_config.name,
+                                    credentials={}
+                                ))
+                            continue
+
+                        provider_model_credentials_cache = ProviderCredentialsCache(
+                            tenant_id=load_balancing_model_config.tenant_id,
+                            identity_id=load_balancing_model_config.id,
+                            cache_type=ProviderCredentialsCacheType.LOAD_BALANCING_MODEL
+                        )
+
+                        # Get cached provider model credentials
+                        cached_provider_model_credentials = provider_model_credentials_cache.get()
+
+                        if not cached_provider_model_credentials:
+                            try:
+                                provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config)
+                            except JSONDecodeError:
+                                continue
+
+                            # Get decoding rsa key and cipher for decrypting credentials
+                            if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
+                                self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(
+                                    load_balancing_model_config.tenant_id)
+
+                            for variable in model_credential_secret_variables:
+                                if variable in provider_model_credentials:
+                                    try:
+                                        provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
+                                            provider_model_credentials.get(variable),
+                                            self.decoding_rsa_key,
+                                            self.decoding_cipher_rsa
+                                        )
+                                    except ValueError:
+                                        pass
+
+                            # cache provider model credentials
+                            provider_model_credentials_cache.set(
+                                credentials=provider_model_credentials
+                            )
+                        else:
+                            provider_model_credentials = cached_provider_model_credentials
+
+                        load_balancing_configs.append(ModelLoadBalancingConfiguration(
+                            id=load_balancing_model_config.id,
+                            name=load_balancing_model_config.name,
+                            credentials=provider_model_credentials
+                        ))
+
+            model_settings.append(
+                ModelSettings(
+                    model=provider_model_setting.model_name,
+                    model_type=ModelType.value_of(provider_model_setting.model_type),
+                    enabled=provider_model_setting.enabled,
+                    load_balancing_configs=load_balancing_configs if len(load_balancing_configs) > 1 else []
+                )
+            )
+
+        return model_settings

api/core/rag/data_post_processor/data_post_processor.py

@@ -5,7 +5,7 @@ from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.rag.data_post_processor.reorder import ReorderRunner
 from core.rag.models.document import Document
-from core.rerank.rerank import RerankRunner
+from core.rag.rerank.rerank import RerankRunner
 
 
 class DataPostProcessor:

api/core/rag/datasource/retrieval_service.py

@@ -33,6 +33,7 @@ class RetrievalService:
             return []
         all_documents = []
         threads = []
+        exceptions = []
         # retrieval_model source with keyword
         if retrival_method == 'keyword_search':
             keyword_thread = threading.Thread(target=RetrievalService.keyword_search, kwargs={
@@ -40,7 +41,8 @@ class RetrievalService:
                 'dataset_id': dataset_id,
                 'query': query,
                 'top_k': top_k,
-                'all_documents': all_documents
+                'all_documents': all_documents,
+                'exceptions': exceptions,
             })
             threads.append(keyword_thread)
             keyword_thread.start()
@@ -54,7 +56,8 @@ class RetrievalService:
                 'score_threshold': score_threshold,
                 'reranking_model': reranking_model,
                 'all_documents': all_documents,
-                'retrival_method': retrival_method
+                'retrival_method': retrival_method,
+                'exceptions': exceptions,
             })
             threads.append(embedding_thread)
             embedding_thread.start()
@@ -69,7 +72,8 @@ class RetrievalService:
                 'score_threshold': score_threshold,
                 'top_k': top_k,
                 'reranking_model': reranking_model,
-                'all_documents': all_documents
+                'all_documents': all_documents,
+                'exceptions': exceptions,
             })
             threads.append(full_text_index_thread)
             full_text_index_thread.start()
@@ -77,6 +81,10 @@ class RetrievalService:
         for thread in threads:
             thread.join()
 
+        if exceptions:
+            exception_message = ';\n'.join(exceptions)
+            raise Exception(exception_message)
+
         if retrival_method == 'hybrid_search':
             data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
             all_documents = data_post_processor.invoke(
@@ -89,82 +97,91 @@ class RetrievalService:
 
     @classmethod
     def keyword_search(cls, flask_app: Flask, dataset_id: str, query: str,
-                       top_k: int, all_documents: list):
+                       top_k: int, all_documents: list, exceptions: list):
         with flask_app.app_context():
-            dataset = db.session.query(Dataset).filter(
-                Dataset.id == dataset_id
-            ).first()
+            try:
+                dataset = db.session.query(Dataset).filter(
+                    Dataset.id == dataset_id
+                ).first()
 
-            keyword = Keyword(
-                dataset=dataset
-            )
+                keyword = Keyword(
+                    dataset=dataset
+                )
 
-            documents = keyword.search(
-                query,
-                top_k=top_k
-            )
-            all_documents.extend(documents)
+                documents = keyword.search(
+                    query,
+                    top_k=top_k
+                )
+                all_documents.extend(documents)
+            except Exception as e:
+                exceptions.append(str(e))
 
     @classmethod
     def embedding_search(cls, flask_app: Flask, dataset_id: str, query: str,
                          top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
-                         all_documents: list, retrival_method: str):
+                         all_documents: list, retrival_method: str, exceptions: list):
         with flask_app.app_context():
-            dataset = db.session.query(Dataset).filter(
-                Dataset.id == dataset_id
-            ).first()
+            try:
+                dataset = db.session.query(Dataset).filter(
+                    Dataset.id == dataset_id
+                ).first()
 
-            vector = Vector(
-                dataset=dataset
-            )
+                vector = Vector(
+                    dataset=dataset
+                )
 
-            documents = vector.search_by_vector(
-                query,
-                search_type='similarity_score_threshold',
-                top_k=top_k,
-                score_threshold=score_threshold,
-                filter={
-                    'group_id': [dataset.id]
-                }
-            )
+                documents = vector.search_by_vector(
+                    query,
+                    search_type='similarity_score_threshold',
+                    top_k=top_k,
+                    score_threshold=score_threshold,
+                    filter={
+                        'group_id': [dataset.id]
+                    }
+                )
 
-            if documents:
-                if reranking_model and retrival_method == 'semantic_search':
-                    data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
-                    all_documents.extend(data_post_processor.invoke(
-                        query=query,
-                        documents=documents,
-                        score_threshold=score_threshold,
-                        top_n=len(documents)
-                    ))
-                else:
-                    all_documents.extend(documents)
+                if documents:
+                    if reranking_model and retrival_method == 'semantic_search':
+                        data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
+                        all_documents.extend(data_post_processor.invoke(
+                            query=query,
+                            documents=documents,
+                            score_threshold=score_threshold,
+                            top_n=len(documents)
+                        ))
+                    else:
+                        all_documents.extend(documents)
+            except Exception as e:
+                exceptions.append(str(e))
 
     @classmethod
     def full_text_index_search(cls, flask_app: Flask, dataset_id: str, query: str,
                                top_k: int, score_threshold: Optional[float], reranking_model: Optional[dict],
-                               all_documents: list, retrival_method: str):
+                               all_documents: list, retrival_method: str, exceptions: list):
         with flask_app.app_context():
-            dataset = db.session.query(Dataset).filter(
-                Dataset.id == dataset_id
-            ).first()
+            try:
+                dataset = db.session.query(Dataset).filter(
+                    Dataset.id == dataset_id
+                ).first()
 
-            vector_processor = Vector(
-                dataset=dataset,
-            )
+                vector_processor = Vector(
+                    dataset=dataset,
+                )
 
-            documents = vector_processor.search_by_full_text(
-                query,
-                top_k=top_k
-            )
-            if documents:
-                if reranking_model and retrival_method == 'full_text_search':
-                    data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
-                    all_documents.extend(data_post_processor.invoke(
-                        query=query,
-                        documents=documents,
-                        score_threshold=score_threshold,
-                        top_n=len(documents)
-                    ))
-                else:
-                    all_documents.extend(documents)
+                documents = vector_processor.search_by_full_text(
+                    query,
+                    top_k=top_k
+                )
+                if documents:
+                    if reranking_model and retrival_method == 'full_text_search':
+                        data_post_processor = DataPostProcessor(str(dataset.tenant_id), reranking_model, False)
+                        all_documents.extend(data_post_processor.invoke(
+                            query=query,
+                            documents=documents,
+                            score_threshold=score_threshold,
+                            top_n=len(documents)
+                        ))
+                    else:
+                        all_documents.extend(documents)
+            except Exception as e:
+                exceptions.append(str(e))

api/core/rag/docstore/dataset_docstore.py

@@ -1,11 +1,10 @@
 from collections.abc import Sequence
-from typing import Any, Optional, cast
+from typing import Any, Optional
 
 from sqlalchemy import func
 
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
-from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from core.rag.models.document import Document
 from extensions.ext_database import db
 from models.dataset import Dataset, DocumentSegment
@@ -95,11 +94,7 @@ class DatasetDocumentStore:
 
             # calc embedding use tokens
             if embedding_model:
-                model_type_instance = embedding_model.model_type_instance
-                model_type_instance = cast(TextEmbeddingModel, model_type_instance)
-                tokens = model_type_instance.get_num_tokens(
-                    model=embedding_model.model,
-                    credentials=embedding_model.credentials,
+                tokens = embedding_model.get_text_embedding_num_tokens(
                     texts=[doc.page_content]
                 )
             else:

api/core/rag/extractor/csv_extractor.py

@@ -57,7 +57,7 @@ class CSVExtractor(BaseExtractor):
         docs = []
         try:
             # load csv file into pandas dataframe
-            df = pd.read_csv(csvfile, error_bad_lines=False, **self.csv_args)
+            df = pd.read_csv(csvfile, on_bad_lines='skip', **self.csv_args)
 
             # check source column exists
             if self.source_column and self.source_column not in df.columns:

api/core/rag/extractor/extract_processor.py

@@ -29,7 +29,7 @@ from core.rag.models.document import Document
 from extensions.ext_storage import storage
 from models.model import UploadFile
 
-SUPPORT_URL_CONTENT_TYPES = ['application/pdf', 'text/plain']
+SUPPORT_URL_CONTENT_TYPES = ['application/pdf', 'text/plain', 'application/json']
 USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
 
 

api/core/rag/index_processor/index_processor_base.py

@@ -7,8 +7,11 @@ from flask import current_app
 from core.model_manager import ModelInstance
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.models.document import Document
-from core.splitter.fixed_text_splitter import EnhanceRecursiveCharacterTextSplitter, FixedRecursiveCharacterTextSplitter
-from core.splitter.text_splitter import TextSplitter
+from core.rag.splitter.fixed_text_splitter import (
+    EnhanceRecursiveCharacterTextSplitter,
+    FixedRecursiveCharacterTextSplitter,
+)
+from core.rag.splitter.text_splitter import TextSplitter
 from models.dataset import Dataset, DatasetProcessRule
 
 

api/core/rag/retrieval/dataset_retrieval.py

@@ -14,9 +14,9 @@ from core.model_runtime.entities.model_entities import ModelFeature, ModelType
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.models.document import Document
+from core.rag.rerank.rerank import RerankRunner
 from core.rag.retrieval.router.multi_dataset_function_call_router import FunctionCallMultiDatasetRouter
 from core.rag.retrieval.router.multi_dataset_react_route import ReactMultiDatasetRouter
-from core.rerank.rerank import RerankRunner
 from core.tools.tool.dataset_retriever.dataset_multi_retriever_tool import DatasetMultiRetrieverTool
 from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
 from core.tools.tool.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool

api/core/rag/splitter/fixed_text_splitter.py

@@ -1,12 +1,11 @@
 """Functionality for splitting text."""
 from __future__ import annotations
 
-from typing import Any, Optional, cast
+from typing import Any, Optional
 
 from core.model_manager import ModelInstance
-from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
 from core.model_runtime.model_providers.__base.tokenizers.gpt2_tokenzier import GPT2Tokenizer
-from core.splitter.text_splitter import (
+from core.rag.splitter.text_splitter import (
     TS,
     Collection,
     Literal,
@@ -35,11 +34,7 @@ class EnhanceRecursiveCharacterTextSplitter(RecursiveCharacterTextSplitter):
                 return 0
 
             if embedding_model_instance:
-                embedding_model_type_instance = embedding_model_instance.model_type_instance
-                embedding_model_type_instance = cast(TextEmbeddingModel, embedding_model_type_instance)
-                return embedding_model_type_instance.get_num_tokens(
-                    model=embedding_model_instance.model,
-                    credentials=embedding_model_instance.credentials,
+                return embedding_model_instance.get_text_embedding_num_tokens(
                     texts=[text]
                 )
             else:

api/core/rag/splitter/text_splitter.py

@@ -6,7 +6,6 @@ import re
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Collection, Iterable, Sequence, Set
 from dataclasses import dataclass
-from enum import Enum
 from typing import (
     Any,
     Literal,
@@ -477,27 +476,6 @@ class TokenTextSplitter(TextSplitter):
         return split_text_on_tokens(text=text, tokenizer=tokenizer)
 
 
-class Language(str, Enum):
-    """Enum of the programming languages."""
-
-    CPP = "cpp"
-    GO = "go"
-    JAVA = "java"
-    JS = "js"
-    PHP = "php"
-    PROTO = "proto"
-    PYTHON = "python"
-    RST = "rst"
-    RUBY = "ruby"
-    RUST = "rust"
-    SCALA = "scala"
-    SWIFT = "swift"
-    MARKDOWN = "markdown"
-    LATEX = "latex"
-    HTML = "html"
-    SOL = "sol"
-
-
 class RecursiveCharacterTextSplitter(TextSplitter):
     """Splitting text by recursively look at characters.
 
@@ -554,350 +532,3 @@ class RecursiveCharacterTextSplitter(TextSplitter):
 
     def split_text(self, text: str) -> list[str]:
         return self._split_text(text, self._separators)
-
-    @classmethod
-    def from_language(
-            cls, language: Language, **kwargs: Any
-    ) -> RecursiveCharacterTextSplitter:
-        separators = cls.get_separators_for_language(language)
-        return cls(separators=separators, **kwargs)
-
-    @staticmethod
-    def get_separators_for_language(language: Language) -> list[str]:
-        if language == Language.CPP:
-            return [
-                # Split along class definitions
-                "\nclass ",
-                # Split along function definitions
-                "\nvoid ",
-                "\nint ",
-                "\nfloat ",
-                "\ndouble ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\nswitch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.GO:
-            return [
-                # Split along function definitions
-                "\nfunc ",
-                "\nvar ",
-                "\nconst ",
-                "\ntype ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nswitch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.JAVA:
-            return [
-                # Split along class definitions
-                "\nclass ",
-                # Split along method definitions
-                "\npublic ",
-                "\nprotected ",
-                "\nprivate ",
-                "\nstatic ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\nswitch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.JS:
-            return [
-                # Split along function definitions
-                "\nfunction ",
-                "\nconst ",
-                "\nlet ",
-                "\nvar ",
-                "\nclass ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\nswitch ",
-                "\ncase ",
-                "\ndefault ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.PHP:
-            return [
-                # Split along function definitions
-                "\nfunction ",
-                # Split along class definitions
-                "\nclass ",
-                # Split along control flow statements
-                "\nif ",
-                "\nforeach ",
-                "\nwhile ",
-                "\ndo ",
-                "\nswitch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.PROTO:
-            return [
-                # Split along message definitions
-                "\nmessage ",
-                # Split along service definitions
-                "\nservice ",
-                # Split along enum definitions
-                "\nenum ",
-                # Split along option definitions
-                "\noption ",
-                # Split along import statements
-                "\nimport ",
-                # Split along syntax declarations
-                "\nsyntax ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.PYTHON:
-            return [
-                # First, try to split along class definitions
-                "\nclass ",
-                "\ndef ",
-                "\n\tdef ",
-                # Now split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.RST:
-            return [
-                # Split along section titles
-                "\n=+\n",
-                "\n-+\n",
-                "\n\*+\n",
-                # Split along directive markers
-                "\n\n.. *\n\n",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.RUBY:
-            return [
-                # Split along method definitions
-                "\ndef ",
-                "\nclass ",
-                # Split along control flow statements
-                "\nif ",
-                "\nunless ",
-                "\nwhile ",
-                "\nfor ",
-                "\ndo ",
-                "\nbegin ",
-                "\nrescue ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.RUST:
-            return [
-                # Split along function definitions
-                "\nfn ",
-                "\nconst ",
-                "\nlet ",
-                # Split along control flow statements
-                "\nif ",
-                "\nwhile ",
-                "\nfor ",
-                "\nloop ",
-                "\nmatch ",
-                "\nconst ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.SCALA:
-            return [
-                # Split along class definitions
-                "\nclass ",
-                "\nobject ",
-                # Split along method definitions
-                "\ndef ",
-                "\nval ",
-                "\nvar ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\nmatch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.SWIFT:
-            return [
-                # Split along function definitions
-                "\nfunc ",
-                # Split along class definitions
-                "\nclass ",
-                "\nstruct ",
-                "\nenum ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\ndo ",
-                "\nswitch ",
-                "\ncase ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.MARKDOWN:
-            return [
-                # First, try to split along Markdown headings (starting with level 2)
-                "\n#{1,6} ",
-                # Note the alternative syntax for headings (below) is not handled here
-                # Heading level 2
-                # ---------------
-                # End of code block
-                "```\n",
-                # Horizontal lines
-                "\n\*\*\*+\n",
-                "\n---+\n",
-                "\n___+\n",
-                # Note that this splitter doesn't handle horizontal lines defined
-                # by *three or more* of ***, ---, or ___, but this is not handled
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        elif language == Language.LATEX:
-            return [
-                # First, try to split along Latex sections
-                "\n\\\chapter{",
-                "\n\\\section{",
-                "\n\\\subsection{",
-                "\n\\\subsubsection{",
-                # Now split by environments
-                "\n\\\begin{enumerate}",
-                "\n\\\begin{itemize}",
-                "\n\\\begin{description}",
-                "\n\\\begin{list}",
-                "\n\\\begin{quote}",
-                "\n\\\begin{quotation}",
-                "\n\\\begin{verse}",
-                "\n\\\begin{verbatim}",
-                # Now split by math environments
-                "\n\\\begin{align}",
-                "$$",
-                "$",
-                # Now split by the normal type of lines
-                " ",
-                "",
-            ]
-        elif language == Language.HTML:
-            return [
-                # First, try to split along HTML tags
-                "<body",
-                "<div",
-                "<p",
-                "<br",
-                "<li",
-                "<h1",
-                "<h2",
-                "<h3",
-                "<h4",
-                "<h5",
-                "<h6",
-                "<span",
-                "<table",
-                "<tr",
-                "<td",
-                "<th",
-                "<ul",
-                "<ol",
-                "<header",
-                "<footer",
-                "<nav",
-                # Head
-                "<head",
-                "<style",
-                "<script",
-                "<meta",
-                "<title",
-                "",
-            ]
-        elif language == Language.SOL:
-            return [
-                # Split along compiler information definitions
-                "\npragma ",
-                "\nusing ",
-                # Split along contract definitions
-                "\ncontract ",
-                "\ninterface ",
-                "\nlibrary ",
-                # Split along method definitions
-                "\nconstructor ",
-                "\ntype ",
-                "\nfunction ",
-                "\nevent ",
-                "\nmodifier ",
-                "\nerror ",
-                "\nstruct ",
-                "\nenum ",
-                # Split along control flow statements
-                "\nif ",
-                "\nfor ",
-                "\nwhile ",
-                "\ndo while ",
-                "\nassembly ",
-                # Split by the normal type of lines
-                "\n\n",
-                "\n",
-                " ",
-                "",
-            ]
-        else:
-            raise ValueError(
-                f"Language {language} is not supported! "
-                f"Please choose from {list(Language)}"
-            )

api/core/tools/docs/en_US/tool_scale_out.md

@@ -24,10 +24,32 @@ identity: # Basic information of the tool provider
     en_US: Google # English description
     zh_Hans: Google # Chinese description
   icon: icon.svg # Icon, needs to be placed in the _assets folder of the current module
+  tags:
+    - search
 
 ```
  - The `identity` field is mandatory, it contains the basic information of the tool provider, including author, name, label, description, icon, etc.
     - The icon needs to be placed in the `_assets` folder of the current module, you can refer to [here](../../provider/builtin/google/_assets/icon.svg).
+    - The `tags` field is optional, it is used to classify the provider, and the frontend can filter the provider according to the tag, for all tags, they have been listed below:
+    ```python
+    class ToolLabelEnum(Enum):
+      SEARCH = 'search'
+      IMAGE = 'image'
+      VIDEOS = 'videos'
+      WEATHER = 'weather'
+      FINANCE = 'finance'
+      DESIGN = 'design'
+      TRAVEL = 'travel'
+      SOCIAL = 'social'
+      NEWS = 'news'
+      MEDICAL = 'medical'
+      PRODUCTIVITY = 'productivity'
+      EDUCATION = 'education'
+      BUSINESS = 'business'
+      ENTERTAINMENT = 'entertainment'
+      UTILITIES = 'utilities'
+      OTHER = 'other'
+    ```
 
 ## 2. Prepare Provider Credentials
 

api/core/tools/docs/zh_Hans/tool_scale_out.md

@@ -24,10 +24,32 @@ identity: # 工具供应商的基本信息
     en_US: Google # 英文描述
     zh_Hans: Google # 中文描述
   icon: icon.svg # 图标，需要放置在当前模块的_assets文件夹下
+  tags: # 标签，用于前端展示
+    - search
 
 ```
  - `identity` 字段是必须的，它包含了工具供应商的基本信息，包括作者、名称、标签、描述、图标等
     - 图标需要放置在当前模块的`_assets`文件夹下，可以参考[这里](../../provider/builtin/google/_assets/icon.svg)。
+    - 标签用于前端展示，可以帮助用户快速找到这个工具供应商，下面列出了目前所支持的所有标签
+    ```python
+    class ToolLabelEnum(Enum):
+      SEARCH = 'search'
+      IMAGE = 'image'
+      VIDEOS = 'videos'
+      WEATHER = 'weather'
+      FINANCE = 'finance'
+      DESIGN = 'design'
+      TRAVEL = 'travel'
+      SOCIAL = 'social'
+      NEWS = 'news'
+      MEDICAL = 'medical'
+      PRODUCTIVITY = 'productivity'
+      EDUCATION = 'education'
+      BUSINESS = 'business'
+      ENTERTAINMENT = 'entertainment'
+      UTILITIES = 'utilities'
+      OTHER = 'other'
+    ```
 
 ## 2. 准备供应商凭据
 

api/core/tools/entities/tool_entities.py

@@ -6,6 +6,24 @@ from pydantic import BaseModel, Field
 from core.tools.entities.common_entities import I18nObject
 
 
+class ToolLabelEnum(Enum):
+    SEARCH = 'search'
+    IMAGE = 'image'
+    VIDEOS = 'videos'
+    WEATHER = 'weather'
+    FINANCE = 'finance'
+    DESIGN = 'design'
+    TRAVEL = 'travel'
+    SOCIAL = 'social'
+    NEWS = 'news'
+    MEDICAL = 'medical'
+    PRODUCTIVITY = 'productivity'
+    EDUCATION = 'education'
+    BUSINESS = 'business'
+    ENTERTAINMENT = 'entertainment'
+    UTILITIES = 'utilities'
+    OTHER = 'other'
+
 class ToolProviderType(Enum):
     """
         Enum class for tool provider
@@ -98,8 +116,9 @@ class ToolParameterOption(BaseModel):
     value: str = Field(..., description="The value of the option")
     label: I18nObject = Field(..., description="The label of the option")
 
+
 class ToolParameter(BaseModel):
-    class ToolParameterType(Enum):
+    class ToolParameterType(str, Enum):
         STRING = "string"
         NUMBER = "number"
         BOOLEAN = "boolean"
@@ -157,6 +176,7 @@ class ToolProviderIdentity(BaseModel):
     description: I18nObject = Field(..., description="The description of the tool")
     icon: str = Field(..., description="The icon of the tool")
     label: I18nObject = Field(..., description="The label of the tool")
+    tags: Optional[list[ToolLabelEnum]] = Field(default=[], description="The tags of the tool", )
 
 class ToolDescription(BaseModel):
     human: I18nObject = Field(..., description="The description presented to the user")

api/core/tools/entities/values.py

@@ -1,26 +1,5 @@
-from enum import Enum
-
 from core.tools.entities.common_entities import I18nObject
-from core.tools.entities.tool_entities import ToolLabel
-
-
-class ToolLabelEnum(Enum):
-    SEARCH = 'search'
-    IMAGE = 'image'
-    VIDEOS = 'videos'
-    WEATHER = 'weather'
-    FINANCE = 'finance'
-    DESIGN = 'design'
-    TRAVEL = 'travel'
-    SOCIAL = 'social'
-    NEWS = 'news'
-    MEDICAL = 'medical'
-    PRODUCTIVITY = 'productivity'
-    EDUCATION = 'education'
-    BUSINESS = 'business'
-    ENTERTAINMENT = 'entertainment'
-    UTILITIES = 'utilities'
-    OTHER = 'other'
+from core.tools.entities.tool_entities import ToolLabel, ToolLabelEnum
 
 ICONS = {
     ToolLabelEnum.SEARCH: '''<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 16 16" fill="none">

api/core/tools/provider/_position.yaml

@@ -1,6 +1,7 @@
 - google
 - bing
 - duckduckgo
+- searchapi
 - searxng
 - dalle
 - azuredalle

api/core/tools/provider/builtin/_positions.py

@@ -1,7 +1,7 @@
 import os.path
 
+from core.helper.position_helper import get_position_map, sort_by_position_map
 from core.tools.entities.api_entities import UserToolProvider
-from core.utils.position_helper import get_position_map, sort_by_position_map
 
 
 class BuiltinToolProviderSort:

api/core/tools/provider/builtin/aippt/aippt.py

@@ -1,4 +1,3 @@
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.aippt.tools.aippt import AIPPTGenerateTool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -10,9 +9,3 @@ class AIPPTProvider(BuiltinToolProviderController):
             AIPPTGenerateTool._get_api_token(credentials, user_id='__dify_system__')
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.PRODUCTIVITY, 
-            ToolLabelEnum.DESIGN,
-        ]

api/core/tools/provider/builtin/aippt/aippt.yaml

@@ -8,6 +8,9 @@ identity:
     en_US: AI-generated PPT with one click, input your content topic, and let AI serve you one-stop
     zh_Hans: AI一键生成PPT，输入你的内容主题，让AI为你一站式服务到底
   icon: icon.png
+  tags:
+    - productivity
+    - design
 credentials_for_provider:
   aippt_access_key:
     type: secret-input

api/core/tools/provider/builtin/arxiv/arxiv.py

@@ -1,4 +1,3 @@
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.arxiv.tools.arxiv_search import ArxivSearchTool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -19,8 +18,4 @@ class ArxivProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-        
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.SEARCH,
-        ]
+        

api/core/tools/provider/builtin/arxiv/arxiv.yaml

@@ -8,3 +8,5 @@ identity:
     en_US: Access to a vast repository of scientific papers and articles in various fields of research.
     zh_Hans: 访问各个研究领域大量科学论文和文章的存储库。
   icon: icon.svg
+  tags:
+    - search

api/core/tools/provider/builtin/azuredalle/azuredalle.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.azuredalle.tools.dalle3 import DallE3Tool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -23,8 +22,3 @@ class AzureDALLEProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.IMAGE
-        ]

api/core/tools/provider/builtin/azuredalle/azuredalle.yaml

@@ -10,6 +10,9 @@ identity:
     zh_Hans: Azure DALL-E 绘画
     pt_BR: Azure DALL-E art
   icon: icon.png
+  tags:
+    - image
+    - productivity
 credentials_for_provider:
   azure_openai_api_key:
     type: secret-input

api/core/tools/provider/builtin/bing/bing.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.bing.tools.bing_web_search import BingSearchTool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -22,8 +21,3 @@ class BingProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.SEARCH
-        ]

api/core/tools/provider/builtin/bing/bing.yaml

@@ -10,6 +10,8 @@ identity:
     zh_Hans: Bing 搜索
     pt_BR: Bing Search
   icon: icon.svg
+  tags:
+    - search
 credentials_for_provider:
   subscription_key:
     type: secret-input

api/core/tools/provider/builtin/brave/brave.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.brave.tools.brave_search import BraveSearchTool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -21,8 +20,4 @@ class BraveProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-        
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.SEARCH,
-        ]
+        

api/core/tools/provider/builtin/brave/brave.yaml

@@ -10,6 +10,8 @@ identity:
     zh_Hans: Brave
     pt_BR: Brave
   icon: icon.svg
+  tags:
+    - search
 credentials_for_provider:
   brave_search_api_key:
     type: secret-input

api/core/tools/provider/builtin/chart/chart.py

@@ -2,7 +2,6 @@ import matplotlib.pyplot as plt
 from fontTools.ttLib import TTFont
 from matplotlib.font_manager import findSystemFonts
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.chart.tools.line import LinearChartTool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -56,8 +55,4 @@ class ChartProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-        
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.DESIGN, ToolLabelEnum.PRODUCTIVITY, ToolLabelEnum.UTILITIES
-        ]
+        

api/core/tools/provider/builtin/chart/chart.yaml

@@ -10,4 +10,8 @@ identity:
     zh_Hans: 图表生成是一个用于生成可视化图表的工具，你可以通过它来生成柱状图、折线图、饼图等各类图表
     pt_BR: O Gerador de gráficos é uma ferramenta para gerar gráficos estatísticos como gráfico de barras, gráfico de linhas, gráfico de pizza, etc.
   icon: icon.png
+  tags:
+    - design
+    - productivity
+    - utilities
 credentials_for_provider:

api/core/tools/provider/builtin/chart/tools/bar.yaml

@@ -21,9 +21,9 @@ parameters:
       zh_Hans: 数据
       pt_BR: dados
     human_description:
-      en_US: data for generating bar chart
-      zh_Hans: 用于生成柱状图的数据
-      pt_BR: dados para gerar gráfico de barras
+      en_US: data for generating chart, each number should be separated by ";"
+      zh_Hans: 用于生成柱状图的数据，每个数字之间用 ";" 分隔
+      pt_BR: dados para gerar gráfico de barras, cada número deve ser separado por ";"
     llm_description: data for generating bar chart, data should be a string contains a list of numbers like "1;2;3;4;5"
     form: llm
   - name: x_axis
@@ -34,8 +34,8 @@ parameters:
       zh_Hans: x 轴
       pt_BR: Eixo X
     human_description:
-      en_US: X axis for bar chart
-      zh_Hans: 柱状图的 x 轴
-      pt_BR: Eixo X para gráfico de barras
+      en_US: X axis for chart, each text should be separated by ";"
+      zh_Hans: 柱状图的 x 轴，每个文本之间用 ";" 分隔
+      pt_BR: Eixo X para gráfico de barras, cada texto deve ser separado por ";"
     llm_description: x axis for bar chart, x axis should be a string contains a list of texts like "a;b;c;1;2" in order to match the data
     form: llm

api/core/tools/provider/builtin/chart/tools/line.yaml

@@ -21,9 +21,9 @@ parameters:
       zh_Hans: 数据
       pt_BR: dados
     human_description:
-      en_US: data for generating linear chart
-      zh_Hans: 用于生成线性图表的数据
-      pt_BR: dados para gerar gráfico linear
+      en_US: data for generating chart, each number should be separated by ";"
+      zh_Hans: 用于生成线性图表的数据，每个数字之间用 ";" 分隔
+      pt_BR: dados para gerar gráfico linear, cada número deve ser separado por ";"
     llm_description: data for generating linear chart, data should be a string contains a list of numbers like "1;2;3;4;5"
     form: llm
   - name: x_axis
@@ -34,8 +34,8 @@ parameters:
       zh_Hans: x 轴
       pt_BR: Eixo X
     human_description:
-      en_US: X axis for linear chart
-      zh_Hans: 线性图表的 x 轴
-      pt_BR: Eixo X para gráfico linear
+      en_US: X axis for chart, each text should be separated by ";"
+      zh_Hans: 线性图表的 x 轴，每个文本之间用 ";" 分隔
+      pt_BR: Eixo X para gráfico linear, cada texto deve ser separado por ";"
     llm_description: x axis for linear chart, x axis should be a string contains a list of texts like "a;b;c;1;2" in order to match the data
     form: llm

api/core/tools/provider/builtin/chart/tools/pie.yaml

@@ -21,9 +21,9 @@ parameters:
       zh_Hans: 数据
       pt_BR: dados
     human_description:
-      en_US: data for generating pie chart
-      zh_Hans: 用于生成饼图的数据
-      pt_BR: dados para gerar gráfico de pizza
+      en_US: data for generating chart, each number should be separated by ";"
+      zh_Hans: 用于生成饼图的数据，每个数字之间用 ";" 分隔
+      pt_BR: dados para gerar gráfico de pizza, cada número deve ser separado por ";"
     llm_description: data for generating pie chart, data should be a string contains a list of numbers like "1;2;3;4;5"
     form: llm
   - name: categories
@@ -34,8 +34,8 @@ parameters:
       zh_Hans: 分类
       pt_BR: Categorias
     human_description:
-      en_US: Categories for pie chart
-      zh_Hans: 饼图的分类
-      pt_BR: Categorias para gráfico de pizza
+      en_US: Categories for chart, each category should be separated by ";"
+      zh_Hans: 饼图的分类，每个分类之间用 ";" 分隔
+      pt_BR: Categorias para gráfico de pizza, cada categoria deve ser separada por ";"
     llm_description: categories for pie chart, categories should be a string contains a list of texts like "a;b;c;1;2" in order to match the data, each category should be split by ";"
     form: llm

api/core/tools/provider/builtin/code/code.py

@@ -1,14 +1,8 @@
 from typing import Any
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
 
 
 class CodeToolProvider(BuiltinToolProviderController):
     def _validate_credentials(self, credentials: dict[str, Any]) -> None:
         pass
-
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.PRODUCTIVITY
-        ]

api/core/tools/provider/builtin/code/code.yaml

@@ -10,4 +10,6 @@ identity:
     zh_Hans: 运行一段代码并返回结果。
     pt_BR: Execute um trecho de código e obtenha o resultado de volta.
   icon: icon.svg
+  tags:
+    - productivity
 credentials_for_provider:

api/core/tools/provider/builtin/dalle/dalle.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from core.tools.entities.values import ToolLabelEnum
 from core.tools.errors import ToolProviderCredentialValidationError
 from core.tools.provider.builtin.dalle.tools.dalle2 import DallE2Tool
 from core.tools.provider.builtin_tool_provider import BuiltinToolProviderController
@@ -23,8 +22,4 @@ class DALLEProvider(BuiltinToolProviderController):
             )
         except Exception as e:
             raise ToolProviderCredentialValidationError(str(e))
-        
-    def _get_tool_labels(self) -> list[ToolLabelEnum]:
-        return [
-            ToolLabelEnum.IMAGE, ToolLabelEnum.PRODUCTIVITY
-        ]
+