fix error weaviate vector

Co-authored-by: jyong <jyong@dify.ai>

CONTRIBUTING.md

@@ -53,9 +53,9 @@ Did you have an issue, like a merge conflict, or don't know how to open a pull r
 
 ## Community channels
 
-Stuck somewhere? Have any questions? Join the [Discord Community Server](https://discord.gg/AhzKf7dNgk). We are here to help!
+Stuck somewhere? Have any questions? Join the [Discord Community Server](https://discord.gg/j3XRWSPBf7). We are here to help!
 
 ### i18n (Internationalization) Support
 
 We are looking for contributors to help with translations in other languages. If you are interested in helping, please join the [Discord Community Server](https://discord.gg/AhzKf7dNgk) and let us know.  
-Also check out the [Frontend i18n README]((web/i18n/README_EN.md)) for more information.
+Also check out the [Frontend i18n README]((web/i18n/README_EN.md)) for more information.

CONTRIBUTING_CN.md

@@ -16,15 +16,15 @@
 
 ## 本地开发
 
-要设置一个可工作的开发环境，只需 fork 项目的 git 存储库，并使用适当的软件包管理器安装后端和前端依赖项，然后创建并运行 docker-compose 堆栈。
+要设置一个可工作的开发环境，只需 fork 项目的 git 存储库，并使用适当的软件包管理器安装后端和前端依赖项，然后创建并运行 docker-compose。
 
 ### Fork存储库
 
-您需要 fork [存储库](https://github.com/langgenius/dify)。
+您需要 fork [Git 仓库](https://github.com/langgenius/dify)。
 
 ### 克隆存储库
 
-克隆您在 GitHub 上 fork 的存储库：
+克隆您在 GitHub 上 fork 的仓库：
 
 ```
 git clone git@github.com:<github_username>/dify.git

CONTRIBUTING_JA.md

@@ -52,4 +52,4 @@ git clone git@github.com:<github_username>/dify.git
 
 ## コミュニティチャンネル
 
-お困りですか？何か質問がありますか？ [Discord Community サーバ](https://discord.gg/AhzKf7dNgk)に参加してください。私たちがお手伝いします！
+お困りですか？何か質問がありますか？ [Discord Community サーバ](https://discord.gg/j3XRWSPBf7) に参加してください。私たちがお手伝いします！

api/Dockerfile

@@ -16,7 +16,7 @@ EXPOSE 5001
 WORKDIR /app/api
 
 RUN apt-get update && \
-    apt-get install -y bash curl wget vim gcc g++ python3-dev libc-dev libffi-dev
+    apt-get install -y bash curl wget vim gcc g++ python3-dev libc-dev libffi-dev nodejs
 
 COPY requirements.txt /app/api/requirements.txt
 

api/commands.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 import math
 import random
 import string
@@ -6,10 +7,16 @@ import time
 
 import click
 from flask import current_app
+from langchain.embeddings import OpenAIEmbeddings
 from werkzeug.exceptions import NotFound
 
+from core.embedding.cached_embedding import CacheEmbedding
 from core.index.index import IndexBuilder
+from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.embedding.openai_embedding import OpenAIEmbedding
+from core.model_providers.models.entity.model_params import ModelType
 from core.model_providers.providers.hosted import hosted_model_providers
+from core.model_providers.providers.openai_provider import OpenAIProvider
 from libs.password import password_pattern, valid_password, hash_password
 from libs.helper import email as email_validate
 from extensions.ext_database import db
@@ -20,7 +27,7 @@ from models.model import Account
 import secrets
 import base64
 
-from models.provider import Provider, ProviderType, ProviderQuotaType
+from models.provider import Provider, ProviderType, ProviderQuotaType, ProviderModel
 
 
 @click.command('reset-password', help='Reset the account password.')
@@ -102,6 +109,7 @@ def reset_encrypt_key_pair():
     tenant.encrypt_public_key = generate_key_pair(tenant.id)
 
     db.session.query(Provider).filter(Provider.provider_type == 'custom').delete()
+    db.session.query(ProviderModel).delete()
     db.session.commit()
 
     click.echo(click.style('Congratulations! '
@@ -295,6 +303,142 @@ def sync_anthropic_hosted_providers():
     click.echo(click.style('Congratulations! Synced {} anthropic hosted providers.'.format(count), fg='green'))
 
 
+@click.command('create-qdrant-indexes', help='Create qdrant indexes.')
+def create_qdrant_indexes():
+    click.echo(click.style('Start create qdrant indexes.', fg='green'))
+    create_count = 0
+
+    page = 1
+    while True:
+        try:
+            datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality') \
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+        except NotFound:
+            break
+
+        page += 1
+        for dataset in datasets:
+            if dataset.index_struct_dict:
+                if dataset.index_struct_dict['type'] != 'qdrant':
+                    try:
+                        click.echo('Create dataset qdrant index: {}'.format(dataset.id))
+                        try:
+                            embedding_model = ModelFactory.get_embedding_model(
+                                tenant_id=dataset.tenant_id,
+                                model_provider_name=dataset.embedding_model_provider,
+                                model_name=dataset.embedding_model
+                            )
+                        except Exception:
+                            try:
+                                embedding_model = ModelFactory.get_embedding_model(
+                                    tenant_id=dataset.tenant_id
+                                )
+                                dataset.embedding_model = embedding_model.name
+                                dataset.embedding_model_provider = embedding_model.model_provider.provider_name
+                            except Exception:
+                                provider = Provider(
+                                    id='provider_id',
+                                    tenant_id=dataset.tenant_id,
+                                    provider_name='openai',
+                                    provider_type=ProviderType.SYSTEM.value,
+                                    encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
+                                    is_valid=True,
+                                )
+                                model_provider = OpenAIProvider(provider=provider)
+                                embedding_model = OpenAIEmbedding(name="text-embedding-ada-002", model_provider=model_provider)
+                        embeddings = CacheEmbedding(embedding_model)
+
+                        from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+                        index = QdrantVectorIndex(
+                            dataset=dataset,
+                            config=QdrantConfig(
+                                endpoint=current_app.config.get('QDRANT_URL'),
+                                api_key=current_app.config.get('QDRANT_API_KEY'),
+                                root_path=current_app.root_path
+                            ),
+                            embeddings=embeddings
+                        )
+                        if index:
+                            index.create_qdrant_dataset(dataset)
+                            index_struct = {
+                                "type": 'qdrant',
+                                "vector_store": {"class_prefix": dataset.index_struct_dict['vector_store']['class_prefix']}
+                            }
+                            dataset.index_struct = json.dumps(index_struct)
+                            db.session.commit()
+                            create_count += 1
+                        else:
+                            click.echo('passed.')
+                    except Exception as e:
+                        click.echo(
+                            click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                        continue
+
+    click.echo(click.style('Congratulations! Create {} dataset indexes.'.format(create_count), fg='green'))
+
+
+@click.command('update-qdrant-indexes', help='Update qdrant indexes.')
+def update_qdrant_indexes():
+    click.echo(click.style('Start Update qdrant indexes.', fg='green'))
+    create_count = 0
+
+    page = 1
+    while True:
+        try:
+            datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality') \
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+        except NotFound:
+            break
+
+        page += 1
+        for dataset in datasets:
+            if dataset.index_struct_dict:
+                if dataset.index_struct_dict['type'] != 'qdrant':
+                    try:
+                        click.echo('Update dataset qdrant index: {}'.format(dataset.id))
+                        try:
+                            embedding_model = ModelFactory.get_embedding_model(
+                                tenant_id=dataset.tenant_id,
+                                model_provider_name=dataset.embedding_model_provider,
+                                model_name=dataset.embedding_model
+                            )
+                        except Exception:
+                            provider = Provider(
+                                id='provider_id',
+                                tenant_id=dataset.tenant_id,
+                                provider_name='openai',
+                                provider_type=ProviderType.CUSTOM.value,
+                                encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
+                                is_valid=True,
+                            )
+                            model_provider = OpenAIProvider(provider=provider)
+                            embedding_model = OpenAIEmbedding(name="text-embedding-ada-002", model_provider=model_provider)
+                        embeddings = CacheEmbedding(embedding_model)
+
+                        from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+                        index = QdrantVectorIndex(
+                            dataset=dataset,
+                            config=QdrantConfig(
+                                endpoint=current_app.config.get('QDRANT_URL'),
+                                api_key=current_app.config.get('QDRANT_API_KEY'),
+                                root_path=current_app.root_path
+                            ),
+                            embeddings=embeddings
+                        )
+                        if index:
+                            index.update_qdrant_dataset(dataset)
+                            create_count += 1
+                        else:
+                            click.echo('passed.')
+                    except Exception as e:
+                        click.echo(
+                            click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                        continue
+
+    click.echo(click.style('Congratulations! Update {} dataset indexes.'.format(create_count), fg='green'))
+
 def register_commands(app):
     app.cli.add_command(reset_password)
     app.cli.add_command(reset_email)
@@ -303,3 +447,5 @@ def register_commands(app):
     app.cli.add_command(recreate_all_dataset_indexes)
     app.cli.add_command(sync_anthropic_hosted_providers)
     app.cli.add_command(clean_unused_dataset_indexes)
+    app.cli.add_command(create_qdrant_indexes)
+    app.cli.add_command(update_qdrant_indexes)

api/config.py

@@ -100,7 +100,7 @@ class Config:
         self.CONSOLE_URL = get_env('CONSOLE_URL')
         self.API_URL = get_env('API_URL')
         self.APP_URL = get_env('APP_URL')
-        self.CURRENT_VERSION = "0.3.14"
+        self.CURRENT_VERSION = "0.3.18"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = "SELF_HOSTED"
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')

api/controllers/console/apikey.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 import flask_restful
 from flask_restful import Resource, fields, marshal_with
 from werkzeug.exceptions import Forbidden

api/controllers/console/app/app.py

@@ -3,7 +3,9 @@ import json
 import logging
 from datetime import datetime
 
-from flask_login import login_required, current_user
+import flask
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with, abort, inputs
 from werkzeug.exceptions import Forbidden
 
@@ -316,7 +318,7 @@ class AppApi(Resource):
 
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
-        
+
         app = _get_app(app_id, current_user.current_tenant_id)
 
         db.session.delete(app)

api/controllers/console/app/audio.py

@@ -2,7 +2,7 @@
 import logging
 
 from flask import request
-from flask_login import login_required
+from core.login.login import login_required
 from werkzeug.exceptions import InternalServerError, NotFound
 
 import services

api/controllers/console/app/completion.py

@@ -5,7 +5,7 @@ from typing import Generator, Union
 
 import flask_login
 from flask import Response, stream_with_context
-from flask_login import login_required
+from core.login.login import login_required
 from werkzeug.exceptions import InternalServerError, NotFound
 
 import services

api/controllers/console/app/conversation.py

@@ -1,7 +1,8 @@
 from datetime import datetime
 
 import pytz
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with
 from flask_restful.inputs import int_range
 from sqlalchemy import or_, func

api/controllers/console/app/generator.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/app/message.py

@@ -3,7 +3,7 @@ import logging
 from typing import Union, Generator
 
 from flask import Response, stream_with_context
-from flask_login import current_user, login_required
+from flask_login import current_user
 from flask_restful import Resource, reqparse, marshal_with, fields
 from flask_restful.inputs import int_range
 from werkzeug.exceptions import InternalServerError, NotFound
@@ -16,6 +16,7 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
+from core.login.login import login_required
 from libs.helper import uuid_value, TimestampField
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from extensions.ext_database import db

api/controllers/console/app/model_config.py

@@ -3,12 +3,13 @@ import json
 
 from flask import request
 from flask_restful import Resource
-from flask_login import login_required, current_user
+from flask_login import current_user
 
 from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from core.login.login import login_required
 from events.app_event import app_model_config_was_updated
 from extensions.ext_database import db
 from models.model import AppModelConfig

api/controllers/console/app/site.py

@@ -1,5 +1,6 @@
 # -*- coding:utf-8 -*-
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 

api/controllers/console/app/statistic.py

@@ -4,7 +4,8 @@ from datetime import datetime
 
 import pytz
 from flask import jsonify
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/auth/data_source_oauth.py

@@ -5,9 +5,12 @@ from typing import Optional
 import flask_login
 import requests
 from flask import request, redirect, current_app, session
-from flask_login import current_user, login_required
+from flask_login import current_user
+
 from flask_restful import Resource
 from werkzeug.exceptions import Forbidden
+
+from core.login.login import login_required
 from libs.oauth_data_source import NotionOAuth
 from controllers.console import api
 from ..setup import setup_required

api/controllers/console/datasets/data_source.py

@@ -3,7 +3,8 @@ import json
 
 from cachetools import TTLCache
 from flask import request, current_app
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, marshal_with, fields, reqparse, marshal
 from werkzeug.exceptions import NotFound
 

api/controllers/console/datasets/datasets.py

@@ -1,6 +1,7 @@
 # -*- coding:utf-8 -*-
 from flask import request
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 import services
@@ -86,13 +87,19 @@ class DatasetListApi(Resource):
         #     raise ProviderNotInitializeError(
         #         f"No Embedding Model available. Please configure a valid provider "
         #         f"in the Settings -> Model Provider.")
-        model_names = [item['model_name'] for item in valid_model_list]
+        model_names = []
+        for valid_model in valid_model_list:
+            model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
         data = marshal(datasets, dataset_detail_fields)
         for item in data:
-            if item['embedding_model'] in model_names:
-                item['embedding_available'] = True
+            if item['indexing_technique'] == 'high_quality':
+                item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
+                if item_model in model_names:
+                    item['embedding_available'] = True
+                else:
+                    item['embedding_available'] = False
             else:
-                item['embedding_available'] = False
+                item['embedding_available'] = True
         response = {
             'data': data,
             'has_more': len(datasets) == limit,
@@ -118,14 +125,6 @@ class DatasetListApi(Resource):
         # The role of the current user in the ta table must be admin or owner
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
 
         try:
             dataset = DatasetService.create_empty_dataset(
@@ -149,20 +148,39 @@ class DatasetApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id_str)
         if dataset is None:
             raise NotFound("Dataset not found.")
-
         try:
             DatasetService.check_dataset_permission(
                 dataset, current_user)
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
-
-        return marshal(dataset, dataset_detail_fields), 200
+        data = marshal(dataset, dataset_detail_fields)
+        # check embedding setting
+        provider_service = ProviderService()
+        # get valid model list
+        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id, ModelType.EMBEDDINGS.value)
+        model_names = []
+        for valid_model in valid_model_list:
+            model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
+        if data['indexing_technique'] == 'high_quality':
+            item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}"
+            if item_model in model_names:
+                data['embedding_available'] = True
+            else:
+                data['embedding_available'] = False
+        else:
+            data['embedding_available'] = True
+        return data, 200
 
     @setup_required
     @login_required
     @account_initialization_required
     def patch(self, dataset_id):
         dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
 
         parser = reqparse.RequestParser()
         parser.add_argument('name', nullable=False,
@@ -250,6 +268,7 @@ class DatasetIndexingEstimateApi(Resource):
         parser = reqparse.RequestParser()
         parser.add_argument('info_list', type=dict, required=True, nullable=True, location='json')
         parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
+        parser.add_argument('indexing_technique', type=str, required=True, nullable=True, location='json')
         parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
         parser.add_argument('dataset_id', type=str, required=False, nullable=False, location='json')
         parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False, location='json')
@@ -271,7 +290,8 @@ class DatasetIndexingEstimateApi(Resource):
             try:
                 response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
                                                                   args['process_rule'], args['doc_form'],
-                                                                  args['doc_language'], args['dataset_id'])
+                                                                  args['doc_language'], args['dataset_id'],
+                                                                  args['indexing_technique'])
             except LLMBadRequestError:
                 raise ProviderNotInitializeError(
                     f"No Embedding Model available. Please configure a valid provider "
@@ -286,7 +306,8 @@ class DatasetIndexingEstimateApi(Resource):
                 response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
                                                                     args['info_list']['notion_info_list'],
                                                                     args['process_rule'], args['doc_form'],
-                                                                    args['doc_language'], args['dataset_id'])
+                                                                    args['doc_language'], args['dataset_id'],
+                                                                    args['indexing_technique'])
             except LLMBadRequestError:
                 raise ProviderNotInitializeError(
                     f"No Embedding Model available. Please configure a valid provider "

api/controllers/console/datasets/datasets_document.py

@@ -3,8 +3,9 @@ import random
 from datetime import datetime
 from typing import List
 
-from flask import request
-from flask_login import login_required, current_user
+from flask import request, current_app
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, fields, marshal, marshal_with, reqparse
 from sqlalchemy import desc, asc
 from werkzeug.exceptions import NotFound, Forbidden
@@ -137,6 +138,10 @@ class GetProcessRuleApi(Resource):
         req_data = request.args
 
         document_id = req_data.get('document_id')
+        
+        # get default rules
+        mode = DocumentService.DEFAULT_RULES['mode']
+        rules = DocumentService.DEFAULT_RULES['rules']
         if document_id:
             # get the latest process rule
             document = Document.query.get_or_404(document_id)
@@ -157,11 +162,9 @@ class GetProcessRuleApi(Resource):
                 order_by(DatasetProcessRule.created_at.desc()). \
                 limit(1). \
                 one_or_none()
-            mode = dataset_process_rule.mode
-            rules = dataset_process_rule.rules_dict
-        else:
-            mode = DocumentService.DEFAULT_RULES['mode']
-            rules = DocumentService.DEFAULT_RULES['rules']
+            if dataset_process_rule:
+                mode = dataset_process_rule.mode
+                rules = dataset_process_rule.rules_dict
 
         return {
             'mode': mode,
@@ -274,7 +277,8 @@ class DatasetDocumentListApi(Resource):
         parser.add_argument('duplicate', type=bool, nullable=False, location='json')
         parser.add_argument('original_document_id', type=str, required=False, location='json')
         parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
-        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False, location='json')
+        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
+                            location='json')
         args = parser.parse_args()
 
         if not dataset.indexing_technique and not args['indexing_technique']:
@@ -283,20 +287,6 @@ class DatasetDocumentListApi(Resource):
         # validate args
         DocumentService.document_create_args_validate(args)
 
-        # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-
         try:
             documents, batch = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
         except ProviderTokenNotInitError as ex:
@@ -334,17 +324,20 @@ class DatasetInitApi(Resource):
         parser.add_argument('data_source', type=dict, required=True, nullable=True, location='json')
         parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
         parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
-        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False, location='json')
+        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
+                            location='json')
         args = parser.parse_args()
-
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
+        if args['indexing_technique'] == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
 
         # validate args
         DocumentService.document_create_args_validate(args)
@@ -413,7 +406,8 @@ class DocumentIndexingEstimateApi(DocumentResource):
 
                 try:
                     response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, [file],
-                                                                      data_process_rule_dict, None, dataset_id)
+                                                                      data_process_rule_dict, None,
+                                                                      'English', dataset_id)
                 except LLMBadRequestError:
                     raise ProviderNotInitializeError(
                         f"No Embedding Model available. Please configure a valid provider "
@@ -482,7 +476,8 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
             indexing_runner = IndexingRunner()
             try:
                 response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
-                                                                  data_process_rule_dict,  None, dataset_id)
+                                                                  data_process_rule_dict, None,
+                                                                  'English', dataset_id)
             except LLMBadRequestError:
                 raise ProviderNotInitializeError(
                     f"No Embedding Model available. Please configure a valid provider "
@@ -496,7 +491,7 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
                 response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
                                                                     info_list,
                                                                     data_process_rule_dict,
-                                                                    None, dataset_id)
+                                                                    None, 'English', dataset_id)
             except LLMBadRequestError:
                 raise ProviderNotInitializeError(
                     f"No Embedding Model available. Please configure a valid provider "
@@ -724,6 +719,12 @@ class DocumentDeleteApi(DocumentResource):
     def delete(self, dataset_id, document_id):
         dataset_id = str(dataset_id)
         document_id = str(document_id)
+        dataset = DatasetService.get_dataset(dataset_id)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
+
         document = self.get_document(dataset_id, document_id)
 
         try:
@@ -764,11 +765,13 @@ class DocumentMetadataApi(DocumentResource):
         metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
 
         document.doc_metadata = {}
-
-        for key, value_type in metadata_schema.items():
-            value = doc_metadata.get(key)
-            if value is not None and isinstance(value, value_type):
-                document.doc_metadata[key] = value
+        if doc_type == 'others':
+            document.doc_metadata = doc_metadata
+        else:
+            for key, value_type in metadata_schema.items():
+                value = doc_metadata.get(key)
+                if value is not None and isinstance(value, value_type):
+                    document.doc_metadata[key] = value
 
         document.doc_type = doc_type
         document.updated_at = datetime.utcnow()
@@ -784,6 +787,12 @@ class DocumentStatusApi(DocumentResource):
     def patch(self, dataset_id, document_id, action):
         dataset_id = str(dataset_id)
         document_id = str(document_id)
+        dataset = DatasetService.get_dataset(dataset_id)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
+
         document = self.get_document(dataset_id, document_id)
 
         # The role of the current user in the ta table must be admin or owner
@@ -852,6 +861,14 @@ class DocumentStatusApi(DocumentResource):
             if not document.archived:
                 raise InvalidActionError('Document is not archived.')
 
+            # check document limit
+            if current_app.config['EDITION'] == 'CLOUD':
+                documents_count = DocumentService.get_tenant_documents_count()
+                total_count = documents_count + 1
+                tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
+                if total_count > tenant_document_count:
+                    raise ValueError(f"All your documents have overed limit {tenant_document_count}.")
+
             document.archived = False
             document.archived_at = None
             document.archived_by = None
@@ -869,6 +886,10 @@ class DocumentStatusApi(DocumentResource):
 
 
 class DocumentPauseApi(DocumentResource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
     def patch(self, dataset_id, document_id):
         """pause document."""
         dataset_id = str(dataset_id)
@@ -898,6 +919,9 @@ class DocumentPauseApi(DocumentResource):
 
 
 class DocumentRecoverApi(DocumentResource):
+    @setup_required
+    @login_required
+    @account_initialization_required
     def patch(self, dataset_id, document_id):
         """recover document."""
         dataset_id = str(dataset_id)
@@ -923,6 +947,21 @@ class DocumentRecoverApi(DocumentResource):
         return {'result': 'success'}, 204
 
 
+class DocumentLimitApi(DocumentResource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        """get document limit"""
+        documents_count = DocumentService.get_tenant_documents_count()
+        tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
+
+        return {
+            'documents_count': documents_count,
+            'documents_limit': tenant_document_count
+                }, 200
+
+
 api.add_resource(GetProcessRuleApi, '/datasets/process-rule')
 api.add_resource(DatasetDocumentListApi,
                  '/datasets/<uuid:dataset_id>/documents')
@@ -948,3 +987,4 @@ api.add_resource(DocumentStatusApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/status/<string:action>')
 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(DocumentLimitApi, '/datasets/limit')

api/controllers/console/datasets/datasets_segments.py

@@ -1,9 +1,8 @@
 # -*- coding:utf-8 -*-
 import uuid
 from datetime import datetime
-
 from flask import request
-from flask_login import login_required, current_user
+from flask_login import current_user
 from flask_restful import Resource, reqparse, fields, marshal
 from werkzeug.exceptions import NotFound, Forbidden
 
@@ -15,6 +14,7 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_providers.model_factory import ModelFactory
+from core.login.login import login_required
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
@@ -149,7 +149,8 @@ class DatasetDocumentSegmentApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
-
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # The role of the current user in the ta table must be admin or owner
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
@@ -158,20 +159,20 @@ class DatasetDocumentSegmentApi(Resource):
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
-
-        # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
+        if dataset.indexing_technique == 'high_quality':
+            # check embedding model setting
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
 
         segment = DocumentSegment.query.filter(
             DocumentSegment.id == str(segment_id),
@@ -244,18 +245,19 @@ class DatasetDocumentSegmentAddApi(Resource):
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
         # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
+        if dataset.indexing_technique == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
         try:
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
@@ -284,25 +286,28 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # check document
         document_id = str(document_id)
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound('Document not found.')
-        # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        # check segment
+        if dataset.indexing_technique == 'high_quality':
+            # check embedding model setting
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+            # check segment
         segment_id = str(segment_id)
         segment = DocumentSegment.query.filter(
             DocumentSegment.id == str(segment_id),
@@ -339,6 +344,8 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # check document
         document_id = str(document_id)
         document = DocumentService.get_document(dataset_id, document_id)
@@ -378,18 +385,6 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound('Document not found.')
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
         # get file from request
         file = request.files['file']
         # check file

api/controllers/console/datasets/file.py

@@ -8,7 +8,8 @@ from pathlib import Path
 
 from cachetools import TTLCache
 from flask import request, current_app
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, marshal_with, fields
 from werkzeug.exceptions import NotFound
 

api/controllers/console/datasets/hit_testing.py

@@ -1,6 +1,7 @@
 import logging
 
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, marshal, fields
 from werkzeug.exceptions import InternalServerError, NotFound, Forbidden
 

api/controllers/console/explore/installed_app.py

@@ -1,7 +1,8 @@
 # -*- coding:utf-8 -*-
 from datetime import datetime
 
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with, inputs
 from sqlalchemy import and_
 from werkzeug.exceptions import NotFound, Forbidden, BadRequest

api/controllers/console/explore/recommended_app.py

@@ -1,5 +1,6 @@
 # -*- coding:utf-8 -*-
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, fields, marshal_with
 from sqlalchemy import and_
 

api/controllers/console/explore/wraps.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource
 from functools import wraps
 

api/controllers/console/universal_chat/wraps.py

@@ -1,7 +1,8 @@
 import json
 from functools import wraps
 
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required

api/controllers/console/workspace/account.py

@@ -3,7 +3,8 @@ from datetime import datetime
 
 import pytz
 from flask import current_app, request
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with
 
 from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError

api/controllers/console/workspace/members.py

@@ -1,6 +1,7 @@
 # -*- coding:utf-8 -*-
 from flask import current_app
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse, marshal_with, abort, fields, marshal
 
 import services

api/controllers/console/workspace/model_providers.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/models.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/workspace/providers.py

@@ -1,5 +1,6 @@
 # -*- coding:utf-8 -*-
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/tool_providers.py

@@ -1,6 +1,7 @@
 import json
 
-from flask_login import login_required, current_user
+from flask_login import current_user
+from core.login.login import login_required
 from flask_restful import Resource, abort, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/workspace.py

@@ -2,10 +2,13 @@
 import logging
 
 from flask import request
-from flask_login import login_required, current_user
-from flask_restful import Resource, fields, marshal_with, reqparse, marshal
+from flask_login import current_user
+from core.login.login import login_required
+from flask_restful import Resource, fields, marshal_with, reqparse, marshal, inputs
+from flask_restful.inputs import int_range
 
 from controllers.console import api
+from controllers.console.admin import admin_required
 from controllers.console.setup import setup_required
 from controllers.console.error import AccountNotLinkTenantError
 from controllers.console.wraps import account_initialization_required
@@ -43,6 +46,13 @@ tenants_fields = {
     'current': fields.Boolean
 }
 
+workspace_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'status': fields.String,
+    'created_at': TimestampField
+}
+
 
 class TenantListApi(Resource):
     @setup_required
@@ -57,6 +67,38 @@ class TenantListApi(Resource):
         return {'workspaces': marshal(tenants, tenants_fields)}, 200
 
 
+class WorkspaceListApi(Resource):
+    @setup_required
+    @admin_required
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('page', type=inputs.int_range(1, 99999), required=False, default=1, location='args')
+        parser.add_argument('limit', type=inputs.int_range(1, 100), required=False, default=20, location='args')
+        args = parser.parse_args()
+
+        tenants = db.session.query(Tenant).order_by(Tenant.created_at.desc())\
+            .paginate(page=args['page'], per_page=args['limit'])
+
+        has_more = False
+        if len(tenants.items) == args['limit']:
+            current_page_first_tenant = tenants[-1]
+            rest_count = db.session.query(Tenant).filter(
+                Tenant.created_at < current_page_first_tenant.created_at,
+                Tenant.id != current_page_first_tenant.id
+            ).count()
+
+            if rest_count > 0:
+                has_more = True
+        total = db.session.query(Tenant).count()
+        return {
+            'data': marshal(tenants.items, workspace_fields),
+            'has_more': has_more,
+            'limit': args['limit'],
+            'page': args['page'],
+            'total': total
+                }, 200
+
+
 class TenantApi(Resource):
     @setup_required
     @login_required
@@ -92,6 +134,7 @@ class SwitchWorkspaceApi(Resource):
 
 
 api.add_resource(TenantListApi, '/workspaces')  # GET for getting all tenants
+api.add_resource(WorkspaceListApi, '/all-workspaces')  # GET for getting all tenants
 api.add_resource(TenantApi, '/workspaces/current', endpoint='workspaces_current')  # GET for getting current tenant info
 api.add_resource(TenantApi, '/info', endpoint='info')  # Deprecated
 api.add_resource(SwitchWorkspaceApi, '/workspaces/switch')  # POST for switching tenant

api/core/agent/agent/multi_dataset_router_agent.py

@@ -52,7 +52,7 @@ class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
         elif len(self.tools) == 1:
             tool = next(iter(self.tools))
             tool = cast(DatasetRetrieverTool, tool)
-            rst = tool.run(tool_input={'dataset_id': tool.dataset_id, 'query': kwargs['input']})
+            rst = tool.run(tool_input={'query': kwargs['input']})
             return AgentFinish(return_values={"output": rst}, log=rst)
 
         if intermediate_steps:
@@ -60,7 +60,13 @@ class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
             return AgentFinish(return_values={"output": observation}, log=observation)
 
         try:
-            return super().plan(intermediate_steps, callbacks, **kwargs)
+            agent_decision = super().plan(intermediate_steps, callbacks, **kwargs)
+            if isinstance(agent_decision, AgentAction):
+                tool_inputs = agent_decision.tool_input
+                if isinstance(tool_inputs, dict) and 'query' in tool_inputs:
+                    tool_inputs['query'] = kwargs['input']
+                    agent_decision.tool_input = tool_inputs
+            return agent_decision
         except Exception as e:
             new_exception = self.model_instance.handle_exceptions(e)
             raise new_exception

api/core/agent/agent/openai_function_call.py

@@ -45,7 +45,7 @@ class AutoSummarizingOpenAIFunctionCallAgent(OpenAIFunctionsAgent, OpenAIFunctio
         :return:
         """
         original_max_tokens = self.llm.max_tokens
-        self.llm.max_tokens = 15
+        self.llm.max_tokens = 40
 
         prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
         messages = prompt.to_messages()
@@ -97,6 +97,13 @@ class AutoSummarizingOpenAIFunctionCallAgent(OpenAIFunctionsAgent, OpenAIFunctio
             messages, functions=self.functions, callbacks=callbacks
         )
         agent_decision = _parse_ai_message(predicted_message)
+
+        if isinstance(agent_decision, AgentAction) and agent_decision.tool == 'dataset':
+            tool_inputs = agent_decision.tool_input
+            if isinstance(tool_inputs, dict) and 'query' in tool_inputs:
+                tool_inputs['query'] = kwargs['input']
+                agent_decision.tool_input = tool_inputs
+
         return agent_decision
 
     @classmethod

api/core/agent/agent/structed_multi_dataset_router_agent.py

@@ -90,7 +90,7 @@ class StructuredMultiDatasetRouterAgent(StructuredChatAgent):
         elif len(self.dataset_tools) == 1:
             tool = next(iter(self.dataset_tools))
             tool = cast(DatasetRetrieverTool, tool)
-            rst = tool.run(tool_input={'dataset_id': tool.dataset_id, 'query': kwargs['input']})
+            rst = tool.run(tool_input={'query': kwargs['input']})
             return AgentFinish(return_values={"output": rst}, log=rst)
 
         full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
@@ -102,7 +102,13 @@ class StructuredMultiDatasetRouterAgent(StructuredChatAgent):
             raise new_exception
 
         try:
-            return self.output_parser.parse(full_output)
+            agent_decision = self.output_parser.parse(full_output)
+            if isinstance(agent_decision, AgentAction):
+                tool_inputs = agent_decision.tool_input
+                if isinstance(tool_inputs, dict) and 'query' in tool_inputs:
+                    tool_inputs['query'] = kwargs['input']
+                    agent_decision.tool_input = tool_inputs
+            return agent_decision
         except OutputParserException:
             return AgentFinish({"output": "I'm sorry, the answer of model is invalid, "
                                           "I don't know how to respond to that."}, "")

api/core/agent/agent/structured_chat.py

@@ -106,7 +106,13 @@ class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
             raise new_exception
 
         try:
-            return self.output_parser.parse(full_output)
+            agent_decision = self.output_parser.parse(full_output)
+            if isinstance(agent_decision, AgentAction) and agent_decision.tool == 'dataset':
+                tool_inputs = agent_decision.tool_input
+                if isinstance(tool_inputs, dict) and 'query' in tool_inputs:
+                    tool_inputs['query'] = kwargs['input']
+                    agent_decision.tool_input = tool_inputs
+            return agent_decision
         except OutputParserException:
             return AgentFinish({"output": "I'm sorry, the answer of model is invalid, "
                                           "I don't know how to respond to that."}, "")

api/core/callback_handler/dataset_tool_callback_handler.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from json import JSONDecodeError
 
 from typing import Any, Dict, List, Union, Optional
 
@@ -44,10 +45,15 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
         input_str: str,
         **kwargs: Any,
     ) -> None:
-        # tool_name = serialized.get('name')
-        input_dict = json.loads(input_str.replace("'", "\""))
-        dataset_id = input_dict.get('dataset_id')
-        query = input_dict.get('query')
+        tool_name: str = serialized.get('name')
+        dataset_id = tool_name.removeprefix('dataset-')
+
+        try:
+            input_dict = json.loads(input_str.replace("'", "\""))
+            query = input_dict.get('query')
+        except JSONDecodeError:
+            query = input_str
+
         self.conversation_message_task.on_dataset_query_end(DatasetQueryObj(dataset_id=dataset_id, query=query))
 
     def on_tool_end(

api/core/completion.py

@@ -130,13 +130,12 @@ class Completion:
             fake_response = agent_execute_result.output
 
         # get llm prompt
-        prompt_messages, stop_words = cls.get_main_llm_prompt(
+        prompt_messages, stop_words = model_instance.get_prompt(
             mode=mode,
-            model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
-            query=query,
             inputs=inputs,
-            agent_execute_result=agent_execute_result,
+            query=query,
+            context=agent_execute_result.output if agent_execute_result else None,
             memory=memory
         )
 
@@ -154,113 +153,6 @@ class Completion:
 
         return response
 
-    @classmethod
-    def get_main_llm_prompt(cls, mode: str, model: dict,
-                            pre_prompt: str, query: str, inputs: dict,
-                            agent_execute_result: Optional[AgentExecuteResult],
-                            memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]) -> \
-            Tuple[List[PromptMessage], Optional[List[str]]]:
-        if mode == 'completion':
-            prompt_template = JinjaPromptTemplate.from_template(
-                template=("""Use the following context as your learned knowledge, inside <context></context> XML tags.
-
-<context>
-{{context}}
-</context>
-
-When answer to user:
-- If you don't know, just say that you don't know.
-- If you don't know when you are not sure, ask for clarification. 
-Avoid mentioning that you obtained the information from the context.
-And answer according to the language of the user's question.
-""" if agent_execute_result else "")
-                         + (pre_prompt + "\n" if pre_prompt else "")
-                         + "{{query}}\n"
-            )
-
-            if agent_execute_result:
-                inputs['context'] = agent_execute_result.output
-
-            prompt_inputs = {k: inputs[k] for k in prompt_template.input_variables if k in inputs}
-            prompt_content = prompt_template.format(
-                query=query,
-                **prompt_inputs
-            )
-
-            return [PromptMessage(content=prompt_content)], None
-        else:
-            messages: List[BaseMessage] = []
-
-            human_inputs = {
-                "query": query
-            }
-
-            human_message_prompt = ""
-
-            if pre_prompt:
-                pre_prompt_inputs = {k: inputs[k] for k in
-                                     JinjaPromptTemplate.from_template(template=pre_prompt).input_variables
-                                     if k in inputs}
-
-                if pre_prompt_inputs:
-                    human_inputs.update(pre_prompt_inputs)
-
-            if agent_execute_result:
-                human_inputs['context'] = agent_execute_result.output
-                human_message_prompt += """Use the following context as your learned knowledge, inside <context></context> XML tags.
-
-<context>
-{{context}}
-</context>
-
-When answer to user:
-- If you don't know, just say that you don't know.
-- If you don't know when you are not sure, ask for clarification. 
-Avoid mentioning that you obtained the information from the context.
-And answer according to the language of the user's question.
-"""
-
-            if pre_prompt:
-                human_message_prompt += pre_prompt
-
-            query_prompt = "\n\nHuman: {{query}}\n\nAssistant: "
-
-            if memory:
-                # append chat histories
-                tmp_human_message = PromptBuilder.to_human_message(
-                    prompt_content=human_message_prompt + query_prompt,
-                    inputs=human_inputs
-                )
-
-                if memory.model_instance.model_rules.max_tokens.max:
-                    curr_message_tokens = memory.model_instance.get_num_tokens(to_prompt_messages([tmp_human_message]))
-                    max_tokens = model.get("completion_params").get('max_tokens')
-                    rest_tokens = memory.model_instance.model_rules.max_tokens.max - max_tokens - curr_message_tokens
-                    rest_tokens = max(rest_tokens, 0)
-                else:
-                    rest_tokens = 2000
-
-                histories = cls.get_history_messages_from_memory(memory, rest_tokens)
-                human_message_prompt += "\n\n" if human_message_prompt else ""
-                human_message_prompt += "Here is the chat histories between human and assistant, " \
-                                        "inside <histories></histories> XML tags.\n\n<histories>\n"
-                human_message_prompt += histories + "\n</histories>"
-
-            human_message_prompt += query_prompt
-
-            # construct main prompt
-            human_message = PromptBuilder.to_human_message(
-                prompt_content=human_message_prompt,
-                inputs=human_inputs
-            )
-
-            messages.append(human_message)
-
-            for message in messages:
-                message.content = re.sub(r'<\|.*?\|>', '', message.content)
-
-            return to_prompt_messages(messages), ['\nHuman:', '</histories>']
-
     @classmethod
     def get_history_messages_from_memory(cls, memory: ReadOnlyConversationTokenDBBufferSharedMemory,
                                          max_token_limit: int) -> str:
@@ -307,13 +199,12 @@ And answer according to the language of the user's question.
             max_tokens = 0
 
         # get prompt without memory and context
-        prompt_messages, _ = cls.get_main_llm_prompt(
+        prompt_messages, _ = model_instance.get_prompt(
             mode=mode,
-            model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
-            query=query,
             inputs=inputs,
-            agent_execute_result=None,
+            query=query,
+            context=None,
             memory=None
         )
 
@@ -358,13 +249,12 @@ And answer according to the language of the user's question.
         )
 
         # get llm prompt
-        old_prompt_messages, _ = cls.get_main_llm_prompt(
-            mode="completion",
-            model=app_model_config.model_dict,
+        old_prompt_messages, _ = final_model_instance.get_prompt(
+            mode='completion',
             pre_prompt=pre_prompt,
-            query=message.query,
             inputs=message.inputs,
-            agent_execute_result=None,
+            query=message.query,
+            context=None,
             memory=None
         )
 

api/core/conversation_message_task.py

@@ -137,7 +137,8 @@ class ConversationMessageTask:
         db.session.flush()
 
     def append_message_text(self, text: str):
-        self._pub_handler.pub_text(text)
+        if text is not None:
+            self._pub_handler.pub_text(text)
 
     def save_message(self, llm_message: LLMMessage, by_stopped: bool = False):
         message_tokens = llm_message.prompt_tokens

api/core/data_loader/loader/excel.py

@@ -30,6 +30,8 @@ class ExcelLoader(BaseLoader):
         wb = load_workbook(filename=self._file_path, read_only=True)
         # loop over all sheets
         for sheet in wb:
+            if 'A1:A1' == sheet.calculate_dimension():
+                sheet.reset_dimensions()
             for row in sheet.iter_rows(values_only=True):
                 if all(v is None for v in row):
                     continue
@@ -38,7 +40,7 @@ class ExcelLoader(BaseLoader):
                 else:
                     row_dict = dict(zip(keys, list(map(str, row))))
                     row_dict = {k: v for k, v in row_dict.items() if v}
-                    item = ''.join(f'{k}:{v}\n' for k, v in row_dict.items())
+                    item = ''.join(f'{k}:{v};' for k, v in row_dict.items())
                     document = Document(page_content=item, metadata={'source': self._file_path})
                     data.append(document)
 

api/core/docstore/dataset_docstore.py

@@ -67,12 +67,13 @@ class DatesetDocumentStore:
 
         if max_position is None:
             max_position = 0
-
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=self._dataset.tenant_id,
-            model_provider_name=self._dataset.embedding_model_provider,
-            model_name=self._dataset.embedding_model
-        )
+        embedding_model = None
+        if self._dataset.indexing_technique == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=self._dataset.tenant_id,
+                model_provider_name=self._dataset.embedding_model_provider,
+                model_name=self._dataset.embedding_model
+            )
 
         for doc in docs:
             if not isinstance(doc, Document):
@@ -88,7 +89,7 @@ class DatesetDocumentStore:
                 )
 
             # calc embedding use tokens
-            tokens = embedding_model.get_num_tokens(doc.page_content)
+            tokens = embedding_model.get_num_tokens(doc.page_content) if embedding_model else 0
 
             if not segment_document:
                 max_position += 1
@@ -103,6 +104,7 @@ class DatesetDocumentStore:
                     content=doc.page_content,
                     word_count=len(doc.page_content),
                     tokens=tokens,
+                    enabled=False,
                     created_by=self._user_id,
                 )
                 if 'answer' in doc.metadata and doc.metadata['answer']:

api/core/embedding/cached_embedding.py

@@ -1,6 +1,7 @@
 import logging
 from typing import List
 
+import numpy as np
 from langchain.embeddings.base import Embeddings
 from sqlalchemy.exc import IntegrityError
 
@@ -32,14 +33,17 @@ class CacheEmbedding(Embeddings):
                 embedding_results = self._embeddings.client.embed_documents(embedding_queue_texts)
             except Exception as ex:
                 raise self._embeddings.handle_exceptions(ex)
-
             i = 0
+            normalized_embedding_results = []
             for text in embedding_queue_texts:
                 hash = helper.generate_text_hash(text)
 
                 try:
                     embedding = Embedding(model_name=self._embeddings.name, hash=hash)
-                    embedding.set_embedding(embedding_results[i])
+                    vector = embedding_results[i]
+                    normalized_embedding = (vector / np.linalg.norm(vector)).tolist()
+                    normalized_embedding_results.append(normalized_embedding)
+                    embedding.set_embedding(normalized_embedding)
                     db.session.add(embedding)
                     db.session.commit()
                 except IntegrityError:
@@ -51,7 +55,7 @@ class CacheEmbedding(Embeddings):
                 finally:
                     i += 1
 
-            text_embeddings.extend(embedding_results)
+            text_embeddings.extend(normalized_embedding_results)
         return text_embeddings
 
     def embed_query(self, text: str) -> List[float]:
@@ -64,6 +68,7 @@ class CacheEmbedding(Embeddings):
 
         try:
             embedding_results = self._embeddings.client.embed_query(text)
+            embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist()
         except Exception as ex:
             raise self._embeddings.handle_exceptions(ex)
 
@@ -79,4 +84,3 @@ class CacheEmbedding(Embeddings):
 
         return embedding_results
 
-

api/core/index/index.py

@@ -1,10 +1,18 @@
+import json
+
 from flask import current_app
+from langchain.embeddings import OpenAIEmbeddings
 
 from core.embedding.cached_embedding import CacheEmbedding
 from core.index.keyword_table_index.keyword_table_index import KeywordTableIndex, KeywordTableConfig
 from core.index.vector_index.vector_index import VectorIndex
 from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.embedding.openai_embedding import OpenAIEmbedding
+from core.model_providers.models.entity.model_params import ModelKwargs
+from core.model_providers.models.llm.openai_model import OpenAIModel
+from core.model_providers.providers.openai_provider import OpenAIProvider
 from models.dataset import Dataset
+from models.provider import Provider, ProviderType
 
 
 class IndexBuilder:
@@ -35,4 +43,13 @@ class IndexBuilder:
                 )
             )
         else:
-            raise ValueError('Unknown indexing technique')
+            raise ValueError('Unknown indexing technique')
+
+    @classmethod
+    def get_default_high_quality_index(cls, dataset: Dataset):
+        embeddings = OpenAIEmbeddings(openai_api_key=' ')
+        return VectorIndex(
+            dataset=dataset,
+            config=current_app.config,
+            embeddings=embeddings
+        )

api/core/index/vector_index/base.py

@@ -15,12 +15,12 @@ from models.dataset import Document as DatasetDocument
 
 
 class BaseVectorIndex(BaseIndex):
-    
+
     def __init__(self, dataset: Dataset, embeddings: Embeddings):
         super().__init__(dataset)
         self._embeddings = embeddings
         self._vector_store = None
-        
+
     def get_type(self) -> str:
         raise NotImplementedError
 
@@ -143,7 +143,7 @@ class BaseVectorIndex(BaseIndex):
                 DocumentSegment.status == 'completed',
                 DocumentSegment.enabled == True
             ).all()
-            
+
             for segment in segments:
                 document = Document(
                     page_content=segment.content,
@@ -173,3 +173,73 @@ class BaseVectorIndex(BaseIndex):
 
         self.dataset = dataset
         logging.info(f"Dataset {dataset.id} recreate successfully.")
+
+    def create_qdrant_dataset(self, dataset: Dataset):
+        logging.info(f"create_qdrant_dataset {dataset.id}")
+
+        try:
+            self.delete()
+        except UnexpectedStatusCodeException as e:
+            if e.status_code != 400:
+                # 400 means index not exists
+                raise e
+
+        dataset_documents = db.session.query(DatasetDocument).filter(
+            DatasetDocument.dataset_id == dataset.id,
+            DatasetDocument.indexing_status == 'completed',
+            DatasetDocument.enabled == True,
+            DatasetDocument.archived == False,
+        ).all()
+
+        documents = []
+        for dataset_document in dataset_documents:
+            segments = db.session.query(DocumentSegment).filter(
+                DocumentSegment.document_id == dataset_document.id,
+                DocumentSegment.status == 'completed',
+                DocumentSegment.enabled == True
+            ).all()
+
+            for segment in segments:
+                document = Document(
+                    page_content=segment.content,
+                    metadata={
+                        "doc_id": segment.index_node_id,
+                        "doc_hash": segment.index_node_hash,
+                        "document_id": segment.document_id,
+                        "dataset_id": segment.dataset_id,
+                    }
+                )
+
+                documents.append(document)
+
+        if documents:
+            try:
+                self.create(documents)
+            except Exception as e:
+                raise e
+
+        logging.info(f"Dataset {dataset.id} recreate successfully.")
+
+    def update_qdrant_dataset(self, dataset: Dataset):
+        logging.info(f"update_qdrant_dataset {dataset.id}")
+
+        segment = db.session.query(DocumentSegment).filter(
+            DocumentSegment.dataset_id == dataset.id,
+            DocumentSegment.status == 'completed',
+            DocumentSegment.enabled == True
+        ).first()
+
+        if segment:
+            try:
+                exist = self.text_exists(segment.index_node_id)
+                if exist:
+                    index_struct = {
+                        "type": 'qdrant',
+                        "vector_store": {"class_prefix": dataset.index_struct_dict['vector_store']['class_prefix']}
+                    }
+                    dataset.index_struct = json.dumps(index_struct)
+                    db.session.commit()
+            except Exception as e:
+                raise e
+
+        logging.info(f"Dataset {dataset.id} recreate successfully.")

api/core/index/vector_index/milvus_vector_index.py

@@ -0,0 +1,114 @@
+from typing import Optional, cast
+
+from langchain.embeddings.base import Embeddings
+from langchain.schema import Document, BaseRetriever
+from langchain.vectorstores import VectorStore, milvus
+from pydantic import BaseModel, root_validator
+
+from core.index.base import BaseIndex
+from core.index.vector_index.base import BaseVectorIndex
+from core.vector_store.milvus_vector_store import MilvusVectorStore
+from core.vector_store.weaviate_vector_store import WeaviateVectorStore
+from models.dataset import Dataset
+
+
+class MilvusConfig(BaseModel):
+    endpoint: str
+    user: str
+    password: str
+    batch_size: int = 100
+
+    @root_validator()
+    def validate_config(cls, values: dict) -> dict:
+        if not values['endpoint']:
+            raise ValueError("config MILVUS_ENDPOINT is required")
+        if not values['user']:
+            raise ValueError("config MILVUS_USER is required")
+        if not values['password']:
+            raise ValueError("config MILVUS_PASSWORD is required")
+        return values
+
+
+class MilvusVectorIndex(BaseVectorIndex):
+    def __init__(self, dataset: Dataset, config: MilvusConfig, embeddings: Embeddings):
+        super().__init__(dataset, embeddings)
+        self._client = self._init_client(config)
+
+    def get_type(self) -> str:
+        return 'milvus'
+
+    def get_index_name(self, dataset: Dataset) -> str:
+        if self.dataset.index_struct_dict:
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
+                # original class_prefix
+                class_prefix += '_Node'
+
+            return class_prefix
+
+        dataset_id = dataset.id
+        return "Vector_index_" + dataset_id.replace("-", "_") + '_Node'
+
+
+    def to_index_struct(self) -> dict:
+        return {
+            "type": self.get_type(),
+            "vector_store": {"class_prefix": self.get_index_name(self.dataset)}
+        }
+
+    def create(self, texts: list[Document], **kwargs) -> BaseIndex:
+        uuids = self._get_uuids(texts)
+        self._vector_store = WeaviateVectorStore.from_documents(
+            texts,
+            self._embeddings,
+            client=self._client,
+            index_name=self.get_index_name(self.dataset),
+            uuids=uuids,
+            by_text=False
+        )
+
+        return self
+
+    def _get_vector_store(self) -> VectorStore:
+        """Only for created index."""
+        if self._vector_store:
+            return self._vector_store
+
+        attributes = ['doc_id', 'dataset_id', 'document_id']
+        if self._is_origin():
+            attributes = ['doc_id']
+
+        return WeaviateVectorStore(
+            client=self._client,
+            index_name=self.get_index_name(self.dataset),
+            text_key='text',
+            embedding=self._embeddings,
+            attributes=attributes,
+            by_text=False
+        )
+
+    def _get_vector_store_class(self) -> type:
+        return MilvusVectorStore
+
+    def delete_by_document_id(self, document_id: str):
+        if self._is_origin():
+            self.recreate_dataset(self.dataset)
+            return
+
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        vector_store.del_texts({
+            "operator": "Equal",
+            "path": ["document_id"],
+            "valueText": document_id
+        })
+
+    def _is_origin(self):
+        if self.dataset.index_struct_dict:
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
+                # original class_prefix
+                return True
+
+        return False

api/core/index/vector_index/qdrant_vector_index.py

@@ -44,15 +44,20 @@ class QdrantVectorIndex(BaseVectorIndex):
 
     def get_index_name(self, dataset: Dataset) -> str:
         if self.dataset.index_struct_dict:
-            return self.dataset.index_struct_dict['vector_store']['collection_name']
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
+                # original class_prefix
+                class_prefix += '_Node'
+
+            return class_prefix
 
         dataset_id = dataset.id
-        return "Index_" + dataset_id.replace("-", "_")
+        return "Vector_index_" + dataset_id.replace("-", "_") + '_Node'
 
     def to_index_struct(self) -> dict:
         return {
             "type": self.get_type(),
-            "vector_store": {"collection_name": self.get_index_name(self.dataset)}
+            "vector_store": {"class_prefix": self.get_index_name(self.dataset)}
         }
 
     def create(self, texts: list[Document], **kwargs) -> BaseIndex:
@@ -62,7 +67,7 @@ class QdrantVectorIndex(BaseVectorIndex):
             self._embeddings,
             collection_name=self.get_index_name(self.dataset),
             ids=uuids,
-            content_payload_key='text',
+            content_payload_key='page_content',
             **self._client_config.to_qdrant_params()
         )
 
@@ -72,7 +77,9 @@ class QdrantVectorIndex(BaseVectorIndex):
         """Only for created index."""
         if self._vector_store:
             return self._vector_store
-        
+        attributes = ['doc_id', 'dataset_id', 'document_id']
+        if self._is_origin():
+            attributes = ['doc_id']
         client = qdrant_client.QdrantClient(
             **self._client_config.to_qdrant_params()
         )
@@ -81,7 +88,7 @@ class QdrantVectorIndex(BaseVectorIndex):
             client=client,
             collection_name=self.get_index_name(self.dataset),
             embeddings=self._embeddings,
-            content_payload_key='text'
+            content_payload_key='page_content'
         )
 
     def _get_vector_store_class(self) -> type:
@@ -108,8 +115,8 @@ class QdrantVectorIndex(BaseVectorIndex):
 
     def _is_origin(self):
         if self.dataset.index_struct_dict:
-            class_prefix: str = self.dataset.index_struct_dict['vector_store']['collection_name']
-            if class_prefix.startswith('Vector_'):
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
                 # original class_prefix
                 return True
 

api/core/indexing_runner.py

@@ -217,25 +217,29 @@ class IndexingRunner:
             db.session.commit()
 
     def file_indexing_estimate(self, tenant_id: str, file_details: List[UploadFile], tmp_processing_rule: dict,
-                               doc_form: str = None, doc_language: str = 'English', dataset_id: str = None) -> dict:
+                               doc_form: str = None, doc_language: str = 'English', dataset_id: str = None,
+                               indexing_technique: str = 'economy') -> dict:
         """
         Estimate the indexing for the document.
         """
+        embedding_model = None
         if dataset_id:
             dataset = Dataset.query.filter_by(
                 id=dataset_id
             ).first()
             if not dataset:
                 raise ValueError('Dataset not found.')
-            embedding_model = ModelFactory.get_embedding_model(
-                tenant_id=dataset.tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
+            if dataset.indexing_technique == 'high_quality' or indexing_technique == 'high_quality':
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=dataset.tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
         else:
-            embedding_model = ModelFactory.get_embedding_model(
-                tenant_id=tenant_id
-            )
+            if indexing_technique == 'high_quality':
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=tenant_id
+                )
         tokens = 0
         preview_texts = []
         total_segments = 0
@@ -263,8 +267,8 @@ class IndexingRunner:
             for document in documents:
                 if len(preview_texts) < 5:
                     preview_texts.append(document.page_content)
-
-                tokens += embedding_model.get_num_tokens(self.filter_string(document.page_content))
+                if indexing_technique == 'high_quality' or embedding_model:
+                    tokens += embedding_model.get_num_tokens(self.filter_string(document.page_content))
 
         if doc_form and doc_form == 'qa_model':
             text_generation_model = ModelFactory.get_text_generation_model(
@@ -286,32 +290,35 @@ class IndexingRunner:
         return {
             "total_segments": total_segments,
             "tokens": tokens,
-            "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)),
-            "currency": embedding_model.get_currency(),
+            "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)) if embedding_model else 0,
+            "currency": embedding_model.get_currency() if embedding_model else 'USD',
             "preview": preview_texts
         }
 
     def notion_indexing_estimate(self, tenant_id: str, notion_info_list: list, tmp_processing_rule: dict,
-                                 doc_form: str = None, doc_language: str = 'English', dataset_id: str = None) -> dict:
+                                 doc_form: str = None, doc_language: str = 'English', dataset_id: str = None,
+                                 indexing_technique: str = 'economy') -> dict:
         """
         Estimate the indexing for the document.
         """
+        embedding_model = None
         if dataset_id:
             dataset = Dataset.query.filter_by(
                 id=dataset_id
             ).first()
             if not dataset:
                 raise ValueError('Dataset not found.')
-            embedding_model = ModelFactory.get_embedding_model(
-                tenant_id=dataset.tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
+            if dataset.indexing_technique == 'high_quality' or indexing_technique == 'high_quality':
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=dataset.tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
         else:
-            embedding_model = ModelFactory.get_embedding_model(
-                tenant_id=tenant_id
-            )
-
+            if indexing_technique == 'high_quality':
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=tenant_id
+                )
         # load data from notion
         tokens = 0
         preview_texts = []
@@ -356,8 +363,8 @@ class IndexingRunner:
                 for document in documents:
                     if len(preview_texts) < 5:
                         preview_texts.append(document.page_content)
-
-                    tokens += embedding_model.get_num_tokens(document.page_content)
+                    if indexing_technique == 'high_quality' or embedding_model:
+                        tokens += embedding_model.get_num_tokens(document.page_content)
 
         if doc_form and doc_form == 'qa_model':
             text_generation_model = ModelFactory.get_text_generation_model(
@@ -379,8 +386,8 @@ class IndexingRunner:
         return {
             "total_segments": total_segments,
             "tokens": tokens,
-            "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)),
-            "currency": embedding_model.get_currency(),
+            "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)) if embedding_model else 0,
+            "currency": embedding_model.get_currency() if embedding_model else 'USD',
             "preview": preview_texts
         }
 
@@ -399,7 +406,8 @@ class IndexingRunner:
                 filter(UploadFile.id == data_source_info['upload_file_id']). \
                 one_or_none()
 
-            text_docs = FileExtractor.load(file_detail)
+            if file_detail:
+                text_docs = FileExtractor.load(file_detail)
         elif dataset_document.data_source_type == 'notion_import':
             loader = NotionLoader.from_document(dataset_document)
             text_docs = loader.load()
@@ -525,12 +533,13 @@ class IndexingRunner:
             documents = splitter.split_documents([text_doc])
             split_documents = []
             for document_node in documents:
-                doc_id = str(uuid.uuid4())
-                hash = helper.generate_text_hash(document_node.page_content)
-                document_node.metadata['doc_id'] = doc_id
-                document_node.metadata['doc_hash'] = hash
 
-                split_documents.append(document_node)
+                if document_node.page_content.strip():
+                    doc_id = str(uuid.uuid4())
+                    hash = helper.generate_text_hash(document_node.page_content)
+                    document_node.metadata['doc_id'] = doc_id
+                    document_node.metadata['doc_hash'] = hash
+                    split_documents.append(document_node)
             all_documents.extend(split_documents)
         # processing qa document
         if document_form == 'qa_model':
@@ -656,12 +665,13 @@ class IndexingRunner:
         """
         vector_index = IndexBuilder.get_index(dataset, 'high_quality')
         keyword_table_index = IndexBuilder.get_index(dataset, 'economy')
-
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=dataset.tenant_id,
-            model_provider_name=dataset.embedding_model_provider,
-            model_name=dataset.embedding_model
-        )
+        embedding_model = None
+        if dataset.indexing_technique == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=dataset.tenant_id,
+                model_provider_name=dataset.embedding_model_provider,
+                model_name=dataset.embedding_model
+            )
 
         # chunk nodes by chunk size
         indexing_start_at = time.perf_counter()
@@ -671,11 +681,11 @@ class IndexingRunner:
             # check document is paused
             self._check_document_paused_status(dataset_document.id)
             chunk_documents = documents[i:i + chunk_size]
-
-            tokens += sum(
-                embedding_model.get_num_tokens(document.page_content)
-                for document in chunk_documents
-            )
+            if dataset.indexing_technique == 'high_quality' or embedding_model:
+                tokens += sum(
+                    embedding_model.get_num_tokens(document.page_content)
+                    for document in chunk_documents
+                )
 
             # save vector index
             if vector_index:
@@ -691,6 +701,7 @@ class IndexingRunner:
                 DocumentSegment.status == "indexing"
             ).update({
                 DocumentSegment.status: "completed",
+                DocumentSegment.enabled: True,
                 DocumentSegment.completed_at: datetime.datetime.utcnow()
             })
 

api/core/login/login.py

@@ -0,0 +1,108 @@
+import os
+from functools import wraps
+
+import flask_login
+from flask import current_app
+from flask import g
+from flask import has_request_context
+from flask import request
+from flask_login import user_logged_in
+from flask_login.config import EXEMPT_METHODS
+from werkzeug.exceptions import Unauthorized
+from werkzeug.local import LocalProxy
+
+from extensions.ext_database import db
+from models.account import Account, Tenant, TenantAccountJoin
+
+#: A proxy for the current user. If no user is logged in, this will be an
+#: anonymous user
+current_user = LocalProxy(lambda: _get_user())
+
+
+def login_required(func):
+    """
+    If you decorate a view with this, it will ensure that the current user is
+    logged in and authenticated before calling the actual view. (If they are
+    not, it calls the :attr:`LoginManager.unauthorized` callback.) For
+    example::
+
+        @app.route('/post')
+        @login_required
+        def post():
+            pass
+
+    If there are only certain times you need to require that your user is
+    logged in, you can do so with::
+
+        if not current_user.is_authenticated:
+            return current_app.login_manager.unauthorized()
+
+    ...which is essentially the code that this function adds to your views.
+
+    It can be convenient to globally turn off authentication when unit testing.
+    To enable this, if the application configuration variable `LOGIN_DISABLED`
+    is set to `True`, this decorator will be ignored.
+
+    .. Note ::
+
+        Per `W3 guidelines for CORS preflight requests
+        <http://www.w3.org/TR/cors/#cross-origin-request-with-preflight-0>`_,
+        HTTP ``OPTIONS`` requests are exempt from login checks.
+
+    :param func: The view function to decorate.
+    :type func: function
+    """
+
+    @wraps(func)
+    def decorated_view(*args, **kwargs):
+        auth_header = request.headers.get('Authorization')
+        admin_api_key_enable = os.getenv('ADMIN_API_KEY_ENABLE', default='False')
+        if admin_api_key_enable:
+            if auth_header:
+                if ' ' not in auth_header:
+                    raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+                auth_scheme, auth_token = auth_header.split(None, 1)
+                auth_scheme = auth_scheme.lower()
+                if auth_scheme != 'bearer':
+                    raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+                admin_api_key = os.getenv('ADMIN_API_KEY')
+
+                if admin_api_key:
+                    if os.getenv('ADMIN_API_KEY') == auth_token:
+                        workspace_id = request.headers.get('X-WORKSPACE-ID')
+                        if workspace_id:
+                            tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \
+                                .filter(Tenant.id == workspace_id) \
+                                .filter(TenantAccountJoin.tenant_id == Tenant.id) \
+                                .filter(TenantAccountJoin.role == 'owner') \
+                                .one_or_none()
+                            if tenant_account_join:
+                                tenant, ta = tenant_account_join
+                                account = Account.query.filter_by(id=ta.account_id).first()
+                                # Login admin
+                                if account:
+                                    account.current_tenant = tenant
+                                    current_app.login_manager._update_request_context_with_user(account)
+                                    user_logged_in.send(current_app._get_current_object(), user=_get_user())
+        if request.method in EXEMPT_METHODS or current_app.config.get("LOGIN_DISABLED"):
+            pass
+        elif not current_user.is_authenticated:
+            return current_app.login_manager.unauthorized()
+
+        # flask 1.x compatibility
+        # current_app.ensure_sync is only available in Flask >= 2.0
+        if callable(getattr(current_app, "ensure_sync", None)):
+            return current_app.ensure_sync(func)(*args, **kwargs)
+        return func(*args, **kwargs)
+
+    return decorated_view
+
+
+def _get_user():
+    if has_request_context():
+        if "_login_user" not in g:
+            current_app.login_manager._load_user()
+
+        return g._login_user
+
+    return None

api/core/model_providers/model_provider_factory.py

@@ -63,6 +63,9 @@ class ModelProviderFactory:
         elif provider_name == 'openllm':
             from core.model_providers.providers.openllm_provider import OpenLLMProvider
             return OpenLLMProvider
+        elif provider_name == 'localai':
+            from core.model_providers.providers.localai_provider import LocalAIProvider
+            return LocalAIProvider
         else:
             raise NotImplementedError
 

api/core/model_providers/models/embedding/localai_embedding.py

@@ -0,0 +1,29 @@
+from langchain.embeddings import LocalAIEmbeddings
+
+from replicate.exceptions import ModelError, ReplicateError
+
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import BaseModelProvider
+from core.model_providers.models.embedding.base import BaseEmbedding
+
+
+class LocalAIEmbedding(BaseEmbedding):
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = LocalAIEmbeddings(
+            model=name,
+            openai_api_key="1",
+            openai_api_base=credentials['server_url'],
+        )
+
+        super().__init__(model_provider, client, name)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        if isinstance(ex, (ModelError, ReplicateError)):
+            return LLMBadRequestError(f"LocalAI embedding: {str(ex)}")
+        else:
+            return ex

api/core/model_providers/models/embedding/xinference_embedding.py

@@ -1,4 +1,4 @@
-from langchain.embeddings import XinferenceEmbeddings
+from core.third_party.langchain.embeddings.xinference_embedding import XinferenceEmbedding as XinferenceEmbeddings
 from replicate.exceptions import ModelError, ReplicateError
 
 from core.model_providers.error import LLMBadRequestError
@@ -14,7 +14,8 @@ class XinferenceEmbedding(BaseEmbedding):
         )
 
         client = XinferenceEmbeddings(
-            **credentials,
+            server_url=credentials['server_url'],
+            model_uid=credentials['model_uid'],
         )
 
         super().__init__(model_provider, client, name)

api/core/model_providers/models/llm/anthropic_model.py

@@ -75,7 +75,7 @@ class AnthropicModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True
 

api/core/model_providers/models/llm/azure_openai_model.py

@@ -141,6 +141,6 @@ class AzureOpenAIModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
-        return True
+    @property
+    def support_streaming(self):
+        return True

api/core/model_providers/models/llm/base.py

@@ -1,17 +1,24 @@
+import json
+import os
+import re
 from abc import abstractmethod
-from typing import List, Optional, Any, Union
+from typing import List, Optional, Any, Union, Tuple
 import decimal
 
 from langchain.callbacks.manager import Callbacks
+from langchain.memory.chat_memory import BaseChatMemory
 from langchain.schema import LLMResult, SystemMessage, AIMessage, HumanMessage, BaseMessage, ChatGeneration
 
 from core.callback_handler.std_out_callback_handler import DifyStreamingStdOutCallbackHandler, DifyStdOutCallbackHandler
 from core.model_providers.models.base import BaseProviderModel
-from core.model_providers.models.entity.message import PromptMessage, MessageType, LLMRunResult
+from core.model_providers.models.entity.message import PromptMessage, MessageType, LLMRunResult, to_prompt_messages
 from core.model_providers.models.entity.model_params import ModelType, ModelKwargs, ModelMode, ModelKwargsRules
 from core.model_providers.providers.base import BaseModelProvider
+from core.prompt.prompt_builder import PromptBuilder
+from core.prompt.prompt_template import JinjaPromptTemplate
 from core.third_party.langchain.llms.fake import FakeLLM
 import logging
+
 logger = logging.getLogger(__name__)
 
 
@@ -76,13 +83,14 @@ class BaseLLM(BaseProviderModel):
     def price_config(self) -> dict:
         def get_or_default():
             default_price_config = {
-                    'prompt': decimal.Decimal('0'),
-                    'completion': decimal.Decimal('0'),
-                    'unit': decimal.Decimal('0'),
-                    'currency': 'USD'
-                }
+                'prompt': decimal.Decimal('0'),
+                'completion': decimal.Decimal('0'),
+                'unit': decimal.Decimal('0'),
+                'currency': 'USD'
+            }
             rules = self.model_provider.get_rules()
-            price_config = rules['price_config'][self.base_model_name] if 'price_config' in rules else default_price_config
+            price_config = rules['price_config'][
+                self.base_model_name] if 'price_config' in rules else default_price_config
             price_config = {
                 'prompt': decimal.Decimal(price_config['prompt']),
                 'completion': decimal.Decimal(price_config['completion']),
@@ -90,7 +98,7 @@ class BaseLLM(BaseProviderModel):
                 'currency': price_config['currency']
             }
             return price_config
-        
+
         self._price_config = self._price_config if hasattr(self, '_price_config') else get_or_default()
 
         logger.debug(f"model: {self.name} price_config: {self._price_config}")
@@ -130,7 +138,7 @@ class BaseLLM(BaseProviderModel):
                 result = self._run(
                     messages=messages,
                     stop=stop,
-                    callbacks=callbacks if not (self.streaming and not self.support_streaming()) else None,
+                    callbacks=callbacks if not (self.streaming and not self.support_streaming) else None,
                     **kwargs
                 )
             except Exception as ex:
@@ -141,7 +149,7 @@ class BaseLLM(BaseProviderModel):
         else:
             completion_content = result.generations[0][0].text
 
-        if self.streaming and not self.support_streaming():
+        if self.streaming and not self.support_streaming:
             # use FakeLLM to simulate streaming when current model not support streaming but streaming is True
             prompts = self._get_prompt_from_messages(messages, ModelMode.CHAT)
             fake_llm = FakeLLM(
@@ -158,7 +166,8 @@ class BaseLLM(BaseProviderModel):
             total_tokens = result.llm_output['token_usage']['total_tokens']
         else:
             prompt_tokens = self.get_num_tokens(messages)
-            completion_tokens = self.get_num_tokens([PromptMessage(content=completion_content, type=MessageType.ASSISTANT)])
+            completion_tokens = self.get_num_tokens(
+                [PromptMessage(content=completion_content, type=MessageType.ASSISTANT)])
             total_tokens = prompt_tokens + completion_tokens
 
         self.model_provider.update_last_used()
@@ -289,10 +298,123 @@ class BaseLLM(BaseProviderModel):
         else:
             self.client.callbacks.extend(callbacks)
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return False
 
+    def get_prompt(self, mode: str,
+                   pre_prompt: str, inputs: dict,
+                   query: str,
+                   context: Optional[str],
+                   memory: Optional[BaseChatMemory]) -> \
+            Tuple[List[PromptMessage], Optional[List[str]]]:
+        prompt_rules = self._read_prompt_rules_from_file(self.prompt_file_name(mode))
+        prompt, stops = self._get_prompt_and_stop(prompt_rules, pre_prompt, inputs, query, context, memory)
+        return [PromptMessage(content=prompt)], stops
+
+    def prompt_file_name(self, mode: str) -> str:
+        if mode == 'completion':
+            return 'common_completion'
+        else:
+            return 'common_chat'
+
+    def _get_prompt_and_stop(self, prompt_rules: dict, pre_prompt: str, inputs: dict,
+                             query: str,
+                             context: Optional[str],
+                             memory: Optional[BaseChatMemory]) -> Tuple[str, Optional[list]]:
+        context_prompt_content = ''
+        if context and 'context_prompt' in prompt_rules:
+            prompt_template = JinjaPromptTemplate.from_template(template=prompt_rules['context_prompt'])
+            context_prompt_content = prompt_template.format(
+                context=context
+            )
+
+        pre_prompt_content = ''
+        if pre_prompt:
+            prompt_template = JinjaPromptTemplate.from_template(template=pre_prompt)
+            prompt_inputs = {k: inputs[k] for k in prompt_template.input_variables if k in inputs}
+            pre_prompt_content = prompt_template.format(
+                **prompt_inputs
+            )
+
+        prompt = ''
+        for order in prompt_rules['system_prompt_orders']:
+            if order == 'context_prompt':
+                prompt += context_prompt_content
+            elif order == 'pre_prompt':
+                prompt += (pre_prompt_content + '\n\n') if pre_prompt_content else ''
+
+        query_prompt = prompt_rules['query_prompt'] if 'query_prompt' in prompt_rules else '{{query}}'
+
+        if memory and 'histories_prompt' in prompt_rules:
+            # append chat histories
+            tmp_human_message = PromptBuilder.to_human_message(
+                prompt_content=prompt + query_prompt,
+                inputs={
+                    'query': query
+                }
+            )
+
+            if self.model_rules.max_tokens.max:
+                curr_message_tokens = self.get_num_tokens(to_prompt_messages([tmp_human_message]))
+                max_tokens = self.model_kwargs.max_tokens
+                rest_tokens = self.model_rules.max_tokens.max - max_tokens - curr_message_tokens
+                rest_tokens = max(rest_tokens, 0)
+            else:
+                rest_tokens = 2000
+
+            memory.human_prefix = prompt_rules['human_prefix'] if 'human_prefix' in prompt_rules else 'Human'
+            memory.ai_prefix = prompt_rules['assistant_prefix'] if 'assistant_prefix' in prompt_rules else 'Assistant'
+
+            histories = self._get_history_messages_from_memory(memory, rest_tokens)
+            prompt_template = JinjaPromptTemplate.from_template(template=prompt_rules['histories_prompt'])
+            histories_prompt_content = prompt_template.format(
+                histories=histories
+            )
+
+            prompt = ''
+            for order in prompt_rules['system_prompt_orders']:
+                if order == 'context_prompt':
+                    prompt += context_prompt_content
+                elif order == 'pre_prompt':
+                    prompt += (pre_prompt_content + '\n') if pre_prompt_content else ''
+                elif order == 'histories_prompt':
+                    prompt += histories_prompt_content
+
+        prompt_template = JinjaPromptTemplate.from_template(template=query_prompt)
+        query_prompt_content = prompt_template.format(
+            query=query
+        )
+
+        prompt += query_prompt_content
+
+        prompt = re.sub(r'<\|.*?\|>', '', prompt)
+
+        stops = prompt_rules.get('stops')
+        if stops is not None and len(stops) == 0:
+            stops = None
+
+        return prompt, stops
+
+    def _read_prompt_rules_from_file(self, prompt_name: str) -> dict:
+        # Get the absolute path of the subdirectory
+        prompt_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))),
+            'prompt/generate_prompts')
+
+        json_file_path = os.path.join(prompt_path, f'{prompt_name}.json')
+        # Open the JSON file and read its content
+        with open(json_file_path, 'r') as json_file:
+            return json.load(json_file)
+
+    def _get_history_messages_from_memory(self, memory: BaseChatMemory,
+                                          max_token_limit: int) -> str:
+        """Get memory messages."""
+        memory.max_token_limit = max_token_limit
+        memory_key = memory.memory_variables[0]
+        external_context = memory.load_memory_variables({})
+        return external_context[memory_key]
+
     def _get_prompt_from_messages(self, messages: List[PromptMessage],
                                   model_mode: Optional[ModelMode] = None) -> Union[str | List[BaseMessage]]:
         if not model_mode:

api/core/model_providers/models/llm/chatglm_model.py

@@ -61,7 +61,3 @@ class ChatGLMModel(BaseLLM):
             return LLMBadRequestError(f"ChatGLM: {str(ex)}")
         else:
             return ex
-
-    @classmethod
-    def support_streaming(cls):
-        return False

api/core/model_providers/models/llm/huggingface_hub_model.py

@@ -1,16 +1,14 @@
-import decimal
-from functools import wraps
 from typing import List, Optional, Any
 
 from langchain import HuggingFaceHub
 from langchain.callbacks.manager import Callbacks
-from langchain.llms import HuggingFaceEndpoint
 from langchain.schema import LLMResult
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.models.llm.base import BaseLLM
-from core.model_providers.models.entity.message import PromptMessage, MessageType
+from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.huggingface_endpoint_llm import HuggingFaceEndpointLLM
 
 
 class HuggingfaceHubModel(BaseLLM):
@@ -19,12 +17,18 @@ class HuggingfaceHubModel(BaseLLM):
     def _init_client(self) -> Any:
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
         if self.credentials['huggingfacehub_api_type'] == 'inference_endpoints':
-            client = HuggingFaceEndpoint(
+            streaming = self.streaming
+
+            if 'baichuan' in self.name.lower():
+                streaming = False
+
+            client = HuggingFaceEndpointLLM(
                 endpoint_url=self.credentials['huggingfacehub_endpoint_url'],
-                task='text2text-generation',
+                task=self.credentials['task_type'],
                 model_kwargs=provider_model_kwargs,
                 huggingfacehub_api_token=self.credentials['huggingfacehub_api_token'],
                 callbacks=self.callbacks,
+                streaming=streaming
             )
         else:
             client = HuggingFaceHub(
@@ -62,6 +66,15 @@ class HuggingfaceHubModel(BaseLLM):
         prompts = self._get_prompt_from_messages(messages)
         return self._client.get_num_tokens(prompts)
 
+    def prompt_file_name(self, mode: str) -> str:
+        if 'baichuan' in self.name.lower():
+            if mode == 'completion':
+                return 'baichuan_completion'
+            else:
+                return 'baichuan_chat'
+        else:
+            return super().prompt_file_name(mode)
+
     def _set_model_kwargs(self, model_kwargs: ModelKwargs):
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, model_kwargs)
         self.client.model_kwargs = provider_model_kwargs
@@ -69,7 +82,10 @@ class HuggingfaceHubModel(BaseLLM):
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"Huggingface Hub: {str(ex)}")
 
-    @classmethod
-    def support_streaming(cls):
-        return False
+    @property
+    def support_streaming(self):
+        if self.credentials['huggingfacehub_api_type'] == 'inference_endpoints':
+            if 'baichuan' in self.name.lower():
+                return False
 
+        return True

api/core/model_providers/models/llm/localai_model.py

@@ -0,0 +1,131 @@
+import logging
+from typing import List, Optional, Any
+
+import openai
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import LLMResult, get_buffer_string
+
+from core.model_providers.error import LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError, \
+    LLMRateLimitError, LLMAuthorizationError
+from core.model_providers.providers.base import BaseModelProvider
+from core.third_party.langchain.llms.chat_open_ai import EnhanceChatOpenAI
+from core.third_party.langchain.llms.open_ai import EnhanceOpenAI
+from core.model_providers.models.llm.base import BaseLLM
+from core.model_providers.models.entity.message import PromptMessage
+from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+
+
+class LocalAIModel(BaseLLM):
+    def __init__(self, model_provider: BaseModelProvider,
+                 name: str,
+                 model_kwargs: ModelKwargs,
+                 streaming: bool = False,
+                 callbacks: Callbacks = None):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        if credentials['completion_type'] == 'chat_completion':
+            self.model_mode = ModelMode.CHAT
+        else:
+            self.model_mode = ModelMode.COMPLETION
+
+        super().__init__(model_provider, name, model_kwargs, streaming, callbacks)
+
+    def _init_client(self) -> Any:
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
+        if self.model_mode == ModelMode.COMPLETION:
+            client = EnhanceOpenAI(
+                model_name=self.name,
+                streaming=self.streaming,
+                callbacks=self.callbacks,
+                request_timeout=60,
+                openai_api_key="1",
+                openai_api_base=self.credentials['server_url'] + '/v1',
+                **provider_model_kwargs
+            )
+        else:
+            extra_model_kwargs = {
+                'top_p': provider_model_kwargs.get('top_p')
+            }
+
+            client = EnhanceChatOpenAI(
+                model_name=self.name,
+                temperature=provider_model_kwargs.get('temperature'),
+                max_tokens=provider_model_kwargs.get('max_tokens'),
+                model_kwargs=extra_model_kwargs,
+                streaming=self.streaming,
+                callbacks=self.callbacks,
+                request_timeout=60,
+                openai_api_key="1",
+                openai_api_base=self.credentials['server_url'] + '/v1'
+            )
+
+        return client
+
+    def _run(self, messages: List[PromptMessage],
+             stop: Optional[List[str]] = None,
+             callbacks: Callbacks = None,
+             **kwargs) -> LLMResult:
+        """
+        run predict by prompt messages and stop words.
+
+        :param messages:
+        :param stop:
+        :param callbacks:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        return self._client.generate([prompts], stop, callbacks)
+
+    def get_num_tokens(self, messages: List[PromptMessage]) -> int:
+        """
+        get num tokens of prompt messages.
+
+        :param messages:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        if isinstance(prompts, str):
+            return self._client.get_num_tokens(prompts)
+        else:
+            return max(sum([self._client.get_num_tokens(get_buffer_string([m])) for m in prompts]) - len(prompts), 0)
+
+    def _set_model_kwargs(self, model_kwargs: ModelKwargs):
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, model_kwargs)
+        if self.model_mode == ModelMode.COMPLETION:
+            for k, v in provider_model_kwargs.items():
+                if hasattr(self.client, k):
+                    setattr(self.client, k, v)
+        else:
+            extra_model_kwargs = {
+                'top_p': provider_model_kwargs.get('top_p')
+            }
+
+            self.client.temperature = provider_model_kwargs.get('temperature')
+            self.client.max_tokens = provider_model_kwargs.get('max_tokens')
+            self.client.model_kwargs = extra_model_kwargs
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        if isinstance(ex, openai.error.InvalidRequestError):
+            logging.warning("Invalid request to LocalAI API.")
+            return LLMBadRequestError(str(ex))
+        elif isinstance(ex, openai.error.APIConnectionError):
+            logging.warning("Failed to connect to LocalAI API.")
+            return LLMAPIConnectionError(ex.__class__.__name__ + ":" + str(ex))
+        elif isinstance(ex, (openai.error.APIError, openai.error.ServiceUnavailableError, openai.error.Timeout)):
+            logging.warning("LocalAI service unavailable.")
+            return LLMAPIUnavailableError(ex.__class__.__name__ + ":" + str(ex))
+        elif isinstance(ex, openai.error.RateLimitError):
+            return LLMRateLimitError(str(ex))
+        elif isinstance(ex, openai.error.AuthenticationError):
+            return LLMAuthorizationError(str(ex))
+        elif isinstance(ex, openai.error.OpenAIError):
+            return LLMBadRequestError(ex.__class__.__name__ + ":" + str(ex))
+        else:
+            return ex
+
+    @classmethod
+    def support_streaming(cls):
+        return True

api/core/model_providers/models/llm/openai_model.py

@@ -154,8 +154,8 @@ class OpenAIModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True
 
     # def is_model_valid_or_raise(self):

api/core/model_providers/models/llm/openllm_model.py

@@ -1,13 +1,13 @@
 from typing import List, Optional, Any
 
 from langchain.callbacks.manager import Callbacks
-from langchain.llms import OpenLLM
 from langchain.schema import LLMResult
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.openllm import OpenLLM
 
 
 class OpenLLMModel(BaseLLM):
@@ -19,7 +19,7 @@ class OpenLLMModel(BaseLLM):
         client = OpenLLM(
             server_url=self.credentials.get('server_url'),
             callbacks=self.callbacks,
-            **self.provider_model_kwargs
+            llm_kwargs=self.provider_model_kwargs
         )
 
         return client
@@ -49,12 +49,17 @@ class OpenLLMModel(BaseLLM):
         prompts = self._get_prompt_from_messages(messages)
         return max(self._client.get_num_tokens(prompts), 0)
 
+    def prompt_file_name(self, mode: str) -> str:
+        if 'baichuan' in self.name.lower():
+            if mode == 'completion':
+                return 'baichuan_completion'
+            else:
+                return 'baichuan_chat'
+        else:
+            return super().prompt_file_name(mode)
+
     def _set_model_kwargs(self, model_kwargs: ModelKwargs):
         pass
 
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"OpenLLM: {str(ex)}")
-
-    @classmethod
-    def support_streaming(cls):
-        return False

api/core/model_providers/models/llm/replicate_model.py

@@ -91,6 +91,6 @@ class ReplicateModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
-        return True
+    @property
+    def support_streaming(self):
+        return True

api/core/model_providers/models/llm/spark_model.py

@@ -65,6 +65,6 @@ class SparkModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
-        return True
+    @property
+    def support_streaming(self):
+        return True

api/core/model_providers/models/llm/tongyi_model.py

@@ -69,6 +69,6 @@ class TongyiModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True

api/core/model_providers/models/llm/wenxin_model.py

@@ -57,7 +57,3 @@ class WenxinModel(BaseLLM):
 
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"Wenxin: {str(ex)}")
-
-    @classmethod
-    def support_streaming(cls):
-        return False

api/core/model_providers/models/llm/xinference_model.py

@@ -1,13 +1,13 @@
 from typing import List, Optional, Any
 
 from langchain.callbacks.manager import Callbacks
-from langchain.llms import Xinference
 from langchain.schema import LLMResult
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.xinference_llm import XinferenceLLM
 
 
 class XinferenceModel(BaseLLM):
@@ -16,8 +16,9 @@ class XinferenceModel(BaseLLM):
     def _init_client(self) -> Any:
         self.provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
 
-        client = Xinference(
-            **self.credentials,
+        client = XinferenceLLM(
+            server_url=self.credentials['server_url'],
+            model_uid=self.credentials['model_uid'],
         )
 
         client.callbacks = self.callbacks
@@ -58,12 +59,21 @@ class XinferenceModel(BaseLLM):
         prompts = self._get_prompt_from_messages(messages)
         return max(self._client.get_num_tokens(prompts), 0)
 
+    def prompt_file_name(self, mode: str) -> str:
+        if 'baichuan' in self.name.lower():
+            if mode == 'completion':
+                return 'baichuan_completion'
+            else:
+                return 'baichuan_chat'
+        else:
+            return super().prompt_file_name(mode)
+
     def _set_model_kwargs(self, model_kwargs: ModelKwargs):
         pass
 
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"Xinference: {str(ex)}")
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True

api/core/model_providers/providers/huggingface_hub_provider.py

@@ -2,7 +2,6 @@ import json
 from typing import Type
 
 from huggingface_hub import HfApi
-from langchain.llms import HuggingFaceEndpoint
 
 from core.helper import encrypter
 from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType
@@ -10,6 +9,7 @@ from core.model_providers.models.llm.huggingface_hub_model import HuggingfaceHub
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 
 from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.llms.huggingface_endpoint_llm import HuggingFaceEndpointLLM
 from models.provider import ProviderType
 
 
@@ -51,7 +51,7 @@ class HuggingfaceHubProvider(BaseModelProvider):
             top_p=KwargRule[float](min=0.01, max=0.99, default=0.7),
             presence_penalty=KwargRule[float](enabled=False),
             frequency_penalty=KwargRule[float](enabled=False),
-            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=1500, default=200),
+            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=200),
         )
 
     @classmethod
@@ -85,10 +85,16 @@ class HuggingfaceHubProvider(BaseModelProvider):
             if 'huggingfacehub_endpoint_url' not in credentials:
                 raise CredentialsValidateFailedError('Hugging Face Hub Endpoint URL must be provided.')
 
+            if 'task_type' not in credentials:
+                raise CredentialsValidateFailedError('Task Type must be provided.')
+
+            if credentials['task_type'] not in ("text2text-generation", "text-generation", "summarization"):
+                raise CredentialsValidateFailedError('Task Type must be one of text2text-generation, text-generation, summarization.')
+
             try:
-                llm = HuggingFaceEndpoint(
+                llm = HuggingFaceEndpointLLM(
                     endpoint_url=credentials['huggingfacehub_endpoint_url'],
-                    task="text2text-generation",
+                    task=credentials['task_type'],
                     model_kwargs={"temperature": 0.5, "max_new_tokens": 200},
                     huggingfacehub_api_token=credentials['huggingfacehub_api_token']
                 )
@@ -160,6 +166,10 @@ class HuggingfaceHubProvider(BaseModelProvider):
             }
 
         credentials = json.loads(provider_model.encrypted_config)
+
+        if 'task_type' not in credentials:
+            credentials['task_type'] = 'text-generation'
+
         if credentials['huggingfacehub_api_token']:
             credentials['huggingfacehub_api_token'] = encrypter.decrypt_token(
                 self.provider.tenant_id,

api/core/model_providers/providers/localai_provider.py

@@ -0,0 +1,164 @@
+import json
+from typing import Type
+
+from langchain.embeddings import LocalAIEmbeddings
+from langchain.schema import HumanMessage
+
+from core.helper import encrypter
+from core.model_providers.models.embedding.localai_embedding import LocalAIEmbedding
+from core.model_providers.models.entity.model_params import ModelKwargsRules, ModelType, KwargRule
+from core.model_providers.models.llm.localai_model import LocalAIModel
+from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
+
+from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.llms.chat_open_ai import EnhanceChatOpenAI
+from core.third_party.langchain.llms.open_ai import EnhanceOpenAI
+from models.provider import ProviderType
+
+
+class LocalAIProvider(BaseModelProvider):
+    @property
+    def provider_name(self):
+        """
+        Returns the name of a provider.
+        """
+        return 'localai'
+
+    def _get_fixed_model_list(self, model_type: ModelType) -> list[dict]:
+        return []
+
+    def get_model_class(self, model_type: ModelType) -> Type[BaseProviderModel]:
+        """
+        Returns the model class.
+
+        :param model_type:
+        :return:
+        """
+        if model_type == ModelType.TEXT_GENERATION:
+            model_class = LocalAIModel
+        elif model_type == ModelType.EMBEDDINGS:
+            model_class = LocalAIEmbedding
+        else:
+            raise NotImplementedError
+
+        return model_class
+
+    def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> ModelKwargsRules:
+        """
+        get model parameter rules.
+
+        :param model_name:
+        :param model_type:
+        :return:
+        """
+        return ModelKwargsRules(
+            temperature=KwargRule[float](min=0, max=2, default=0.7),
+            top_p=KwargRule[float](min=0, max=1, default=1),
+            max_tokens=KwargRule[int](min=10, max=4097, default=16),
+        )
+
+    @classmethod
+    def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
+        """
+        check model credentials valid.
+
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        """
+        if 'server_url' not in credentials:
+            raise CredentialsValidateFailedError('LocalAI Server URL must be provided.')
+
+        try:
+            if model_type == ModelType.EMBEDDINGS:
+                model = LocalAIEmbeddings(
+                    model=model_name,
+                    openai_api_key='1',
+                    openai_api_base=credentials['server_url']
+                )
+
+                model.embed_query("ping")
+            else:
+                if ('completion_type' not in credentials
+                        or credentials['completion_type'] not in ['completion', 'chat_completion']):
+                    raise CredentialsValidateFailedError('LocalAI Completion Type must be provided.')
+
+                if credentials['completion_type'] == 'chat_completion':
+                    model = EnhanceChatOpenAI(
+                        model_name=model_name,
+                        openai_api_key='1',
+                        openai_api_base=credentials['server_url'] + '/v1',
+                        max_tokens=10,
+                        request_timeout=60,
+                    )
+
+                    model([HumanMessage(content='ping')])
+                else:
+                    model = EnhanceOpenAI(
+                        model_name=model_name,
+                        openai_api_key='1',
+                        openai_api_base=credentials['server_url'] + '/v1',
+                        max_tokens=10,
+                        request_timeout=60,
+                    )
+
+                    model('ping')
+        except Exception as ex:
+            raise CredentialsValidateFailedError(str(ex))
+
+    @classmethod
+    def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
+                                  credentials: dict) -> dict:
+        """
+        encrypt model credentials for save.
+
+        :param tenant_id:
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        :return:
+        """
+        credentials['server_url'] = encrypter.encrypt_token(tenant_id, credentials['server_url'])
+        return credentials
+
+    def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
+        """
+        get credentials for llm use.
+
+        :param model_name:
+        :param model_type:
+        :param obfuscated:
+        :return:
+        """
+        if self.provider.provider_type != ProviderType.CUSTOM.value:
+            raise NotImplementedError
+
+        provider_model = self._get_provider_model(model_name, model_type)
+
+        if not provider_model.encrypted_config:
+            return {
+                'server_url': None,
+            }
+
+        credentials = json.loads(provider_model.encrypted_config)
+        if credentials['server_url']:
+            credentials['server_url'] = encrypter.decrypt_token(
+                self.provider.tenant_id,
+                credentials['server_url']
+            )
+
+            if obfuscated:
+                credentials['server_url'] = encrypter.obfuscated_token(credentials['server_url'])
+
+        return credentials
+
+    @classmethod
+    def is_provider_credentials_valid_or_raise(cls, credentials: dict):
+        return
+
+    @classmethod
+    def encrypt_provider_credentials(cls, tenant_id: str, credentials: dict) -> dict:
+        return {}
+
+    def get_provider_credentials(self, obfuscated: bool = False) -> dict:
+        return {}

api/core/model_providers/providers/openllm_provider.py

@@ -1,14 +1,13 @@
 import json
 from typing import Type
 
-from langchain.llms import OpenLLM
-
 from core.helper import encrypter
 from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType
 from core.model_providers.models.llm.openllm_model import OpenLLMModel
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 
 from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.llms.openllm import OpenLLM
 from models.provider import ProviderType
 
 
@@ -46,11 +45,11 @@ class OpenLLMProvider(BaseModelProvider):
         :return:
         """
         return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
+            temperature=KwargRule[float](min=0.01, max=2, default=1),
             top_p=KwargRule[float](min=0, max=1, default=0.7),
             presence_penalty=KwargRule[float](min=-2, max=2, default=0),
             frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
-            max_tokens=KwargRule[int](min=10, max=4000, default=128),
+            max_tokens=KwargRule[int](alias='max_new_tokens', min=10, max=4000, default=128),
         )
 
     @classmethod
@@ -71,7 +70,9 @@ class OpenLLMProvider(BaseModelProvider):
             }
 
             llm = OpenLLM(
-                max_tokens=10,
+                llm_kwargs={
+                    'max_new_tokens': 10
+                },
                 **credential_kwargs
             )
 

api/core/model_providers/providers/xinference_provider.py

@@ -1,7 +1,7 @@
 import json
 from typing import Type
 
-from langchain.llms import Xinference
+import requests
 
 from core.helper import encrypter
 from core.model_providers.models.embedding.xinference_embedding import XinferenceEmbedding
@@ -10,6 +10,7 @@ from core.model_providers.models.llm.xinference_model import XinferenceModel
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 
 from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.llms.xinference_llm import XinferenceLLM
 from models.provider import ProviderType
 
 
@@ -48,13 +49,32 @@ class XinferenceProvider(BaseModelProvider):
         :param model_type:
         :return:
         """
-        return ModelKwargsRules(
-            temperature=KwargRule[float](min=0, max=2, default=1),
-            top_p=KwargRule[float](min=0, max=1, default=0.7),
-            presence_penalty=KwargRule[float](min=-2, max=2, default=0),
-            frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
-            max_tokens=KwargRule[int](min=10, max=4000, default=256),
-        )
+        credentials = self.get_model_credentials(model_name, model_type)
+        if credentials['model_format'] == "ggmlv3" and credentials["model_handle_type"] == "chatglm":
+            return ModelKwargsRules(
+                temperature=KwargRule[float](min=0.01, max=2, default=1),
+                top_p=KwargRule[float](min=0, max=1, default=0.7),
+                presence_penalty=KwargRule[float](enabled=False),
+                frequency_penalty=KwargRule[float](enabled=False),
+                max_tokens=KwargRule[int](min=10, max=4000, default=256),
+            )
+        elif credentials['model_format'] == "ggmlv3":
+            return ModelKwargsRules(
+                temperature=KwargRule[float](min=0.01, max=2, default=1),
+                top_p=KwargRule[float](min=0, max=1, default=0.7),
+                presence_penalty=KwargRule[float](min=-2, max=2, default=0),
+                frequency_penalty=KwargRule[float](min=-2, max=2, default=0),
+                max_tokens=KwargRule[int](min=10, max=4000, default=256),
+            )
+        else:
+            return ModelKwargsRules(
+                temperature=KwargRule[float](min=0.01, max=2, default=1),
+                top_p=KwargRule[float](min=0, max=1, default=0.7),
+                presence_penalty=KwargRule[float](enabled=False),
+                frequency_penalty=KwargRule[float](enabled=False),
+                max_tokens=KwargRule[int](min=10, max=4000, default=256),
+            )
+
 
     @classmethod
     def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
@@ -77,11 +97,11 @@ class XinferenceProvider(BaseModelProvider):
                 'model_uid': credentials['model_uid'],
             }
 
-            llm = Xinference(
+            llm = XinferenceLLM(
                 **credential_kwargs
             )
 
-            llm("ping", generate_config={'max_tokens': 10})
+            llm("ping")
         except Exception as ex:
             raise CredentialsValidateFailedError(str(ex))
 
@@ -97,7 +117,11 @@ class XinferenceProvider(BaseModelProvider):
         :param credentials:
         :return:
         """
+        extra_credentials = cls._get_extra_credentials(credentials)
+        credentials.update(extra_credentials)
+
         credentials['server_url'] = encrypter.encrypt_token(tenant_id, credentials['server_url'])
+
         return credentials
 
     def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
@@ -132,6 +156,30 @@ class XinferenceProvider(BaseModelProvider):
 
         return credentials
 
+    @classmethod
+    def _get_extra_credentials(self, credentials: dict) -> dict:
+        url = f"{credentials['server_url']}/v1/models/{credentials['model_uid']}"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get the model description, detail: {response.json()['detail']}"
+            )
+        desc = response.json()
+
+        extra_credentials = {
+            'model_format': desc['model_format'],
+        }
+        if desc["model_format"] == "ggmlv3" and "chatglm" in desc["model_name"]:
+            extra_credentials['model_handle_type'] = 'chatglm'
+        elif "generate" in desc["model_ability"]:
+            extra_credentials['model_handle_type'] = 'generate'
+        elif "chat" in desc["model_ability"]:
+            extra_credentials['model_handle_type'] = 'chat'
+        else:
+            raise NotImplementedError(f"Model handle type not supported.")
+
+        return extra_credentials
+
     @classmethod
     def is_provider_credentials_valid_or_raise(cls, credentials: dict):
         return

api/core/model_providers/rules/_providers.json

@@ -10,5 +10,6 @@
   "replicate",
   "huggingface_hub",
   "xinference",
-  "openllm"
+  "openllm",
+  "localai"
 ]

api/core/model_providers/rules/localai.json

@@ -0,0 +1,7 @@
+{
+    "support_provider_types": [
+        "custom"
+    ],
+    "system_config": null,
+    "model_flexibility": "configurable"
+}

api/core/prompt/generate_prompts/baichuan_chat.json

@@ -0,0 +1,13 @@
+{
+  "human_prefix": "用户",
+  "assistant_prefix": "助手",
+  "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料，给出全面专业的解释，但不会过度演绎。同时回答中不会暴露引用的材料：\n\n```\n{{context}}\n```\n\n",
+  "histories_prompt": "用户和助手的历史对话内容如下：\n```\n{{histories}}\n```\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt",
+    "histories_prompt"
+  ],
+  "query_prompt": "用户：{{query}}",
+  "stops": ["用户："]
+}

api/core/prompt/generate_prompts/baichuan_completion.json

@@ -0,0 +1,9 @@
+{
+  "context_prompt": "用户在与一个客观的助手对话。助手会尊重找到的材料，给出全面专业的解释，但不会过度演绎。同时回答中不会暴露引用的材料：\n\n```\n{{context}}\n```\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt"
+  ],
+  "query_prompt": "{{query}}",
+  "stops": null
+}

api/core/prompt/generate_prompts/common_chat.json

@@ -0,0 +1,13 @@
+{
+  "human_prefix": "Human",
+  "assistant_prefix": "Assistant",
+  "context_prompt": "Use the following context as your learned knowledge, inside <context></context> XML tags.\n\n<context>\n{{context}}\n</context>\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n",
+  "histories_prompt": "Here is the chat histories between human and assistant, inside <histories></histories> XML tags.\n\n<histories>\n{{histories}}\n</histories>\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt",
+    "histories_prompt"
+  ],
+  "query_prompt": "Human: {{query}}\n\nAssistant: ",
+  "stops": ["\nHuman:", "</histories>"]
+}

api/core/prompt/generate_prompts/common_completion.json

@@ -0,0 +1,9 @@
+{
+  "context_prompt": "Use the following context as your learned knowledge, inside <context></context> XML tags.\n\n<context>\n{{context}}\n</context>\n\nWhen answer to user:\n- If you don't know, just say that you don't know.\n- If you don't know when you are not sure, ask for clarification.\nAvoid mentioning that you obtained the information from the context.\nAnd answer according to the language of the user's question.\n\n",
+  "system_prompt_orders": [
+    "context_prompt",
+    "pre_prompt"
+  ],
+  "query_prompt": "{{query}}",
+  "stops": null
+}

api/core/third_party/langchain/embeddings/xinference_embedding.py

@@ -0,0 +1,21 @@
+from typing import List
+
+import numpy as np
+from langchain.embeddings import XinferenceEmbeddings
+
+
+class XinferenceEmbedding(XinferenceEmbeddings):
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        vectors = super().embed_documents(texts)
+
+        normalized_vectors = [(vector / np.linalg.norm(vector)).tolist() for vector in vectors]
+
+        return normalized_vectors
+
+    def embed_query(self, text: str) -> List[float]:
+        vector = super().embed_query(text)
+
+        normalized_vector = (vector / np.linalg.norm(vector)).tolist()
+
+        return normalized_vector

api/core/third_party/langchain/llms/chat_open_ai.py

@@ -42,7 +42,8 @@ class EnhanceChatOpenAI(ChatOpenAI):
         return {
             **super()._default_params,
             "api_type": 'openai',
-            "api_base": os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+            "api_base": self.openai_api_base if self.openai_api_base
+            else os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
             "api_version": None,
             "api_key": self.openai_api_key,
             "organization": self.openai_organization if self.openai_organization else None,

api/core/third_party/langchain/llms/huggingface_endpoint_llm.py

@@ -0,0 +1,128 @@
+from typing import Dict, Any, Optional, List, Iterable, Iterator
+
+from huggingface_hub import InferenceClient
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.embeddings.huggingface_hub import VALID_TASKS
+from langchain.llms import HuggingFaceEndpoint
+from langchain.llms.utils import enforce_stop_tokens
+from pydantic import root_validator
+
+from langchain.utils import get_from_dict_or_env
+
+
+class HuggingFaceEndpointLLM(HuggingFaceEndpoint):
+    """HuggingFace Endpoint models.
+
+    To use, you should have the ``huggingface_hub`` python package installed, and the
+    environment variable ``HUGGINGFACEHUB_API_TOKEN`` set with your API token, or pass
+    it as a named parameter to the constructor.
+
+    Only supports `text-generation` and `text2text-generation` for now.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.llms import HuggingFaceEndpoint
+            endpoint_url = (
+                "https://abcdefghijklmnop.us-east-1.aws.endpoints.huggingface.cloud"
+            )
+            hf = HuggingFaceEndpoint(
+                endpoint_url=endpoint_url,
+                huggingfacehub_api_token="my-api-key"
+            )
+    """
+    client: Any
+    streaming: bool = False
+
+    @root_validator(allow_reuse=True)
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        huggingfacehub_api_token = get_from_dict_or_env(
+            values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
+        )
+
+        values['client'] = InferenceClient(values['endpoint_url'], token=huggingfacehub_api_token)
+
+        values["huggingfacehub_api_token"] = huggingfacehub_api_token
+        return values
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to HuggingFace Hub's inference endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+
+        Example:
+            .. code-block:: python
+
+                response = hf("Tell me a joke.")
+        """
+        _model_kwargs = self.model_kwargs or {}
+
+        # payload samples
+        params = {**_model_kwargs, **kwargs}
+
+        # generation parameter
+        gen_kwargs = {
+            **params,
+            'stop_sequences': stop
+        }
+
+        response = self.client.text_generation(prompt, stream=self.streaming, details=True, **gen_kwargs)
+
+        if self.streaming and isinstance(response, Iterable):
+            combined_text_output = ""
+            for token in self._stream_response(response, run_manager):
+                combined_text_output += token
+            completion = combined_text_output
+        else:
+            completion = response.generated_text
+
+        if self.task == "text-generation":
+            text = completion
+            # Remove prompt if included in generated text.
+            if text.startswith(prompt):
+                text = text[len(prompt) :]
+        elif self.task == "text2text-generation":
+            text = completion
+        else:
+            raise ValueError(
+                f"Got invalid task {self.task}, "
+                f"currently only {VALID_TASKS} are supported"
+            )
+
+        if stop is not None:
+            # This is a bit hacky, but I can't figure out a better way to enforce
+            # stop tokens when making calls to huggingface_hub.
+            text = enforce_stop_tokens(text, stop)
+
+        return text
+
+    def _stream_response(
+            self,
+            response: Iterable,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+    ) -> Iterator[str]:
+        for r in response:
+            # skip special tokens
+            if r.token.special:
+                continue
+
+            token = r.token.text
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    token=token, verbose=self.verbose, log_probs=None
+                )
+
+            # yield the generated token
+            yield token

api/core/third_party/langchain/llms/open_ai.py

@@ -1,7 +1,10 @@
 import os
 
-from typing import Dict, Any, Mapping, Optional, Union, Tuple
+from typing import Dict, Any, Mapping, Optional, Union, Tuple, List, Iterator
 from langchain import OpenAI
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.openai import completion_with_retry, _stream_response_to_generation_chunk
+from langchain.schema.output import GenerationChunk
 from pydantic import root_validator
 
 
@@ -33,7 +36,8 @@ class EnhanceOpenAI(OpenAI):
     def _invocation_params(self) -> Dict[str, Any]:
         return {**super()._invocation_params, **{
             "api_type": 'openai',
-            "api_base": os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+            "api_base": self.openai_api_base if self.openai_api_base
+            else os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
             "api_version": None,
             "api_key": self.openai_api_key,
             "organization": self.openai_organization if self.openai_organization else None,
@@ -43,8 +47,33 @@ class EnhanceOpenAI(OpenAI):
     def _identifying_params(self) -> Mapping[str, Any]:
         return {**super()._identifying_params, **{
             "api_type": 'openai',
-            "api_base": os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+            "api_base": self.openai_api_base if self.openai_api_base
+            else os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
             "api_version": None,
             "api_key": self.openai_api_key,
             "organization": self.openai_organization if self.openai_organization else None,
         }}
+
+    def _stream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        params = {**self._invocation_params, **kwargs, "stream": True}
+        self.get_sub_prompts(params, [prompt], stop)  # this mutates params
+        for stream_resp in completion_with_retry(
+            self, prompt=prompt, run_manager=run_manager, **params
+        ):
+            if 'text' in stream_resp["choices"][0]:
+                chunk = _stream_response_to_generation_chunk(stream_resp)
+                yield chunk
+                if run_manager:
+                    run_manager.on_llm_new_token(
+                        chunk.text,
+                        verbose=self.verbose,
+                        logprobs=chunk.generation_info["logprobs"]
+                        if chunk.generation_info
+                        else None,
+                    )

api/core/third_party/langchain/llms/openllm.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import logging
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+)
+
+import requests
+from langchain.llms.utils import enforce_stop_tokens
+from pydantic import Field
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain.llms.base import LLM
+
+logger = logging.getLogger(__name__)
+
+
+class OpenLLM(LLM):
+    """OpenLLM, supporting both in-process model
+    instance and remote OpenLLM servers.
+
+    If you have a OpenLLM server running, you can also use it remotely:
+        .. code-block:: python
+
+            from langchain.llms import OpenLLM
+            llm = OpenLLM(server_url='http://localhost:3000')
+            llm("What is the difference between a duck and a goose?")
+    """
+
+    server_url: Optional[str] = None
+    """Optional server URL that currently runs a LLMServer with 'openllm start'."""
+    llm_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Key word arguments to be passed to openllm.LLM"""
+
+    @property
+    def _llm_type(self) -> str:
+        return "openllm"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> str:
+        params = {
+            "prompt": prompt,
+            "llm_config": self.llm_kwargs
+        }
+
+        headers = {"Content-Type": "application/json"}
+        response = requests.post(
+            f'{self.server_url}/v1/generate',
+            headers=headers,
+            json=params
+        )
+
+        if not response.ok:
+            raise ValueError(f"OpenLLM HTTP {response.status_code} error: {response.text}")
+
+        json_response = response.json()
+        completion = json_response["responses"][0]
+
+        if stop is not None:
+            completion = enforce_stop_tokens(completion, stop)
+
+        return completion
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        raise NotImplementedError(
+            "Async call is not supported for OpenLLM at the moment."
+        )

api/core/third_party/langchain/llms/xinference_llm.py

@@ -0,0 +1,159 @@
+from typing import Optional, List, Any, Union, Generator
+
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms import Xinference
+from langchain.llms.utils import enforce_stop_tokens
+from xinference.client import (
+    RESTfulChatglmCppChatModelHandle,
+    RESTfulChatModelHandle,
+    RESTfulGenerateModelHandle,
+)
+
+
+class XinferenceLLM(Xinference):
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call the xinference model and return the output.
+
+        Args:
+            prompt: The prompt to use for generation.
+            stop: Optional list of stop words to use when generating.
+            generate_config: Optional dictionary for the configuration used for
+                generation.
+
+        Returns:
+            The generated string by the model.
+        """
+        model = self.client.get_model(self.model_uid)
+
+        if isinstance(model, RESTfulChatModelHandle):
+            generate_config: "LlamaCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
+
+            if stop:
+                generate_config["stop"] = stop
+
+            if generate_config and generate_config.get("stream"):
+                combined_text_output = ""
+                for token in self._stream_generate(
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
+                ):
+                    combined_text_output += token
+                return combined_text_output
+            else:
+                completion = model.chat(prompt=prompt, generate_config=generate_config)
+                return completion["choices"][0]["message"]["content"]
+        elif isinstance(model, RESTfulGenerateModelHandle):
+            generate_config: "LlamaCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
+
+            if stop:
+                generate_config["stop"] = stop
+
+            if generate_config and generate_config.get("stream"):
+                combined_text_output = ""
+                for token in self._stream_generate(
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
+                ):
+                    combined_text_output += token
+                return combined_text_output
+
+            else:
+                completion = model.generate(
+                    prompt=prompt, generate_config=generate_config
+                )
+                return completion["choices"][0]["text"]
+        elif isinstance(model, RESTfulChatglmCppChatModelHandle):
+            generate_config: "ChatglmCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
+
+            if generate_config and generate_config.get("stream"):
+                combined_text_output = ""
+                for token in self._stream_generate(
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
+                ):
+                    combined_text_output += token
+                completion = combined_text_output
+            else:
+                completion = model.chat(prompt=prompt, generate_config=generate_config)
+                completion = completion["choices"][0]["message"]["content"]
+
+            if stop is not None:
+                completion = enforce_stop_tokens(completion, stop)
+
+            return completion
+
+    def _stream_generate(
+        self,
+        model: Union[
+            "RESTfulGenerateModelHandle",
+            "RESTfulChatModelHandle",
+            "RESTfulChatglmCppChatModelHandle",
+        ],
+        prompt: str,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        generate_config: Optional[
+            Union[
+                "LlamaCppGenerateConfig",
+                "PytorchGenerateConfig",
+                "ChatglmCppGenerateConfig",
+            ]
+        ] = None,
+    ) -> Generator[str, None, None]:
+        """
+        Args:
+            prompt: The prompt to use for generation.
+            model: The model used for generation.
+            stop: Optional list of stop words to use when generating.
+            generate_config: Optional dictionary for the configuration used for
+                generation.
+
+        Yields:
+            A string token.
+        """
+        if isinstance(
+            model, (RESTfulChatModelHandle, RESTfulChatglmCppChatModelHandle)
+        ):
+            streaming_response = model.chat(
+                prompt=prompt, generate_config=generate_config
+            )
+        else:
+            streaming_response = model.generate(
+                prompt=prompt, generate_config=generate_config
+            )
+
+        for chunk in streaming_response:
+            if isinstance(chunk, dict):
+                choices = chunk.get("choices", [])
+                if choices:
+                    choice = choices[0]
+                    if isinstance(choice, dict):
+                        if "text" in choice:
+                            token = choice.get("text", "")
+                        elif "delta" in choice and "content" in choice["delta"]:
+                            token = choice.get("delta").get("content")
+                        else:
+                            continue
+                        log_probs = choice.get("logprobs")
+                        if run_manager:
+                            run_manager.on_llm_new_token(
+                                token=token, verbose=self.verbose, log_probs=log_probs
+                            )
+                        yield token

api/core/tool/dataset_retriever_tool.py

@@ -1,4 +1,3 @@
-import re
 from typing import Type
 
 from flask import current_app
@@ -16,7 +15,6 @@ from models.dataset import Dataset, DocumentSegment
 
 
 class DatasetRetrieverToolInput(BaseModel):
-    dataset_id: str = Field(..., description="ID of dataset to be queried. MUST be UUID format.")
     query: str = Field(..., description="Query for the dataset to be used to retrieve the dataset.")
 
 
@@ -37,27 +35,22 @@ class DatasetRetrieverTool(BaseTool):
             description = 'useful for when you want to answer queries about the ' + dataset.name
 
         description = description.replace('\n', '').replace('\r', '')
-        description += '\nID of dataset MUST be ' + dataset.id
         return cls(
+            name=f'dataset-{dataset.id}',
             tenant_id=dataset.tenant_id,
             dataset_id=dataset.id,
             description=description,
             **kwargs
         )
 
-    def _run(self, dataset_id: str, query: str) -> str:
-        pattern = r'\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b'
-        match = re.search(pattern, dataset_id, re.IGNORECASE)
-        if match:
-            dataset_id = match.group()
-
+    def _run(self, query: str) -> str:
         dataset = db.session.query(Dataset).filter(
             Dataset.tenant_id == self.tenant_id,
-            Dataset.id == dataset_id
+            Dataset.id == self.dataset_id
         ).first()
 
         if not dataset:
-            return f'[{self.name} failed to find dataset with id {dataset_id}.]'
+            return f'[{self.name} failed to find dataset with id {self.dataset_id}.]'
 
         if dataset.indexing_technique == "economy":
             # use keyword table query

api/core/tool/web_reader_tool.py

@@ -88,6 +88,11 @@ class WebReaderTool(BaseTool):
             texts = character_splitter.split_text(page_contents)
             docs = [Document(page_content=t) for t in texts]
 
+            if len(docs) == 0:
+                return "No content found."
+
+            docs = docs[1:]
+
             # only use first 5 docs
             if len(docs) > 5:
                 docs = docs[:5]

api/core/vector_store/milvus_vector_store.py

@@ -0,0 +1,38 @@
+from langchain.vectorstores import  Milvus
+
+
+class MilvusVectorStore(Milvus):
+    def del_texts(self, where_filter: dict):
+        if not where_filter:
+            raise ValueError('where_filter must not be empty')
+
+        self._client.batch.delete_objects(
+            class_name=self._index_name,
+            where=where_filter,
+            output='minimal'
+        )
+
+    def del_text(self, uuid: str) -> None:
+        self._client.data_object.delete(
+            uuid,
+            class_name=self._index_name
+        )
+
+    def text_exists(self, uuid: str) -> bool:
+        result = self._client.query.get(self._index_name).with_additional(["id"]).with_where({
+            "path": ["doc_id"],
+            "operator": "Equal",
+            "valueText": uuid,
+        }).with_limit(1).do()
+
+        if "errors" in result:
+            raise ValueError(f"Error during query: {result['errors']}")
+
+        entries = result["data"]["Get"][self._index_name]
+        if len(entries) == 0:
+            return False
+
+        return True
+
+    def delete(self):
+        self._client.schema.delete_class(self._index_name)

api/core/vector_store/qdrant_vector_store.py

@@ -1,10 +1,11 @@
 from typing import cast, Any
 
 from langchain.schema import Document
-from langchain.vectorstores import Qdrant
 from qdrant_client.http.models import Filter, PointIdsList, FilterSelector
 from qdrant_client.local.qdrant_local import QdrantLocal
 
+from core.index.vector_index.qdrant import Qdrant
+
 
 class QdrantVectorStore(Qdrant):
     def del_texts(self, filter: Filter):

api/events/event_handlers/create_document_index.py

@@ -1,6 +1,5 @@
 from events.dataset_event import dataset_was_deleted
 from events.event_handlers.document_index_event import document_index_created
-from tasks.clean_dataset_task import clean_dataset_task
 import datetime
 import logging
 import time

api/migrations/versions/4bcffcd64aa4_update_dataset_model_field_null_.py

@@ -0,0 +1,46 @@
+"""update_dataset_model_field_null_available
+
+Revision ID: 4bcffcd64aa4
+Revises: 853f9b9cd3b6
+Create Date: 2023-08-28 20:58:50.077056
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '4bcffcd64aa4'
+down_revision = '853f9b9cd3b6'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('datasets', schema=None) as batch_op:
+        batch_op.alter_column('embedding_model',
+               existing_type=sa.VARCHAR(length=255),
+               nullable=True,
+               existing_server_default=sa.text("'text-embedding-ada-002'::character varying"))
+        batch_op.alter_column('embedding_model_provider',
+               existing_type=sa.VARCHAR(length=255),
+               nullable=True,
+               existing_server_default=sa.text("'openai'::character varying"))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('datasets', schema=None) as batch_op:
+        batch_op.alter_column('embedding_model_provider',
+               existing_type=sa.VARCHAR(length=255),
+               nullable=False,
+               existing_server_default=sa.text("'openai'::character varying"))
+        batch_op.alter_column('embedding_model',
+               existing_type=sa.VARCHAR(length=255),
+               nullable=False,
+               existing_server_default=sa.text("'text-embedding-ada-002'::character varying"))
+
+    # ### end Alembic commands ###

api/models/dataset.py

@@ -36,10 +36,8 @@ class Dataset(db.Model):
     updated_by = db.Column(UUID, nullable=True)
     updated_at = db.Column(db.DateTime, nullable=False,
                            server_default=db.text('CURRENT_TIMESTAMP(0)'))
-    embedding_model = db.Column(db.String(
-        255), nullable=False, server_default=db.text("'text-embedding-ada-002'::character varying"))
-    embedding_model_provider = db.Column(db.String(
-        255), nullable=False, server_default=db.text("'openai'::character varying"))
+    embedding_model = db.Column(db.String(255), nullable=True)
+    embedding_model_provider = db.Column(db.String(255), nullable=True)
 
     @property
     def dataset_keyword_table(self):

api/requirements.txt

@@ -49,5 +49,5 @@ huggingface_hub~=0.16.4
 transformers~=4.31.0
 stripe~=5.5.0
 pandas==1.5.3
-xinference==0.2.0
-openllm~=0.2.26
+xinference==0.2.1
+safetensors==0.3.2

api/services/dataset_service.py

@@ -10,6 +10,7 @@ from flask import current_app
 from sqlalchemy import func
 
 from core.index.index import IndexBuilder
+from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_providers.model_factory import ModelFactory
 from extensions.ext_redis import redis_client
 from flask_login import current_user
@@ -91,16 +92,18 @@ class DatasetService:
         if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
             raise DatasetNameDuplicateError(
                 f'Dataset with name {name} already exists.')
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=current_user.current_tenant_id
-        )
+        embedding_model = None
+        if indexing_technique == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=current_user.current_tenant_id
+            )
         dataset = Dataset(name=name, indexing_technique=indexing_technique)
         # dataset = Dataset(name=name, provider=provider, config=config)
         dataset.created_by = account.id
         dataset.updated_by = account.id
         dataset.tenant_id = tenant_id
-        dataset.embedding_model_provider = embedding_model.model_provider.provider_name
-        dataset.embedding_model = embedding_model.name
+        dataset.embedding_model_provider = embedding_model.model_provider.provider_name if embedding_model else None
+        dataset.embedding_model = embedding_model.name if embedding_model else None
         db.session.add(dataset)
         db.session.commit()
         return dataset
@@ -115,17 +118,50 @@ class DatasetService:
         else:
             return dataset
 
+    @staticmethod
+    def check_dataset_model_setting(dataset):
+        if dataset.indexing_technique == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=dataset.tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ValueError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ValueError(f"The dataset in unavailable, due to: "
+                                 f"{ex.description}")
+
     @staticmethod
     def update_dataset(dataset_id, data, user):
+        filtered_data = {k: v for k, v in data.items() if v is not None or k == 'description'}
         dataset = DatasetService.get_dataset(dataset_id)
         DatasetService.check_dataset_permission(dataset, user)
+        action = None
         if dataset.indexing_technique != data['indexing_technique']:
             # if update indexing_technique
             if data['indexing_technique'] == 'economy':
-                deal_dataset_vector_index_task.delay(dataset_id, 'remove')
+                action = 'remove'
+                filtered_data['embedding_model'] = None
+                filtered_data['embedding_model_provider'] = None
             elif data['indexing_technique'] == 'high_quality':
-                deal_dataset_vector_index_task.delay(dataset_id, 'add')
-        filtered_data = {k: v for k, v in data.items() if v is not None or k == 'description'}
+                action = 'add'
+                # get embedding model setting
+                try:
+                    embedding_model = ModelFactory.get_embedding_model(
+                        tenant_id=current_user.current_tenant_id
+                    )
+                    filtered_data['embedding_model'] = embedding_model.name
+                    filtered_data['embedding_model_provider'] = embedding_model.model_provider.provider_name
+                except LLMBadRequestError:
+                    raise ValueError(
+                        f"No Embedding Model available. Please configure a valid provider "
+                        f"in the Settings -> Model Provider.")
+                except ProviderTokenNotInitError as ex:
+                    raise ValueError(ex.description)
 
         filtered_data['updated_by'] = user.id
         filtered_data['updated_at'] = datetime.datetime.now()
@@ -133,7 +169,8 @@ class DatasetService:
         dataset.query.filter_by(id=dataset_id).update(filtered_data)
 
         db.session.commit()
-
+        if action:
+            deal_dataset_vector_index_task.delay(dataset_id, action)
         return dataset
 
     @staticmethod
@@ -284,8 +321,9 @@ class DocumentService:
             "github_link": str,
             "open_source_license": str,
             "commit_date": str,
-            "commit_author": str
-        }
+            "commit_author": str,
+        },
+        "others": dict
     }
 
     @staticmethod
@@ -393,16 +431,26 @@ class DocumentService:
     def save_document_with_dataset_id(dataset: Dataset, document_data: dict,
                                       account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
                                       created_from: str = 'web'):
+
         # check document limit
         if current_app.config['EDITION'] == 'CLOUD':
-            documents_count = DocumentService.get_tenant_documents_count()
-            tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
-            if documents_count > tenant_document_count:
-                raise ValueError(f"over document limit {tenant_document_count}.")
+            if 'original_document_id' not in document_data or not document_data['original_document_id']:
+                count = 0
+                if document_data["data_source"]["type"] == "upload_file":
+                    upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
+                    count = len(upload_file_list)
+                elif document_data["data_source"]["type"] == "notion_import":
+                    notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
+                    for notion_info in notion_info_list:
+                        count = count + len(notion_info['pages'])
+                documents_count = DocumentService.get_tenant_documents_count()
+                total_count = documents_count + count
+                tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
+                if total_count > tenant_document_count:
+                    raise ValueError(f"over document limit {tenant_document_count}.")
         # if dataset is empty, update dataset data_source_type
         if not dataset.data_source_type:
             dataset.data_source_type = document_data["data_source"]["type"]
-            db.session.commit()
 
         if not dataset.indexing_technique:
             if 'indexing_technique' not in document_data \
@@ -410,6 +458,13 @@ class DocumentService:
                 raise ValueError("Indexing technique is required")
 
             dataset.indexing_technique = document_data["indexing_technique"]
+            if document_data["indexing_technique"] == 'high_quality':
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=dataset.tenant_id
+                )
+                dataset.embedding_model = embedding_model.name
+                dataset.embedding_model_provider = embedding_model.model_provider.provider_name
+
 
         documents = []
         batch = time.strftime('%Y%m%d%H%M%S') + str(random.randint(100000, 999999))
@@ -454,12 +509,12 @@ class DocumentService:
                     data_source_info = {
                         "upload_file_id": file_id,
                     }
-                    document = DocumentService.save_document(dataset, dataset_process_rule.id,
-                                                             document_data["data_source"]["type"],
-                                                             document_data["doc_form"],
-                                                             document_data["doc_language"],
-                                                             data_source_info, created_from, position,
-                                                             account, file_name, batch)
+                    document = DocumentService.build_document(dataset, dataset_process_rule.id,
+                                                              document_data["data_source"]["type"],
+                                                              document_data["doc_form"],
+                                                              document_data["doc_language"],
+                                                              data_source_info, created_from, position,
+                                                              account, file_name, batch)
                     db.session.add(document)
                     db.session.flush()
                     document_ids.append(document.id)
@@ -500,12 +555,12 @@ class DocumentService:
                                 "notion_page_icon": page['page_icon'],
                                 "type": page['type']
                             }
-                            document = DocumentService.save_document(dataset, dataset_process_rule.id,
-                                                                     document_data["data_source"]["type"],
-                                                                     document_data["doc_form"],
-                                                                     document_data["doc_language"],
-                                                                     data_source_info, created_from, position,
-                                                                     account, page['page_name'], batch)
+                            document = DocumentService.build_document(dataset, dataset_process_rule.id,
+                                                                      document_data["data_source"]["type"],
+                                                                      document_data["doc_form"],
+                                                                      document_data["doc_language"],
+                                                                      data_source_info, created_from, position,
+                                                                      account, page['page_name'], batch)
                             db.session.add(document)
                             db.session.flush()
                             document_ids.append(document.id)
@@ -524,10 +579,10 @@ class DocumentService:
         return documents, batch
 
     @staticmethod
-    def save_document(dataset: Dataset, process_rule_id: str, data_source_type: str, document_form: str,
-                      document_language: str, data_source_info: dict, created_from: str, position: int,
-                      account: Account,
-                      name: str, batch: str):
+    def build_document(dataset: Dataset, process_rule_id: str, data_source_type: str, document_form: str,
+                       document_language: str, data_source_info: dict, created_from: str, position: int,
+                       account: Account,
+                       name: str, batch: str):
         document = Document(
             tenant_id=dataset.tenant_id,
             dataset_id=dataset.id,
@@ -556,6 +611,7 @@ class DocumentService:
     def update_document_with_dataset_id(dataset: Dataset, document_data: dict,
                                         account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
                                         created_from: str = 'web'):
+        DatasetService.check_dataset_model_setting(dataset)
         document = DocumentService.get_document(dataset.id, document_data["original_document_id"])
         if document.display_status != 'available':
             raise ValueError("Document is not available")
@@ -648,15 +704,26 @@ class DocumentService:
 
     @staticmethod
     def save_document_without_dataset_id(tenant_id: str, document_data: dict, account: Account):
+        count = 0
+        if document_data["data_source"]["type"] == "upload_file":
+            upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
+            count = len(upload_file_list)
+        elif document_data["data_source"]["type"] == "notion_import":
+            notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
+            for notion_info in notion_info_list:
+                count = count + len(notion_info['pages'])
         # check document limit
         if current_app.config['EDITION'] == 'CLOUD':
             documents_count = DocumentService.get_tenant_documents_count()
+            total_count = documents_count + count
             tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
-            if documents_count > tenant_document_count:
-                raise ValueError(f"over document limit {tenant_document_count}.")
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=tenant_id
-        )
+            if total_count > tenant_document_count:
+                raise ValueError(f"All your documents have overed limit {tenant_document_count}.")
+        embedding_model = None
+        if document_data['indexing_technique'] == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=tenant_id
+            )
         # save dataset
         dataset = Dataset(
             tenant_id=tenant_id,
@@ -664,8 +731,8 @@ class DocumentService:
             data_source_type=document_data["data_source"]["type"],
             indexing_technique=document_data["indexing_technique"],
             created_by=account.id,
-            embedding_model=embedding_model.name,
-            embedding_model_provider=embedding_model.model_provider.provider_name
+            embedding_model=embedding_model.name if embedding_model else None,
+            embedding_model_provider=embedding_model.model_provider.provider_name if embedding_model else None
         )
 
         db.session.add(dataset)
@@ -873,21 +940,25 @@ class SegmentService:
         if document.doc_form == 'qa_model':
             if 'answer' not in args or not args['answer']:
                 raise ValueError("Answer is required")
+            if not args['answer'].strip():
+                raise ValueError("Answer is empty")
+        if 'content' not in args or not args['content'] or not args['content'].strip():
+            raise ValueError("Content is empty")
 
     @classmethod
     def create_segment(cls, args: dict, document: Document, dataset: Dataset):
         content = args['content']
         doc_id = str(uuid.uuid4())
         segment_hash = helper.generate_text_hash(content)
-
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=dataset.tenant_id,
-            model_provider_name=dataset.embedding_model_provider,
-            model_name=dataset.embedding_model
-        )
-
-        # calc embedding use tokens
-        tokens = embedding_model.get_num_tokens(content)
+        tokens = 0
+        if dataset.indexing_technique == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=dataset.tenant_id,
+                model_provider_name=dataset.embedding_model_provider,
+                model_name=dataset.embedding_model
+            )
+            # calc embedding use tokens
+            tokens = embedding_model.get_num_tokens(content)
         max_position = db.session.query(func.max(DocumentSegment.position)).filter(
             DocumentSegment.document_id == document.id
         ).scalar()
@@ -949,15 +1020,16 @@ class SegmentService:
                     kw_index.update_segment_keywords_index(segment.index_node_id, segment.keywords)
             else:
                 segment_hash = helper.generate_text_hash(content)
+                tokens = 0
+                if dataset.indexing_technique == 'high_quality':
+                    embedding_model = ModelFactory.get_embedding_model(
+                        tenant_id=dataset.tenant_id,
+                        model_provider_name=dataset.embedding_model_provider,
+                        model_name=dataset.embedding_model
+                    )
 
-                embedding_model = ModelFactory.get_embedding_model(
-                    tenant_id=dataset.tenant_id,
-                    model_provider_name=dataset.embedding_model_provider,
-                    model_name=dataset.embedding_model
-                )
-
-                # calc embedding use tokens
-                tokens = embedding_model.get_num_tokens(content)
+                    # calc embedding use tokens
+                    tokens = embedding_model.get_num_tokens(content)
                 segment.content = content
                 segment.index_node_hash = segment_hash
                 segment.word_count = len(content)
@@ -972,7 +1044,7 @@ class SegmentService:
                 db.session.add(segment)
                 db.session.commit()
                 # update segment vector index
-                VectorService.create_segment_vector(args['keywords'], segment, dataset)
+                VectorService.update_segment_vector(args['keywords'], segment, dataset)
         except Exception as e:
             logging.exception("update segment index failed")
             segment.enabled = False
@@ -989,10 +1061,11 @@ class SegmentService:
         cache_result = redis_client.get(indexing_cache_key)
         if cache_result is not None:
             raise ValueError("Segment is deleting.")
-        # send delete segment index task
-        redis_client.setex(indexing_cache_key, 600, 1)
+
         # enabled segment need to delete index
         if segment.enabled:
+            # send delete segment index task
+            redis_client.setex(indexing_cache_key, 600, 1)
             delete_segment_from_index_task.delay(segment.id, segment.index_node_id, dataset.id, document.id)
         db.session.delete(segment)
         db.session.commit()

api/services/hit_testing_service.py

@@ -19,7 +19,7 @@ from models.dataset import Dataset, DocumentSegment, DatasetQuery
 class HitTestingService:
     @classmethod
     def retrieve(cls, dataset: Dataset, query: str, account: Account, limit: int = 10) -> dict:
-        if dataset.available_document_count == 0 or dataset.available_document_count == 0:
+        if dataset.available_document_count == 0 or dataset.available_segment_count == 0:
             return {
                 "query": {
                     "content": query,

api/tasks/batch_create_segment_to_index_task.py

@@ -49,18 +49,20 @@ def batch_create_segment_to_index_task(job_id: str, content: List, dataset_id: s
         if not dataset_document.enabled or dataset_document.archived or dataset_document.indexing_status != 'completed':
             raise ValueError('Document is not available.')
         document_segments = []
-        for segment in content:
-            content = segment['content']
-            doc_id = str(uuid.uuid4())
-            segment_hash = helper.generate_text_hash(content)
+        embedding_model = None
+        if dataset.indexing_technique == 'high_quality':
             embedding_model = ModelFactory.get_embedding_model(
                 tenant_id=dataset.tenant_id,
                 model_provider_name=dataset.embedding_model_provider,
                 model_name=dataset.embedding_model
             )
 
+        for segment in content:
+            content = segment['content']
+            doc_id = str(uuid.uuid4())
+            segment_hash = helper.generate_text_hash(content)
             # calc embedding use tokens
-            tokens = embedding_model.get_num_tokens(content)
+            tokens = embedding_model.get_num_tokens(content) if embedding_model else 0
             max_position = db.session.query(func.max(DocumentSegment.position)).filter(
                 DocumentSegment.document_id == dataset_document.id
             ).scalar()

api/tasks/clean_dataset_task.py

@@ -3,8 +3,10 @@ import time
 
 import click
 from celery import shared_task
+from flask import current_app
 
 from core.index.index import IndexBuilder
+from core.index.vector_index.vector_index import VectorIndex
 from extensions.ext_database import db
 from models.dataset import DocumentSegment, Dataset, DatasetKeywordTable, DatasetQuery, DatasetProcessRule, \
     AppDatasetJoin, Document
@@ -35,11 +37,11 @@ def clean_dataset_task(dataset_id: str, tenant_id: str, indexing_technique: str,
         documents = db.session.query(Document).filter(Document.dataset_id == dataset_id).all()
         segments = db.session.query(DocumentSegment).filter(DocumentSegment.dataset_id == dataset_id).all()
 
-        vector_index = IndexBuilder.get_index(dataset, 'high_quality')
         kw_index = IndexBuilder.get_index(dataset, 'economy')
 
         # delete from vector index
-        if vector_index:
+        if dataset.indexing_technique == 'high_quality':
+            vector_index = IndexBuilder.get_default_high_quality_index(dataset)
             try:
                 vector_index.delete()
             except Exception:

api/tasks/deal_dataset_vector_index_task.py

@@ -31,7 +31,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
             raise Exception('Dataset not found')
 
         if action == "remove":
-            index = IndexBuilder.get_index(dataset, 'high_quality', ignore_high_quality_check=True)
+            index = IndexBuilder.get_index(dataset, 'high_quality', ignore_high_quality_check=False)
             index.delete()
         elif action == "add":
             dataset_documents = db.session.query(DatasetDocument).filter(
@@ -43,7 +43,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
 
             if dataset_documents:
                 # save vector index
-                index = IndexBuilder.get_index(dataset, 'high_quality', ignore_high_quality_check=True)
+                index = IndexBuilder.get_index(dataset, 'high_quality', ignore_high_quality_check=False)
                 documents = []
                 for dataset_document in dataset_documents:
                     # delete from vector index
@@ -65,7 +65,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
                         documents.append(document)
 
                 # save vector index
-                index.add_texts(documents)
+                index.create(documents)
 
         end_at = time.perf_counter()
         logging.info(

api/tests/integration_tests/.env.example

@@ -39,4 +39,7 @@ XINFERENCE_SERVER_URL=
 XINFERENCE_MODEL_UID=
 
 # OpenLLM Credentials
-OPENLLM_SERVER_URL=
+OPENLLM_SERVER_URL=
+
+# LocalAI Credentials
+LOCALAI_SERVER_URL=

api/tests/integration_tests/models/embedding/test_localai_embedding.py

@@ -0,0 +1,61 @@
+import json
+import os
+from unittest.mock import patch, MagicMock
+
+from core.model_providers.models.embedding.localai_embedding import LocalAIEmbedding
+from core.model_providers.models.entity.model_params import ModelType
+from core.model_providers.providers.localai_provider import LocalAIProvider
+from models.provider import Provider, ProviderType, ProviderModel
+
+
+def get_mock_provider():
+    return Provider(
+        id='provider_id',
+        tenant_id='tenant_id',
+        provider_name='localai',
+        provider_type=ProviderType.CUSTOM.value,
+        encrypted_config='',
+        is_valid=True,
+    )
+
+
+def get_mock_embedding_model(mocker):
+    model_name = 'text-embedding-ada-002'
+    server_url = os.environ['LOCALAI_SERVER_URL']
+    model_provider = LocalAIProvider(provider=get_mock_provider())
+
+    mock_query = MagicMock()
+    mock_query.filter.return_value.first.return_value = ProviderModel(
+        provider_name='localai',
+        model_name=model_name,
+        model_type=ModelType.EMBEDDINGS.value,
+        encrypted_config=json.dumps({
+            'server_url': server_url,
+        }),
+        is_valid=True,
+    )
+    mocker.patch('extensions.ext_database.db.session.query', return_value=mock_query)
+
+    return LocalAIEmbedding(
+        model_provider=model_provider,
+        name=model_name
+    )
+
+
+def decrypt_side_effect(tenant_id, encrypted_api_key):
+    return encrypted_api_key
+
+
+@patch('core.helper.encrypter.decrypt_token', side_effect=decrypt_side_effect)
+def test_embed_documents(mock_decrypt, mocker):
+    embedding_model = get_mock_embedding_model(mocker)
+    rst = embedding_model.client.embed_documents(['test', 'test1'])
+    assert isinstance(rst, list)
+    assert len(rst) == 2
+
+
+@patch('core.helper.encrypter.decrypt_token', side_effect=decrypt_side_effect)
+def test_embed_query(mock_decrypt, mocker):
+    embedding_model = get_mock_embedding_model(mocker)
+    rst = embedding_model.client.embed_query('test')
+    assert isinstance(rst, list)

api/tests/integration_tests/models/llm/test_localai_model.py

@@ -0,0 +1,68 @@
+import json
+import os
+from unittest.mock import patch, MagicMock
+
+from core.model_providers.models.llm.localai_model import LocalAIModel
+from core.model_providers.providers.localai_provider import LocalAIProvider
+from core.model_providers.models.entity.message import PromptMessage
+from core.model_providers.models.entity.model_params import ModelKwargs, ModelType
+from models.provider import Provider, ProviderType, ProviderModel
+
+
+def get_mock_provider(server_url):
+    return Provider(
+        id='provider_id',
+        tenant_id='tenant_id',
+        provider_name='localai',
+        provider_type=ProviderType.CUSTOM.value,
+        encrypted_config=json.dumps({}),
+        is_valid=True,
+    )
+
+
+def get_mock_model(model_name, mocker):
+    model_kwargs = ModelKwargs(
+        max_tokens=10,
+        temperature=0
+    )
+    server_url = os.environ['LOCALAI_SERVER_URL']
+
+    mock_query = MagicMock()
+    mock_query.filter.return_value.first.return_value = ProviderModel(
+        provider_name='localai',
+        model_name=model_name,
+        model_type=ModelType.TEXT_GENERATION.value,
+        encrypted_config=json.dumps({'server_url': server_url, 'completion_type': 'completion'}),
+        is_valid=True,
+    )
+    mocker.patch('extensions.ext_database.db.session.query', return_value=mock_query)
+
+    openai_provider = LocalAIProvider(provider=get_mock_provider(server_url))
+    return LocalAIModel(
+        model_provider=openai_provider,
+        name=model_name,
+        model_kwargs=model_kwargs
+    )
+
+
+def decrypt_side_effect(tenant_id, encrypted_openai_api_key):
+    return encrypted_openai_api_key
+
+
+@patch('core.helper.encrypter.decrypt_token', side_effect=decrypt_side_effect)
+def test_get_num_tokens(mock_decrypt, mocker):
+    openai_model = get_mock_model('ggml-gpt4all-j', mocker)
+    rst = openai_model.get_num_tokens([PromptMessage(content='you are a kindness Assistant.')])
+    assert rst > 0
+
+
+@patch('core.helper.encrypter.decrypt_token', side_effect=decrypt_side_effect)
+def test_run(mock_decrypt, mocker):
+    mocker.patch('core.model_providers.providers.base.BaseModelProvider.update_last_used', return_value=None)
+
+    openai_model = get_mock_model('ggml-gpt4all-j', mocker)
+    rst = openai_model.run(
+        [PromptMessage(content='Human: Are you Human? you MUST only answer `y` or `n`? \nAssistant: ')],
+        stop=['\nHuman:'],
+    )
+    assert len(rst.content) > 0

api/tests/unit_tests/model_providers/test_huggingface_hub_provider.py

@@ -17,7 +17,8 @@ HOSTED_INFERENCE_API_VALIDATE_CREDENTIAL = {
 INFERENCE_ENDPOINTS_VALIDATE_CREDENTIAL = {
     'huggingfacehub_api_type': 'inference_endpoints',
     'huggingfacehub_api_token': 'valid_key',
-    'huggingfacehub_endpoint_url': 'valid_url'
+    'huggingfacehub_endpoint_url': 'valid_url',
+    'task_type': 'text-generation'
 }
 
 def encrypt_side_effect(tenant_id, encrypt_key):
@@ -62,7 +63,7 @@ def test_hosted_inference_api_is_credentials_valid_or_raise_invalid(mock_model_i
 
 def test_inference_endpoints_is_credentials_valid_or_raise_valid(mocker):
     mocker.patch('huggingface_hub.hf_api.HfApi.whoami', return_value=None)
-    mocker.patch('langchain.llms.huggingface_endpoint.HuggingFaceEndpoint._call', return_value="abc")
+    mocker.patch('core.third_party.langchain.llms.huggingface_endpoint_llm.HuggingFaceEndpointLLM._call', return_value="abc")
 
     MODEL_PROVIDER_CLASS.is_model_credentials_valid_or_raise(
         model_name='test_model_name',
@@ -70,8 +71,10 @@ def test_inference_endpoints_is_credentials_valid_or_raise_valid(mocker):
         credentials=INFERENCE_ENDPOINTS_VALIDATE_CREDENTIAL
     )
 
+
 def test_inference_endpoints_is_credentials_valid_or_raise_invalid(mocker):
     mocker.patch('huggingface_hub.hf_api.HfApi.whoami', return_value=None)
+    mocker.patch('core.third_party.langchain.llms.huggingface_endpoint_llm.HuggingFaceEndpointLLM._call', return_value="abc")
 
     with pytest.raises(CredentialsValidateFailedError):
         MODEL_PROVIDER_CLASS.is_model_credentials_valid_or_raise(