feat: bump version to 0.3.8 (#559)

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

.github/workflows/stale.yml

@@ -20,8 +20,11 @@ jobs:
     steps:
     - uses: actions/stale@v5
       with:
+        days-before-issue-stale: 30
+        days-before-issue-close: 3
         repo-token: ${{ secrets.GITHUB_TOKEN }}
         stale-issue-message: "Close due to it's no longer active, if you have any questions, you can reopen it."
         stale-pr-message: "Close due to it's no longer active, if you have any questions, you can reopen it."
         stale-issue-label: 'no-issue-activity'
         stale-pr-label: 'no-pr-activity'
+        any-of-labels: 'duplicate,question,invalid,wontfix,no-issue-activity,no-pr-activity,enhancement'

.gitignore

@@ -109,6 +109,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+.conda/
 
 # Spyder project settings
 .spyderproject
@@ -147,3 +148,5 @@ docker/volumes/weaviate/*
 sdks/python-client/build
 sdks/python-client/dist
 sdks/python-client/dify_client.egg-info
+
+.vscode/

CONTRIBUTING.md

@@ -54,3 +54,8 @@ 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!
+
+### 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.

CONTRIBUTING_CN.md

@@ -51,3 +51,7 @@ git clone git@github.com:<github_username>/dify.git
 ## 社区渠道
 
 遇到困难了吗？有任何问题吗? 加入 [Discord Community Server](https://discord.gg/AhzKf7dNgk)，我们将为您提供帮助。
+
+### 多语言支持
+
+需要参与贡献翻译内容，请参阅[前端多语言翻译 README](web/i18n/README_CN.md)。

README.md

@@ -45,6 +45,11 @@ docker compose up -d
 
 After running, you can access the Dify dashboard in your browser at [http://localhost/install](http://localhost/install) and start the initialization installation process.
 
+### Helm Chart
+
+A big thanks to @BorisPolonsky for providing us with a [Helm Chart](https://helm.sh/) version, which allows Dify to be deployed on Kubernetes.
+You can go to https://github.com/BorisPolonsky/dify-helm for deployment information.
+
 ### Configuration
 
 If you need to customize the configuration, please refer to the comments in our [docker-compose.yml](docker/docker-compose.yaml) file and manually set the environment configuration. After making the changes, please run 'docker-compose up -d' again.

README_CN.md

@@ -47,6 +47,11 @@ docker compose up -d
 
 运行后，可以在浏览器上访问 [http://localhost/install](http://localhost/install) 进入 Dify 控制台并开始初始化安装操作。
 
+### Helm Chart
+
+非常感谢 @BorisPolonsky 为我们提供了一个 [Helm Chart](https://helm.sh/) 版本，可以在 Kubernetes 上部署 Dify。
+您可以前往 https://github.com/BorisPolonsky/dify-helm 来获取部署信息。
+
 ### 配置
 
 需要自定义配置，请参考我们的 [docker-compose.yml](docker/docker-compose.yaml) 文件中的注释，并手动设置环境配置，修改完毕后，请再次执行 `docker-compose up -d`。

README_ES.md

@@ -46,6 +46,11 @@ docker compose up -d
 
 Después de ejecutarlo, puedes acceder al panel de control de Dify en tu navegador desde [http://localhost/install](http://localhost/install) y comenzar el proceso de instalación de inicialización.
 
+### Helm Chart
+
+Un gran agradecimiento a @BorisPolonsky por proporcionarnos una versión de [Helm Chart](https://helm.sh/), que permite desplegar Dify en Kubernetes.
+Puede ir a https://github.com/BorisPolonsky/dify-helm para obtener información de despliegue.
+
 ### Configuración
 
 Si necesitas personalizar la configuración, consulta los comentarios en nuestro archivo [docker-compose.yml](docker/docker-compose.yaml) y configura manualmente la configuración del entorno. Después de realizar los cambios, ejecuta nuevamente 'docker-compose up -d'.

README_JA.md

@@ -46,6 +46,11 @@ docker compose up -d
 
 実行後、ブラウザで [http://localhost/install](http://localhost/install) にアクセスし、初期化インストール作業を開始することができます。
 
+### Helm Chart
+
+@BorisPolonsky に大感謝します。彼は Dify を Kubernetes 上にデプロイするための [Helm Chart](https://helm.sh/) バージョンを提供してくれました。
+デプロイ情報については、https://github.com/BorisPolonsky/dify-helm をご覧ください。
+
 ### 構成
 
 カスタマイズが必要な場合は、[docker-compose.yml](docker/docker-compose.yaml) ファイルのコメントを参照し、手動で環境設定をお願いします。変更後、再度 'docker-compose up -d' を実行してください。

api/.env.example

@@ -8,13 +8,19 @@ EDITION=SELF_HOSTED
 SECRET_KEY=
 
 # Console API base URL
-CONSOLE_URL=http://127.0.0.1:5001
+CONSOLE_API_URL=http://127.0.0.1:5001
+
+# Console frontend web base URL
+CONSOLE_WEB_URL=http://127.0.0.1:3000
 
 # Service API base URL
-API_URL=http://127.0.0.1:5001
+SERVICE_API_URL=http://127.0.0.1:5001
 
-# Web APP base URL
-APP_URL=http://127.0.0.1:3000
+# Web APP API base URL
+APP_API_URL=http://127.0.0.1:5001
+
+# Web APP frontend web base URL
+APP_WEB_URL=http://127.0.0.1:3000
 
 # celery configuration
 CELERY_BROKER_URL=redis://:difyai123456@localhost:6379/1
@@ -22,6 +28,7 @@ CELERY_BROKER_URL=redis://:difyai123456@localhost:6379/1
 # redis configuration
 REDIS_HOST=localhost
 REDIS_PORT=6379
+REDIS_USERNAME=
 REDIS_PASSWORD=difyai123456
 REDIS_DB=0
 
@@ -78,9 +85,20 @@ WEAVIATE_BATCH_SIZE=100
 QDRANT_URL=path:storage/qdrant
 QDRANT_API_KEY=your-qdrant-api-key
 
+# Mail configuration, support: resend
+MAIL_TYPE=
+MAIL_DEFAULT_SEND_FROM=no-reply <no-reply@dify.ai>
+RESEND_API_KEY=
+
 # Sentry configuration
 SENTRY_DSN=
 
 # DEBUG
 DEBUG=false
 SQLALCHEMY_ECHO=false
+
+# Notion import configuration, support public and internal
+NOTION_INTEGRATION_TYPE=public
+NOTION_CLIENT_SECRET=you-client-secret
+NOTION_CLIENT_ID=you-client-id
+NOTION_INTERNAL_SECRET=you-internal-secret

api/Dockerfile

@@ -5,9 +5,11 @@ LABEL maintainer="takatost@gmail.com"
 ENV FLASK_APP app.py
 ENV EDITION SELF_HOSTED
 ENV DEPLOY_ENV PRODUCTION
-ENV CONSOLE_URL http://127.0.0.1:5001
-ENV API_URL http://127.0.0.1:5001
-ENV APP_URL http://127.0.0.1:5001
+ENV CONSOLE_API_URL http://127.0.0.1:5001
+ENV CONSOLE_WEB_URL http://127.0.0.1:3000
+ENV SERVICE_API_URL http://127.0.0.1:5001
+ENV APP_API_URL http://127.0.0.1:5001
+ENV APP_WEB_URL http://127.0.0.1:3000
 
 EXPOSE 5001
 

api/README.md

@@ -17,6 +17,11 @@
    ```bash
    openssl rand -base64 42
    ```
+3.5 If you use annaconda, create a new environment and activate it
+   ```bash
+   conda create --name dify python=3.10
+   conda activate dify
+   ```
 4. Install dependencies
    ```bash
    pip install -r requirements.txt

api/app.py

@@ -1,5 +1,7 @@
 # -*- coding:utf-8 -*-
 import os
+from datetime import datetime
+
 if not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true':
     from gevent import monkey
     monkey.patch_all()
@@ -12,13 +14,13 @@ from flask import Flask, request, Response, session
 import flask_login
 from flask_cors import CORS
 
-from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_vector_store, ext_migrate, \
-    ext_database, ext_storage
+from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
+    ext_database, ext_storage, ext_mail
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
 
 # DO NOT REMOVE BELOW
-from models import model, account, dataset, web, task
+from models import model, account, dataset, web, task, source
 from events import event_handlers
 # DO NOT REMOVE ABOVE
 
@@ -77,11 +79,11 @@ def initialize_extensions(app):
     ext_database.init_app(app)
     ext_migrate.init(app, db)
     ext_redis.init_app(app)
-    ext_vector_store.init_app(app)
     ext_storage.init_app(app)
     ext_celery.init_app(app)
     ext_session.init_app(app)
     ext_login.init_app(app)
+    ext_mail.init_app(app)
     ext_sentry.init_app(app)
 
 
@@ -122,6 +124,9 @@ def load_user(user_id):
                     account.current_tenant_id = tenant_account_join.tenant_id
                     session['workspace_id'] = account.current_tenant_id
 
+            account.last_active_at = datetime.utcnow()
+            db.session.commit()
+
             # Log in the user with the updated user_id
             flask_login.login_user(account, remember=True)
 
@@ -145,13 +150,17 @@ def register_blueprints(app):
     from controllers.web import bp as web_bp
     from controllers.console import bp as console_app_bp
 
+    CORS(service_api_bp,
+         allow_headers=['Content-Type', 'Authorization', 'X-App-Code'],
+         methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH']
+         )
     app.register_blueprint(service_api_bp)
 
     CORS(web_bp,
          resources={
              r"/*": {"origins": app.config['WEB_API_CORS_ALLOW_ORIGINS']}},
          supports_credentials=True,
-         allow_headers=['Content-Type', 'Authorization'],
+         allow_headers=['Content-Type', 'Authorization', 'X-App-Code'],
          methods=['GET', 'PUT', 'POST', 'DELETE', 'OPTIONS', 'PATCH'],
          expose_headers=['X-Version', 'X-Env']
          )

api/commands.py

@@ -1,15 +1,19 @@
 import datetime
+import logging
 import random
 import string
 
 import click
 from flask import current_app
+from werkzeug.exceptions import NotFound
 
+from core.index.index import IndexBuilder
 from libs.password import password_pattern, valid_password, hash_password
 from libs.helper import email as email_validate
 from extensions.ext_database import db
 from libs.rsa import generate_key_pair
 from models.account import InvitationCode, Tenant
+from models.dataset import Dataset
 from models.model import Account
 import secrets
 import base64
@@ -159,8 +163,39 @@ def generate_upper_string():
     return result
 
 
+@click.command('recreate-all-dataset-indexes', help='Recreate all dataset indexes.')
+def recreate_all_dataset_indexes():
+    click.echo(click.style('Start recreate all dataset indexes.', fg='green'))
+    recreate_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:
+            try:
+                click.echo('Recreating dataset index: {}'.format(dataset.id))
+                index = IndexBuilder.get_index(dataset, 'high_quality')
+                if index and index._is_origin():
+                    index.recreate_dataset(dataset)
+                    recreate_count += 1
+                else:
+                    click.echo('passed.')
+            except Exception as e:
+                click.echo(click.style('Recreate dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                continue
+
+    click.echo(click.style('Congratulations! Recreate {} dataset indexes.'.format(recreate_count), fg='green'))
+
+
 def register_commands(app):
     app.cli.add_command(reset_password)
     app.cli.add_command(reset_email)
     app.cli.add_command(generate_invitation_codes)
     app.cli.add_command(reset_encrypt_key_pair)
+    app.cli.add_command(recreate_all_dataset_indexes)

api/config.py

@@ -28,9 +28,11 @@ DEFAULTS = {
     'SESSION_REDIS_USE_SSL': 'False',
     'OAUTH_REDIRECT_PATH': '/console/api/oauth/authorize',
     'OAUTH_REDIRECT_INDEX_PATH': '/',
-    'CONSOLE_URL': 'https://cloud.dify.ai',
-    'API_URL': 'https://api.dify.ai',
-    'APP_URL': 'https://udify.app',
+    'CONSOLE_WEB_URL': 'https://cloud.dify.ai',
+    'CONSOLE_API_URL': 'https://cloud.dify.ai',
+    'SERVICE_API_URL': 'https://api.dify.ai',
+    'APP_WEB_URL': 'https://udify.app',
+    'APP_API_URL': 'https://udify.app',
     'STORAGE_TYPE': 'local',
     'STORAGE_LOCAL_PATH': 'storage',
     'CHECK_UPDATE_URL': 'https://updates.dify.ai',
@@ -76,10 +78,15 @@ class Config:
 
     def __init__(self):
         # app settings
+        self.CONSOLE_API_URL = get_env('CONSOLE_URL') if get_env('CONSOLE_URL') else get_env('CONSOLE_API_URL')
+        self.CONSOLE_WEB_URL = get_env('CONSOLE_URL') if get_env('CONSOLE_URL') else get_env('CONSOLE_WEB_URL')
+        self.SERVICE_API_URL = get_env('API_URL') if get_env('API_URL') else get_env('SERVICE_API_URL')
+        self.APP_WEB_URL = get_env('APP_URL') if get_env('APP_URL') else get_env('APP_WEB_URL')
+        self.APP_API_URL = get_env('APP_URL') if get_env('APP_URL') else get_env('APP_API_URL')
         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.3"
+        self.CURRENT_VERSION = "0.3.8"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = "SELF_HOSTED"
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')
@@ -147,10 +154,15 @@ class Config:
 
         # cors settings
         self.CONSOLE_CORS_ALLOW_ORIGINS = get_cors_allow_origins(
-            'CONSOLE_CORS_ALLOW_ORIGINS', self.CONSOLE_URL)
+            'CONSOLE_CORS_ALLOW_ORIGINS', self.CONSOLE_WEB_URL)
         self.WEB_API_CORS_ALLOW_ORIGINS = get_cors_allow_origins(
             'WEB_API_CORS_ALLOW_ORIGINS', '*')
 
+        # mail settings
+        self.MAIL_TYPE = get_env('MAIL_TYPE')
+        self.MAIL_DEFAULT_SEND_FROM = get_env('MAIL_DEFAULT_SEND_FROM')
+        self.RESEND_API_KEY = get_env('RESEND_API_KEY')
+
         # sentry settings
         self.SENTRY_DSN = get_env('SENTRY_DSN')
         self.SENTRY_TRACES_SAMPLE_RATE = float(get_env('SENTRY_TRACES_SAMPLE_RATE'))
@@ -188,6 +200,14 @@ class Config:
         # set default LLM provider, default is 'openai', support `azure_openai`
         self.DEFAULT_LLM_PROVIDER = get_env('DEFAULT_LLM_PROVIDER')
 
+        # notion import setting
+        self.NOTION_CLIENT_ID = get_env('NOTION_CLIENT_ID')
+        self.NOTION_CLIENT_SECRET = get_env('NOTION_CLIENT_SECRET')
+        self.NOTION_INTEGRATION_TYPE = get_env('NOTION_INTEGRATION_TYPE')
+        self.NOTION_INTERNAL_SECRET = get_env('NOTION_INTERNAL_SECRET')
+        self.NOTION_INTEGRATION_TOKEN = get_env('NOTION_INTEGRATION_TOKEN')
+
+
 class CloudEditionConfig(Config):
 
     def __init__(self):

api/controllers/console/__init__.py

@@ -9,16 +9,16 @@ api = ExternalApi(bp)
 from . import setup, version, apikey, admin
 
 # Import app controllers
-from .app import app, site, completion, model_config, statistic, conversation, message, generator
+from .app import app, site, completion, model_config, statistic, conversation, message, generator, audio
 
 # Import auth controllers
-from .auth import login, oauth
+from .auth import login, oauth, data_source_oauth, activate
 
 # Import datasets controllers
-from .datasets import datasets, datasets_document, datasets_segments, file, hit_testing
+from .datasets import datasets, datasets_document, datasets_segments, file, hit_testing, data_source
 
 # Import workspace controllers
 from .workspace import workspace, members, providers, account
 
 # Import explore controllers
-from .explore import installed_app, recommended_app, completion, conversation, message, parameter, saved_message
+from .explore import installed_app, recommended_app, completion, conversation, message, parameter, saved_message, audio

api/controllers/console/app/app.py

@@ -22,6 +22,7 @@ model_config_fields = {
     'opening_statement': fields.String,
     'suggested_questions': fields.Raw(attribute='suggested_questions_list'),
     'suggested_questions_after_answer': fields.Raw(attribute='suggested_questions_after_answer_dict'),
+    'speech_to_text': fields.Raw(attribute='speech_to_text_dict'),
     'more_like_this': fields.Raw(attribute='more_like_this_dict'),
     'model': fields.Raw(attribute='model_dict'),
     'user_input_form': fields.Raw(attribute='user_input_form_list'),
@@ -144,6 +145,7 @@ class AppListApi(Resource):
                 opening_statement=model_configuration['opening_statement'],
                 suggested_questions=json.dumps(model_configuration['suggested_questions']),
                 suggested_questions_after_answer=json.dumps(model_configuration['suggested_questions_after_answer']),
+                speech_to_text=json.dumps(model_configuration['speech_to_text']),
                 more_like_this=json.dumps(model_configuration['more_like_this']),
                 model=json.dumps(model_configuration['model']),
                 user_input_form=json.dumps(model_configuration['user_input_form']),
@@ -434,6 +436,7 @@ class AppCopy(Resource):
             opening_statement=app_config.opening_statement,
             suggested_questions=app_config.suggested_questions,
             suggested_questions_after_answer=app_config.suggested_questions_after_answer,
+            speech_to_text=app_config.speech_to_text,
             more_like_this=app_config.more_like_this,
             model=app_config.model,
             user_input_form=app_config.user_input_form,

api/controllers/console/app/audio.py

@@ -0,0 +1,69 @@
+# -*- coding:utf-8 -*-
+import logging
+
+from flask import request
+from flask_login import login_required
+from werkzeug.exceptions import InternalServerError, NotFound
+
+import services
+from controllers.console import api
+from controllers.console.app import _get_app
+from controllers.console.app.error import AppUnavailableError, \
+    ProviderNotInitializeError, CompletionRequestError, ProviderQuotaExceededError, \
+    ProviderModelCurrentlyNotSupportError, NoAudioUploadedError, AudioTooLargeError, \
+    UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+    LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from flask_restful import Resource
+from services.audio_service import AudioService
+from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \
+    UnsupportedAudioTypeServiceError, ProviderNotSupportSpeechToTextServiceError
+
+
+class ChatMessageAudioApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, app_id):
+        app_id = str(app_id)
+        app_model = _get_app(app_id, 'chat')
+
+        file = request.files['file']
+
+        try:
+            response = AudioService.transcript(
+                tenant_id=app_model.tenant_id,
+                file=file,
+            )
+
+            return response
+        except services.errors.app_model_config.AppModelConfigBrokenError:
+            logging.exception("App model config broken.")
+            raise AppUnavailableError()
+        except NoAudioUploadedServiceError:
+            raise NoAudioUploadedError()
+        except AudioTooLargeServiceError as e:
+            raise AudioTooLargeError(str(e))
+        except UnsupportedAudioTypeServiceError:
+            raise UnsupportedAudioTypeError()
+        except ProviderNotSupportSpeechToTextServiceError:
+            raise ProviderNotSupportSpeechToTextError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except ValueError as e:
+            raise e
+        except Exception as e:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+        
+
+api.add_resource(ChatMessageAudioApi, '/apps/<uuid:app_id>/audio-to-text')

api/controllers/console/app/conversation.py

@@ -209,6 +209,26 @@ class CompletionConversationDetailApi(Resource):
         conversation_id = str(conversation_id)
 
         return _get_conversation(app_id, conversation_id, 'completion')
+    
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, app_id, conversation_id):
+        app_id = str(app_id)
+        conversation_id = str(conversation_id)
+
+        app = _get_app(app_id, 'chat')
+
+        conversation = db.session.query(Conversation) \
+            .filter(Conversation.id == conversation_id, Conversation.app_id == app.id).first()
+
+        if not conversation:
+            raise NotFound("Conversation Not Exists.")
+
+        conversation.is_deleted = True
+        db.session.commit()
+
+        return {'result': 'success'}, 204
 
 
 class ChatConversationApi(Resource):
@@ -356,6 +376,27 @@ class ChatConversationDetailApi(Resource):
         conversation_id = str(conversation_id)
 
         return _get_conversation(app_id, conversation_id, 'chat')
+    
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, app_id, conversation_id):
+        app_id = str(app_id)
+        conversation_id = str(conversation_id)
+
+        # get app info
+        app = _get_app(app_id, 'chat')
+
+        conversation = db.session.query(Conversation) \
+            .filter(Conversation.id == conversation_id, Conversation.app_id == app.id).first()
+
+        if not conversation:
+            raise NotFound("Conversation Not Exists.")
+
+        conversation.is_deleted = True
+        db.session.commit()
+
+        return {'result': 'success'}, 204
 
 
 

api/controllers/console/app/error.py

@@ -49,3 +49,27 @@ class AppMoreLikeThisDisabledError(BaseHTTPException):
     error_code = 'app_more_like_this_disabled'
     description = "The 'More like this' feature is disabled. Please refresh your page."
     code = 403
+
+
+class NoAudioUploadedError(BaseHTTPException):
+    error_code = 'no_audio_uploaded'
+    description = "Please upload your audio."
+    code = 400
+
+
+class AudioTooLargeError(BaseHTTPException):
+    error_code = 'audio_too_large'
+    description = "Audio size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedAudioTypeError(BaseHTTPException):
+    error_code = 'unsupported_audio_type'
+    description = "Audio type not allowed."
+    code = 415
+
+
+class ProviderNotSupportSpeechToTextError(BaseHTTPException):
+    error_code = 'provider_not_support_speech_to_text'
+    description = "Provider not support speech to text."
+    code = 400

api/controllers/console/app/model_config.py

@@ -41,6 +41,7 @@ class ModelConfigResource(Resource):
             opening_statement=model_configuration['opening_statement'],
             suggested_questions=json.dumps(model_configuration['suggested_questions']),
             suggested_questions_after_answer=json.dumps(model_configuration['suggested_questions_after_answer']),
+            speech_to_text=json.dumps(model_configuration['speech_to_text']),
             more_like_this=json.dumps(model_configuration['more_like_this']),
             model=json.dumps(model_configuration['model']),
             user_input_form=json.dumps(model_configuration['user_input_form']),

api/controllers/console/auth/activate.py

@@ -0,0 +1,75 @@
+import base64
+import secrets
+from datetime import datetime
+
+from flask_restful import Resource, reqparse
+
+from controllers.console import api
+from controllers.console.error import AlreadyActivateError
+from extensions.ext_database import db
+from libs.helper import email, str_len, supported_language, timezone
+from libs.password import valid_password, hash_password
+from models.account import AccountStatus, Tenant
+from services.account_service import RegisterService
+
+
+class ActivateCheckApi(Resource):
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('workspace_id', type=str, required=True, nullable=False, location='args')
+        parser.add_argument('email', type=email, required=True, nullable=False, location='args')
+        parser.add_argument('token', type=str, required=True, nullable=False, location='args')
+        args = parser.parse_args()
+
+        account = RegisterService.get_account_if_token_valid(args['workspace_id'], args['email'], args['token'])
+
+        tenant = db.session.query(Tenant).filter(
+            Tenant.id == args['workspace_id'],
+            Tenant.status == 'normal'
+        ).first()
+
+        return {'is_valid': account is not None, 'workspace_name': tenant.name}
+
+
+class ActivateApi(Resource):
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('workspace_id', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('email', type=email, required=True, nullable=False, location='json')
+        parser.add_argument('token', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('name', type=str_len(30), required=True, nullable=False, location='json')
+        parser.add_argument('password', type=valid_password, required=True, nullable=False, location='json')
+        parser.add_argument('interface_language', type=supported_language, required=True, nullable=False,
+                            location='json')
+        parser.add_argument('timezone', type=timezone, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        account = RegisterService.get_account_if_token_valid(args['workspace_id'], args['email'], args['token'])
+        if account is None:
+            raise AlreadyActivateError()
+
+        RegisterService.revoke_token(args['workspace_id'], args['email'], args['token'])
+
+        account.name = args['name']
+
+        # generate password salt
+        salt = secrets.token_bytes(16)
+        base64_salt = base64.b64encode(salt).decode()
+
+        # encrypt password with salt
+        password_hashed = hash_password(args['password'], salt)
+        base64_password_hashed = base64.b64encode(password_hashed).decode()
+        account.password = base64_password_hashed
+        account.password_salt = base64_salt
+        account.interface_language = args['interface_language']
+        account.timezone = args['timezone']
+        account.interface_theme = 'light'
+        account.status = AccountStatus.ACTIVE.value
+        account.initialized_at = datetime.utcnow()
+        db.session.commit()
+
+        return {'result': 'success'}
+
+
+api.add_resource(ActivateCheckApi, '/activate/check')
+api.add_resource(ActivateApi, '/activate')

api/controllers/console/auth/data_source_oauth.py

@@ -0,0 +1,101 @@
+import logging
+from datetime import datetime
+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_restful import Resource
+from werkzeug.exceptions import Forbidden
+from libs.oauth_data_source import NotionOAuth
+from controllers.console import api
+from ..setup import setup_required
+from ..wraps import account_initialization_required
+
+
+def get_oauth_providers():
+    with current_app.app_context():
+        notion_oauth = NotionOAuth(client_id=current_app.config.get('NOTION_CLIENT_ID'),
+                                   client_secret=current_app.config.get(
+                                       'NOTION_CLIENT_SECRET'),
+                                   redirect_uri=current_app.config.get(
+                                       'CONSOLE_API_URL') + '/console/api/oauth/data-source/callback/notion')
+
+        OAUTH_PROVIDERS = {
+            'notion': notion_oauth
+        }
+        return OAUTH_PROVIDERS
+
+
+class OAuthDataSource(Resource):
+    def get(self, provider: str):
+        # The role of the current user in the table must be admin or owner
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+            print(vars(oauth_provider))
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        if current_app.config.get('NOTION_INTEGRATION_TYPE') == 'internal':
+            internal_secret = current_app.config.get('NOTION_INTERNAL_SECRET')
+            oauth_provider.save_internal_access_token(internal_secret)
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
+        else:
+            auth_url = oauth_provider.get_authorization_url()
+            return redirect(auth_url)
+
+
+
+
+class OAuthDataSourceCallback(Resource):
+    def get(self, provider: str):
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        if 'code' in request.args:
+            code = request.args.get('code')
+            try:
+                oauth_provider.get_access_token(code)
+            except requests.exceptions.HTTPError as e:
+                logging.exception(
+                    f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")
+                return {'error': 'OAuth data source process failed'}, 400
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
+        elif 'error' in request.args:
+            error = request.args.get('error')
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source={error}')
+        else:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=access_denied')
+
+
+class OAuthDataSourceSync(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider, binding_id):
+        provider = str(provider)
+        binding_id = str(binding_id)
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        try:
+            oauth_provider.sync_data_source(binding_id)
+        except requests.exceptions.HTTPError as e:
+            logging.exception(
+                f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")
+            return {'error': 'OAuth data source process failed'}, 400
+
+        return {'result': 'success'}, 200
+
+
+api.add_resource(OAuthDataSource, '/oauth/data-source/<string:provider>')
+api.add_resource(OAuthDataSourceCallback, '/oauth/data-source/callback/<string:provider>')
+api.add_resource(OAuthDataSourceSync, '/oauth/data-source/<string:provider>/<uuid:binding_id>/sync')

api/controllers/console/auth/oauth.py

@@ -20,13 +20,13 @@ def get_oauth_providers():
                                    client_secret=current_app.config.get(
                                        'GITHUB_CLIENT_SECRET'),
                                    redirect_uri=current_app.config.get(
-                                       'CONSOLE_URL') + '/console/api/oauth/authorize/github')
+                                       'CONSOLE_API_URL') + '/console/api/oauth/authorize/github')
 
         google_oauth = GoogleOAuth(client_id=current_app.config.get('GOOGLE_CLIENT_ID'),
                                    client_secret=current_app.config.get(
                                        'GOOGLE_CLIENT_SECRET'),
                                    redirect_uri=current_app.config.get(
-                                       'CONSOLE_URL') + '/console/api/oauth/authorize/google')
+                                       'CONSOLE_API_URL') + '/console/api/oauth/authorize/google')
 
         OAUTH_PROVIDERS = {
             'github': github_oauth,
@@ -80,7 +80,7 @@ class OAuthCallback(Resource):
         flask_login.login_user(account, remember=True)
         AccountService.update_last_login(account, request)
 
-        return redirect(f'{current_app.config.get("CONSOLE_URL")}?oauth_login=success')
+        return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_login=success')
 
 
 def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) -> Optional[Account]:

api/controllers/console/datasets/data_source.py

@@ -0,0 +1,304 @@
+import datetime
+import json
+
+from cachetools import TTLCache
+from flask import request, current_app
+from flask_login import login_required, current_user
+from flask_restful import Resource, marshal_with, fields, reqparse, marshal
+from werkzeug.exceptions import NotFound
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.data_loader.loader.notion import NotionLoader
+from core.indexing_runner import IndexingRunner
+from extensions.ext_database import db
+from libs.helper import TimestampField
+from models.dataset import Document
+from models.source import DataSourceBinding
+from services.dataset_service import DatasetService, DocumentService
+from tasks.document_indexing_sync_task import document_indexing_sync_task
+
+cache = TTLCache(maxsize=None, ttl=30)
+
+FILE_SIZE_LIMIT = 15 * 1024 * 1024  # 15MB
+ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm']
+PREVIEW_WORDS_LIMIT = 3000
+
+
+class DataSourceApi(Resource):
+    integrate_icon_fields = {
+        'type': fields.String,
+        'url': fields.String,
+        'emoji': fields.String
+    }
+    integrate_page_fields = {
+        'page_name': fields.String,
+        'page_id': fields.String,
+        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
+        'parent_id': fields.String,
+        'type': fields.String
+    }
+    integrate_workspace_fields = {
+        'workspace_name': fields.String,
+        'workspace_id': fields.String,
+        'workspace_icon': fields.String,
+        'pages': fields.List(fields.Nested(integrate_page_fields)),
+        'total': fields.Integer
+    }
+    integrate_fields = {
+        'id': fields.String,
+        'provider': fields.String,
+        'created_at': TimestampField,
+        'is_bound': fields.Boolean,
+        'disabled': fields.Boolean,
+        'link': fields.String,
+        'source_info': fields.Nested(integrate_workspace_fields)
+    }
+    integrate_list_fields = {
+        'data': fields.List(fields.Nested(integrate_fields)),
+    }
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(integrate_list_fields)
+    def get(self):
+        # get workspace data source integrates
+        data_source_integrates = db.session.query(DataSourceBinding).filter(
+            DataSourceBinding.tenant_id == current_user.current_tenant_id,
+            DataSourceBinding.disabled == False
+        ).all()
+
+        base_url = request.url_root.rstrip('/')
+        data_source_oauth_base_path = "/console/api/oauth/data-source"
+        providers = ["notion"]
+
+        integrate_data = []
+        for provider in providers:
+            # existing_integrate = next((ai for ai in data_source_integrates if ai.provider == provider), None)
+            existing_integrates = filter(lambda item: item.provider == provider, data_source_integrates)
+            if existing_integrates:
+                for existing_integrate in list(existing_integrates):
+                    integrate_data.append({
+                        'id': existing_integrate.id,
+                        'provider': provider,
+                        'created_at': existing_integrate.created_at,
+                        'is_bound': True,
+                        'disabled': existing_integrate.disabled,
+                        'source_info': existing_integrate.source_info,
+                        'link': f'{base_url}{data_source_oauth_base_path}/{provider}'
+                })
+            else:
+                integrate_data.append({
+                    'id': None,
+                    'provider': provider,
+                    'created_at': None,
+                    'source_info': None,
+                    'is_bound': False,
+                    'disabled': None,
+                    'link': f'{base_url}{data_source_oauth_base_path}/{provider}'
+                })
+        return {'data': integrate_data}, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(self, binding_id, action):
+        binding_id = str(binding_id)
+        action = str(action)
+        data_source_binding = DataSourceBinding.query.filter_by(
+            id=binding_id
+        ).first()
+        if data_source_binding is None:
+            raise NotFound('Data source binding not found.')
+        # enable binding
+        if action == 'enable':
+            if data_source_binding.disabled:
+                data_source_binding.disabled = False
+                data_source_binding.updated_at = datetime.datetime.utcnow()
+                db.session.add(data_source_binding)
+                db.session.commit()
+            else:
+                raise ValueError('Data source is not disabled.')
+        # disable binding
+        if action == 'disable':
+            if not data_source_binding.disabled:
+                data_source_binding.disabled = True
+                data_source_binding.updated_at = datetime.datetime.utcnow()
+                db.session.add(data_source_binding)
+                db.session.commit()
+            else:
+                raise ValueError('Data source is disabled.')
+        return {'result': 'success'}, 200
+
+
+class DataSourceNotionListApi(Resource):
+    integrate_icon_fields = {
+        'type': fields.String,
+        'url': fields.String,
+        'emoji': fields.String
+    }
+    integrate_page_fields = {
+        'page_name': fields.String,
+        'page_id': fields.String,
+        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
+        'is_bound': fields.Boolean,
+        'parent_id': fields.String,
+        'type': fields.String
+    }
+    integrate_workspace_fields = {
+        'workspace_name': fields.String,
+        'workspace_id': fields.String,
+        'workspace_icon': fields.String,
+        'pages': fields.List(fields.Nested(integrate_page_fields))
+    }
+    integrate_notion_info_list_fields = {
+        'notion_info': fields.List(fields.Nested(integrate_workspace_fields)),
+    }
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(integrate_notion_info_list_fields)
+    def get(self):
+        dataset_id = request.args.get('dataset_id', default=None, type=str)
+        exist_page_ids = []
+        # import notion in the exist dataset
+        if dataset_id:
+            dataset = DatasetService.get_dataset(dataset_id)
+            if not dataset:
+                raise NotFound('Dataset not found.')
+            if dataset.data_source_type != 'notion_import':
+                raise ValueError('Dataset is not notion type.')
+            documents = Document.query.filter_by(
+                dataset_id=dataset_id,
+                tenant_id=current_user.current_tenant_id,
+                data_source_type='notion_import',
+                enabled=True
+            ).all()
+            if documents:
+                for document in documents:
+                    data_source_info = json.loads(document.data_source_info)
+                    exist_page_ids.append(data_source_info['notion_page_id'])
+        # get all authorized pages
+        data_source_bindings = DataSourceBinding.query.filter_by(
+            tenant_id=current_user.current_tenant_id,
+            provider='notion',
+            disabled=False
+        ).all()
+        if not data_source_bindings:
+            return {
+                'notion_info': []
+            }, 200
+        pre_import_info_list = []
+        for data_source_binding in data_source_bindings:
+            source_info = data_source_binding.source_info
+            pages = source_info['pages']
+            # Filter out already bound pages
+            for page in pages:
+                if page['page_id'] in exist_page_ids:
+                    page['is_bound'] = True
+                else:
+                    page['is_bound'] = False
+            pre_import_info = {
+                'workspace_name': source_info['workspace_name'],
+                'workspace_icon': source_info['workspace_icon'],
+                'workspace_id': source_info['workspace_id'],
+                'pages': pages,
+            }
+            pre_import_info_list.append(pre_import_info)
+        return {
+            'notion_info': pre_import_info_list
+        }, 200
+
+
+class DataSourceNotionApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, workspace_id, page_id, page_type):
+        workspace_id = str(workspace_id)
+        page_id = str(page_id)
+        data_source_binding = DataSourceBinding.query.filter(
+            db.and_(
+                DataSourceBinding.tenant_id == current_user.current_tenant_id,
+                DataSourceBinding.provider == 'notion',
+                DataSourceBinding.disabled == False,
+                DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
+            )
+        ).first()
+        if not data_source_binding:
+            raise NotFound('Data source binding not found.')
+
+        loader = NotionLoader(
+            notion_access_token=data_source_binding.access_token,
+            notion_workspace_id=workspace_id,
+            notion_obj_id=page_id,
+            notion_page_type=page_type
+        )
+
+        text_docs = loader.load()
+        return {
+            'content': "\n".join([doc.page_content for doc in text_docs])
+        }, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('notion_info_list', type=list, required=True, nullable=True, location='json')
+        parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
+        args = parser.parse_args()
+        # validate args
+        DocumentService.estimate_args_validate(args)
+        indexing_runner = IndexingRunner()
+        response = indexing_runner.notion_indexing_estimate(args['notion_info_list'], args['process_rule'])
+        return response, 200
+
+
+class DataSourceNotionDatasetSyncApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(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.")
+
+        documents = DocumentService.get_document_by_dataset_id(dataset_id_str)
+        for document in documents:
+            document_indexing_sync_task.delay(dataset_id_str, document.id)
+        return 200
+
+
+class DataSourceNotionDocumentSyncApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id, document_id):
+        dataset_id_str = str(dataset_id)
+        document_id_str = str(document_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+
+        document = DocumentService.get_document(dataset_id_str, document_id_str)
+        if document is None:
+            raise NotFound("Document not found.")
+        document_indexing_sync_task.delay(dataset_id_str, document_id_str)
+        return 200
+
+
+api.add_resource(DataSourceApi, '/data-source/integrates', '/data-source/integrates/<uuid:binding_id>/<string:action>')
+api.add_resource(DataSourceNotionListApi, '/notion/pre-import/pages')
+api.add_resource(DataSourceNotionApi,
+                 '/notion/workspaces/<uuid:workspace_id>/pages/<uuid:page_id>/<string:page_type>/preview',
+                 '/datasets/notion-indexing-estimate')
+api.add_resource(DataSourceNotionDatasetSyncApi, '/datasets/<uuid:dataset_id>/notion/sync')
+api.add_resource(DataSourceNotionDocumentSyncApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/notion/sync')

api/controllers/console/datasets/datasets.py

@@ -12,8 +12,9 @@ from controllers.console.wraps import account_initialization_required
 from core.indexing_runner import IndexingRunner
 from libs.helper import TimestampField
 from extensions.ext_database import db
+from models.dataset import DocumentSegment, Document
 from models.model import UploadFile
-from services.dataset_service import DatasetService
+from services.dataset_service import DatasetService, DocumentService
 
 dataset_detail_fields = {
     'id': fields.String,
@@ -217,17 +218,31 @@ class DatasetIndexingEstimateApi(Resource):
     @login_required
     @account_initialization_required
     def post(self):
-        segment_rule = request.get_json()
-        file_detail = db.session.query(UploadFile).filter(
-            UploadFile.tenant_id == current_user.current_tenant_id,
-            UploadFile.id == segment_rule["file_id"]
-        ).first()
+        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')
+        args = parser.parse_args()
+        # validate args
+        DocumentService.estimate_args_validate(args)
+        if args['info_list']['data_source_type'] == 'upload_file':
+            file_ids = args['info_list']['file_info_list']['file_ids']
+            file_details = db.session.query(UploadFile).filter(
+                UploadFile.tenant_id == current_user.current_tenant_id,
+                UploadFile.id.in_(file_ids)
+            ).all()
 
-        if file_detail is None:
-            raise NotFound("File not found.")
+            if file_details is None:
+                raise NotFound("File not found.")
 
-        indexing_runner = IndexingRunner()
-        response = indexing_runner.indexing_estimate(file_detail, segment_rule['process_rule'])
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.file_indexing_estimate(file_details, args['process_rule'])
+        elif args['info_list']['data_source_type'] == 'notion_import':
+
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.notion_indexing_estimate(args['info_list']['notion_info_list'],
+                                                                args['process_rule'])
+        else:
+            raise ValueError('Data source type not support')
         return response, 200
 
 
@@ -274,8 +289,54 @@ class DatasetRelatedAppListApi(Resource):
         }, 200
 
 
+class DatasetIndexingStatusApi(Resource):
+    document_status_fields = {
+        'id': fields.String,
+        'indexing_status': fields.String,
+        'processing_started_at': TimestampField,
+        'parsing_completed_at': TimestampField,
+        'cleaning_completed_at': TimestampField,
+        'splitting_completed_at': TimestampField,
+        'completed_at': TimestampField,
+        'paused_at': TimestampField,
+        'error': fields.String,
+        'stopped_at': TimestampField,
+        'completed_segments': fields.Integer,
+        'total_segments': fields.Integer,
+    }
+
+    document_status_fields_list = {
+        'data': fields.List(fields.Nested(document_status_fields))
+    }
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id):
+        dataset_id = str(dataset_id)
+        documents = db.session.query(Document).filter(
+            Document.dataset_id == dataset_id,
+            Document.tenant_id == current_user.current_tenant_id
+        ).all()
+        documents_status = []
+        for document in documents:
+            completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+                                                              DocumentSegment.document_id == str(document.id),
+                                                              DocumentSegment.status != 're_segment').count()
+            total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
+                                                          DocumentSegment.status != 're_segment').count()
+            document.completed_segments = completed_segments
+            document.total_segments = total_segments
+            documents_status.append(marshal(document, self.document_status_fields))
+        data = {
+            'data': documents_status
+        }
+        return data
+
+
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
 api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
-api.add_resource(DatasetIndexingEstimateApi, '/datasets/file-indexing-estimate')
+api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
 api.add_resource(DatasetRelatedAppListApi, '/datasets/<uuid:dataset_id>/related-apps')
+api.add_resource(DatasetIndexingStatusApi, '/datasets/<uuid:dataset_id>/indexing-status')

api/controllers/console/datasets/datasets_document.py

@@ -1,6 +1,7 @@
 # -*- coding:utf-8 -*-
 import random
 from datetime import datetime
+from typing import List
 
 from flask import request
 from flask_login import login_required, current_user
@@ -61,6 +62,29 @@ document_fields = {
     'hit_count': fields.Integer,
 }
 
+document_with_segments_fields = {
+    'id': fields.String,
+    'position': fields.Integer,
+    'data_source_type': fields.String,
+    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
+    'dataset_process_rule_id': fields.String,
+    'name': fields.String,
+    'created_from': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'tokens': fields.Integer,
+    'indexing_status': fields.String,
+    'error': fields.String,
+    'enabled': fields.Boolean,
+    'disabled_at': TimestampField,
+    'disabled_by': fields.String,
+    'archived': fields.Boolean,
+    'display_status': fields.String,
+    'word_count': fields.Integer,
+    'hit_count': fields.Integer,
+    'completed_segments': fields.Integer,
+    'total_segments': fields.Integer
+}
 
 class DocumentResource(Resource):
     def get_document(self, dataset_id: str, document_id: str) -> Document:
@@ -83,6 +107,23 @@ class DocumentResource(Resource):
 
         return document
 
+    def get_batch_documents(self, dataset_id: str, batch: str) -> List[Document]:
+        dataset = DatasetService.get_dataset(dataset_id)
+        if not dataset:
+            raise NotFound('Dataset not found.')
+
+        try:
+            DatasetService.check_dataset_permission(dataset, current_user)
+        except services.errors.account.NoPermissionError as e:
+            raise Forbidden(str(e))
+
+        documents = DocumentService.get_batch_documents(dataset_id, batch)
+
+        if not documents:
+            raise NotFound('Documents not found.')
+
+        return documents
+
 
 class GetProcessRuleApi(Resource):
     @setup_required
@@ -132,9 +173,9 @@ class DatasetDocumentListApi(Resource):
         dataset_id = str(dataset_id)
         page = request.args.get('page', default=1, type=int)
         limit = request.args.get('limit', default=20, type=int)
-        search = request.args.get('search', default=None, type=str)
+        search = request.args.get('keyword', default=None, type=str)
         sort = request.args.get('sort', default='-created_at', type=str)
-
+        fetch = request.args.get('fetch', default=False, type=bool)
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
@@ -173,9 +214,20 @@ class DatasetDocumentListApi(Resource):
         paginated_documents = query.paginate(
             page=page, per_page=limit, max_per_page=100, error_out=False)
         documents = paginated_documents.items
-
+        if fetch:
+            for document in documents:
+                completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+                                                                  DocumentSegment.document_id == str(document.id),
+                                                                  DocumentSegment.status != 're_segment').count()
+                total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
+                                                              DocumentSegment.status != 're_segment').count()
+                document.completed_segments = completed_segments
+                document.total_segments = total_segments
+            data = marshal(documents, document_with_segments_fields)
+        else:
+            data = marshal(documents, document_fields)
         response = {
-            'data': marshal(documents, document_fields),
+            'data': data,
             'has_more': len(documents) == limit,
             'limit': limit,
             'total': paginated_documents.total,
@@ -184,10 +236,15 @@ class DatasetDocumentListApi(Resource):
 
         return response
 
+    documents_and_batch_fields = {
+        'documents': fields.List(fields.Nested(document_fields)),
+        'batch': fields.String
+    }
+
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(document_fields)
+    @marshal_with(documents_and_batch_fields)
     def post(self, dataset_id):
         dataset_id = str(dataset_id)
 
@@ -221,7 +278,7 @@ class DatasetDocumentListApi(Resource):
         DocumentService.document_create_args_validate(args)
 
         try:
-            document = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
+            documents, batch = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
         except ProviderTokenNotInitError:
             raise ProviderNotInitializeError()
         except QuotaExceededError:
@@ -229,13 +286,17 @@ class DatasetDocumentListApi(Resource):
         except ModelCurrentlyNotSupportError:
             raise ProviderModelCurrentlyNotSupportError()
 
-        return document
+        return {
+            'documents': documents,
+            'batch': batch
+        }
 
 
 class DatasetInitApi(Resource):
     dataset_and_document_fields = {
         'dataset': fields.Nested(dataset_fields),
-        'document': fields.Nested(document_fields)
+        'documents': fields.List(fields.Nested(document_fields)),
+        'batch': fields.String
     }
 
     @setup_required
@@ -258,7 +319,7 @@ class DatasetInitApi(Resource):
         DocumentService.document_create_args_validate(args)
 
         try:
-            dataset, document = DocumentService.save_document_without_dataset_id(
+            dataset, documents, batch = DocumentService.save_document_without_dataset_id(
                 tenant_id=current_user.current_tenant_id,
                 document_data=args,
                 account=current_user
@@ -272,7 +333,8 @@ class DatasetInitApi(Resource):
 
         response = {
             'dataset': dataset,
-            'document': document
+            'documents': documents,
+            'batch': batch
         }
 
         return response
@@ -317,11 +379,122 @@ class DocumentIndexingEstimateApi(DocumentResource):
                     raise NotFound('File not found.')
 
                 indexing_runner = IndexingRunner()
-                response = indexing_runner.indexing_estimate(file, data_process_rule_dict)
+
+                response = indexing_runner.file_indexing_estimate([file], data_process_rule_dict)
 
         return response
 
 
+class DocumentBatchIndexingEstimateApi(DocumentResource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id, batch):
+        dataset_id = str(dataset_id)
+        batch = str(batch)
+        dataset = DatasetService.get_dataset(dataset_id)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        documents = self.get_batch_documents(dataset_id, batch)
+        response = {
+            "tokens": 0,
+            "total_price": 0,
+            "currency": "USD",
+            "total_segments": 0,
+            "preview": []
+        }
+        if not documents:
+            return response
+        data_process_rule = documents[0].dataset_process_rule
+        data_process_rule_dict = data_process_rule.to_dict()
+        info_list = []
+        for document in documents:
+            if document.indexing_status in ['completed', 'error']:
+                raise DocumentAlreadyFinishedError()
+            data_source_info = document.data_source_info_dict
+            # format document files info
+            if data_source_info and 'upload_file_id' in data_source_info:
+                file_id = data_source_info['upload_file_id']
+                info_list.append(file_id)
+            # format document notion info
+            elif data_source_info and 'notion_workspace_id' in data_source_info and 'notion_page_id' in data_source_info:
+                pages = []
+                page = {
+                    'page_id': data_source_info['notion_page_id'],
+                    'type': data_source_info['type']
+                }
+                pages.append(page)
+                notion_info = {
+                    'workspace_id': data_source_info['notion_workspace_id'],
+                    'pages': pages
+                }
+                info_list.append(notion_info)
+
+        if dataset.data_source_type == 'upload_file':
+            file_details = db.session.query(UploadFile).filter(
+                UploadFile.tenant_id == current_user.current_tenant_id,
+                UploadFile.id in info_list
+            ).all()
+
+            if file_details is None:
+                raise NotFound("File not found.")
+
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.file_indexing_estimate(file_details, data_process_rule_dict)
+        elif dataset.data_source_type:
+
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.notion_indexing_estimate(info_list,
+                                                                data_process_rule_dict)
+        else:
+            raise ValueError('Data source type not support')
+        return response
+
+
+class DocumentBatchIndexingStatusApi(DocumentResource):
+    document_status_fields = {
+        'id': fields.String,
+        'indexing_status': fields.String,
+        'processing_started_at': TimestampField,
+        'parsing_completed_at': TimestampField,
+        'cleaning_completed_at': TimestampField,
+        'splitting_completed_at': TimestampField,
+        'completed_at': TimestampField,
+        'paused_at': TimestampField,
+        'error': fields.String,
+        'stopped_at': TimestampField,
+        'completed_segments': fields.Integer,
+        'total_segments': fields.Integer,
+    }
+
+    document_status_fields_list = {
+        'data': fields.List(fields.Nested(document_status_fields))
+    }
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id, batch):
+        dataset_id = str(dataset_id)
+        batch = str(batch)
+        documents = self.get_batch_documents(dataset_id, batch)
+        documents_status = []
+        for document in documents:
+            completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+                                                              DocumentSegment.document_id == str(document.id),
+                                                              DocumentSegment.status != 're_segment').count()
+            total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
+                                                          DocumentSegment.status != 're_segment').count()
+            document.completed_segments = completed_segments
+            document.total_segments = total_segments
+            documents_status.append(marshal(document, self.document_status_fields))
+        data = {
+            'data': documents_status
+        }
+        return data
+
+
 class DocumentIndexingStatusApi(DocumentResource):
     document_status_fields = {
         'id': fields.String,
@@ -408,7 +581,7 @@ class DocumentDetailApi(DocumentResource):
                 'disabled_by': document.disabled_by,
                 'archived': document.archived,
                 'segment_count': document.segment_count,
-                'average_segment_length':   document.average_segment_length,
+                'average_segment_length': document.average_segment_length,
                 'hit_count': document.hit_count,
                 'display_status': document.display_status
             }
@@ -428,7 +601,7 @@ class DocumentDetailApi(DocumentResource):
                 'created_at': document.created_at.timestamp(),
                 'tokens': document.tokens,
                 'indexing_status': document.indexing_status,
-                'completed_at': int(document.completed_at.timestamp())if document.completed_at else None,
+                'completed_at': int(document.completed_at.timestamp()) if document.completed_at else None,
                 'updated_at': int(document.updated_at.timestamp()) if document.updated_at else None,
                 'indexing_latency': document.indexing_latency,
                 'error': document.error,
@@ -579,6 +752,8 @@ class DocumentStatusApi(DocumentResource):
             return {'result': 'success'}, 200
 
         elif action == "disable":
+            if not document.completed_at or document.indexing_status != 'completed':
+                raise InvalidActionError('Document is not completed.')
             if not document.enabled:
                 raise InvalidActionError('Document already disabled.')
 
@@ -678,6 +853,10 @@ api.add_resource(DatasetInitApi,
                  '/datasets/init')
 api.add_resource(DocumentIndexingEstimateApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate')
+api.add_resource(DocumentBatchIndexingEstimateApi,
+                 '/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-estimate')
+api.add_resource(DocumentBatchIndexingStatusApi,
+                 '/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-status')
 api.add_resource(DocumentIndexingStatusApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-status')
 api.add_resource(DocumentDetailApi,

api/controllers/console/datasets/file.py

@@ -1,6 +1,7 @@
 import datetime
 import hashlib
 import tempfile
+import chardet
 import time
 import uuid
 from pathlib import Path
@@ -16,9 +17,7 @@ from controllers.console.datasets.error import NoFileUploadedError, TooManyFiles
     UnsupportedFileTypeError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.index.readers.html_parser import HTMLParser
-from core.index.readers.pdf_parser import PDFParser
-from core.index.readers.xlsx_parser import XLSXParser
+from core.data_loader.file_extractor import FileExtractor
 from extensions.ext_storage import storage
 from libs.helper import TimestampField
 from extensions.ext_database import db
@@ -122,27 +121,7 @@ class FilePreviewApi(Resource):
         if extension not in ALLOWED_EXTENSIONS:
             raise UnsupportedFileTypeError()
 
-        with tempfile.TemporaryDirectory() as temp_dir:
-            suffix = Path(upload_file.key).suffix
-            filepath = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
-            storage.download(upload_file.key, filepath)
-
-            if extension == 'pdf':
-                parser = PDFParser({'upload_file': upload_file})
-                text = parser.parse_file(Path(filepath))
-            elif extension in ['html', 'htm']:
-                # Use BeautifulSoup to extract text
-                parser = HTMLParser()
-                text = parser.parse_file(Path(filepath))
-            elif extension == 'xlsx':
-                parser = XLSXParser()
-                text = parser.parse_file(filepath)
-            else:
-                # ['txt', 'markdown', 'md']
-                with open(filepath, "rb") as fp:
-                    data = fp.read()
-                    text = data.decode(encoding='utf-8').strip() if data else ''
-
+        text = FileExtractor.load(upload_file, return_text=True)
         text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
         return {'content': text}
 

api/controllers/console/error.py

@@ -18,3 +18,9 @@ class AccountNotLinkTenantError(BaseHTTPException):
     error_code = 'account_not_link_tenant'
     description = "Account not link tenant."
     code = 403
+
+
+class AlreadyActivateError(BaseHTTPException):
+    error_code = 'already_activate'
+    description = "Auth Token is invalid or account already activated, please check again."
+    code = 403

api/controllers/console/explore/audio.py

@@ -0,0 +1,66 @@
+# -*- coding:utf-8 -*-
+import logging
+
+from flask import request
+from werkzeug.exceptions import InternalServerError
+
+import services
+from controllers.console import api
+from controllers.console.app.error import AppUnavailableError, ProviderNotInitializeError, \
+    ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, CompletionRequestError, \
+    NoAudioUploadedError, AudioTooLargeError, \
+    UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
+from controllers.console.explore.wraps import InstalledAppResource
+from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+    LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from services.audio_service import AudioService
+from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \
+    UnsupportedAudioTypeServiceError, ProviderNotSupportSpeechToTextServiceError
+from models.model import AppModelConfig
+
+
+class ChatAudioApi(InstalledAppResource):
+    def post(self, installed_app):
+        app_model = installed_app.app
+        app_model_config: AppModelConfig = app_model.app_model_config
+
+        if not app_model_config.speech_to_text_dict['enabled']:
+            raise AppUnavailableError()
+
+        file = request.files['file']
+
+        try:
+            response = AudioService.transcript(
+                tenant_id=app_model.tenant_id,
+                file=file,
+            )
+
+            return response
+        except services.errors.app_model_config.AppModelConfigBrokenError:
+            logging.exception("App model config broken.")
+            raise AppUnavailableError()
+        except NoAudioUploadedServiceError:
+            raise NoAudioUploadedError()
+        except AudioTooLargeServiceError as e:
+            raise AudioTooLargeError(str(e))
+        except UnsupportedAudioTypeServiceError:
+            raise UnsupportedAudioTypeError()
+        except ProviderNotSupportSpeechToTextServiceError:
+            raise ProviderNotSupportSpeechToTextError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except ValueError as e:
+            raise e
+        except Exception as e:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+        
+
+api.add_resource(ChatAudioApi, '/installed-apps/<uuid:installed_app_id>/audio-to-text', endpoint='installed_app_audio')

api/controllers/console/explore/parameter.py

@@ -21,6 +21,7 @@ class AppParameterApi(InstalledAppResource):
         'opening_statement': fields.String,
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
+        'speech_to_text': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -35,6 +36,7 @@ class AppParameterApi(InstalledAppResource):
             'opening_statement': app_model_config.opening_statement,
             'suggested_questions': app_model_config.suggested_questions_list,
             'suggested_questions_after_answer': app_model_config.suggested_questions_after_answer_dict,
+            'speech_to_text': app_model_config.speech_to_text_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/console/version.py

@@ -32,8 +32,13 @@ class VersionApi(Resource):
                 'current_version': args.get('current_version')
             })
         except Exception as error:
-            logging.exception("Check update error.")
-            raise InternalServerError()
+            logging.warning("Check update version error: {}.".format(str(error)))
+            return {
+                'version': args.get('current_version'),
+                'release_date': '',
+                'release_notes': '',
+                'can_auto_update': False
+            }
 
         content = json.loads(response.content)
         return {

api/controllers/console/workspace/account.py

@@ -6,22 +6,23 @@ from flask import current_app, request
 from flask_login import login_required, current_user
 from flask_restful import Resource, reqparse, fields, marshal_with
 
+from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError
 from controllers.console import api
 from controllers.console.setup import setup_required
 from controllers.console.workspace.error import AccountAlreadyInitedError, InvalidInvitationCodeError, \
-    RepeatPasswordNotMatchError
+    RepeatPasswordNotMatchError, CurrentPasswordIncorrectError
 from controllers.console.wraps import account_initialization_required
 from libs.helper import TimestampField, supported_language, timezone
 from extensions.ext_database import db
 from models.account import InvitationCode, AccountIntegrate
 from services.account_service import AccountService
 
-
 account_fields = {
     'id': fields.String,
     'name': fields.String,
     'avatar': fields.String,
     'email': fields.String,
+    'is_password_set': fields.Boolean,
     'interface_language': fields.String,
     'interface_theme': fields.String,
     'timezone': fields.String,
@@ -194,8 +195,11 @@ class AccountPasswordApi(Resource):
         if args['new_password'] != args['repeat_new_password']:
             raise RepeatPasswordNotMatchError()
 
-        AccountService.update_account_password(
-            current_user, args['password'], args['new_password'])
+        try:
+            AccountService.update_account_password(
+                current_user, args['password'], args['new_password'])
+        except ServiceCurrentPasswordIncorrectError:
+            raise CurrentPasswordIncorrectError()
 
         return {"result": "success"}
 

api/controllers/console/workspace/error.py

@@ -7,6 +7,12 @@ class RepeatPasswordNotMatchError(BaseHTTPException):
     code = 400
 
 
+class CurrentPasswordIncorrectError(BaseHTTPException):
+    error_code = 'current_password_incorrect'
+    description = "Current password is incorrect."
+    code = 400
+
+
 class ProviderRequestFailedError(BaseHTTPException):
     error_code = 'provider_request_failed'
     description = None

api/controllers/console/workspace/members.py

@@ -1,5 +1,5 @@
 # -*- coding:utf-8 -*-
-
+from flask import current_app
 from flask_login import login_required, current_user
 from flask_restful import Resource, reqparse, marshal_with, abort, fields, marshal
 
@@ -60,7 +60,8 @@ class MemberInviteEmailApi(Resource):
         inviter = current_user
 
         try:
-            RegisterService.invite_new_member(inviter.current_tenant, invitee_email, role=invitee_role, inviter=inviter)
+            token = RegisterService.invite_new_member(inviter.current_tenant, invitee_email, role=invitee_role,
+                                                      inviter=inviter)
             account = db.session.query(Account, TenantAccountJoin.role).join(
                 TenantAccountJoin, Account.id == TenantAccountJoin.account_id
             ).filter(Account.email == args['email']).first()
@@ -78,7 +79,16 @@ class MemberInviteEmailApi(Resource):
 
         # todo:413
 
-        return {'result': 'success', 'account': account}, 201
+        return {
+            'result': 'success',
+            'account': account,
+            'invite_url': '{}/activate?workspace_id={}&email={}&token={}'.format(
+                current_app.config.get("CONSOLE_WEB_URL"),
+                str(current_user.current_tenant_id),
+                invitee_email,
+                token
+            )
+        }, 201
 
 
 class MemberCancelInviteApi(Resource):
@@ -88,7 +98,7 @@ class MemberCancelInviteApi(Resource):
     @login_required
     @account_initialization_required
     def delete(self, member_id):
-        member = Account.query.get(str(member_id))
+        member = db.session.query(Account).filter(Account.id == str(member_id)).first()
         if not member:
             abort(404)
 

api/controllers/service_api/__init__.py

@@ -7,6 +7,6 @@ bp = Blueprint('service_api', __name__, url_prefix='/v1')
 api = ExternalApi(bp)
 
 
-from .app import completion, app, conversation, message
+from .app import completion, app, conversation, message, audio
 
 from .dataset import document

api/controllers/service_api/app/app.py

@@ -22,6 +22,7 @@ class AppParameterApi(AppApiResource):
         'opening_statement': fields.String,
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
+        'speech_to_text': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -35,6 +36,7 @@ class AppParameterApi(AppApiResource):
             'opening_statement': app_model_config.opening_statement,
             'suggested_questions': app_model_config.suggested_questions_list,
             'suggested_questions_after_answer': app_model_config.suggested_questions_after_answer_dict,
+            'speech_to_text': app_model_config.speech_to_text_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/service_api/app/audio.py

@@ -0,0 +1,61 @@
+import logging
+
+from flask import request
+from werkzeug.exceptions import InternalServerError
+
+import services
+from controllers.service_api import api
+from controllers.service_api.app.error import AppUnavailableError, ProviderNotInitializeError, CompletionRequestError, ProviderQuotaExceededError, \
+    ProviderModelCurrentlyNotSupportError, NoAudioUploadedError, AudioTooLargeError, UnsupportedAudioTypeError, \
+    ProviderNotSupportSpeechToTextError
+from controllers.service_api.wraps import AppApiResource
+from core.llm.error import LLMBadRequestError, LLMAuthorizationError, LLMAPIUnavailableError, LLMAPIConnectionError, \
+    LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from models.model import App, AppModelConfig
+from services.audio_service import AudioService
+from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \
+    UnsupportedAudioTypeServiceError, ProviderNotSupportSpeechToTextServiceError
+
+class AudioApi(AppApiResource):
+    def post(self, app_model: App, end_user):
+        app_model_config: AppModelConfig = app_model.app_model_config
+
+        if not app_model_config.speech_to_text_dict['enabled']:
+            raise AppUnavailableError() 
+
+        file = request.files['file']
+
+        try:
+            response = AudioService.transcript(
+                tenant_id=app_model.tenant_id,
+                file=file,
+            )
+
+            return response
+        except services.errors.app_model_config.AppModelConfigBrokenError:
+            logging.exception("App model config broken.")
+            raise AppUnavailableError()
+        except NoAudioUploadedServiceError:
+            raise NoAudioUploadedError()
+        except AudioTooLargeServiceError as e:
+            raise AudioTooLargeError(str(e))
+        except UnsupportedAudioTypeServiceError:
+            raise UnsupportedAudioTypeError()
+        except ProviderNotSupportSpeechToTextServiceError:
+            raise ProviderNotSupportSpeechToTextError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except ValueError as e:
+            raise e
+        except Exception as e:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+        
+api.add_resource(AudioApi, '/audio-to-text')

api/controllers/service_api/app/conversation.py

@@ -1,4 +1,5 @@
 # -*- coding:utf-8 -*-
+from flask import request
 from flask_restful import fields, marshal_with, reqparse
 from flask_restful.inputs import int_range
 from werkzeug.exceptions import NotFound
@@ -48,6 +49,24 @@ class ConversationApi(AppApiResource):
         except services.errors.conversation.LastConversationNotExistsError:
             raise NotFound("Last Conversation Not Exists.")
 
+class ConversationDetailApi(AppApiResource):
+    @marshal_with(conversation_fields)
+    def delete(self, app_model, end_user, c_id):
+        if app_model.mode != 'chat':
+            raise NotChatAppError()
+
+        conversation_id = str(c_id)
+
+        user = request.get_json().get('user')
+
+        if end_user is None and user is not None:
+            end_user = create_or_update_end_user_for_user_id(app_model, user)
+
+        try:
+            ConversationService.delete(app_model, conversation_id, end_user)
+            return {"result": "success"}
+        except services.errors.conversation.ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
 
 class ConversationRenameApi(AppApiResource):
 
@@ -74,3 +93,5 @@ class ConversationRenameApi(AppApiResource):
 
 api.add_resource(ConversationRenameApi, '/conversations/<uuid:c_id>/name', endpoint='conversation_name')
 api.add_resource(ConversationApi, '/conversations')
+api.add_resource(ConversationApi, '/conversations/<uuid:c_id>', endpoint='conversation')
+api.add_resource(ConversationDetailApi, '/conversations/<uuid:c_id>', endpoint='conversation_detail')

api/controllers/service_api/app/error.py

@@ -51,3 +51,27 @@ class CompletionRequestError(BaseHTTPException):
     description = "Completion request failed."
     code = 400
 
+
+class NoAudioUploadedError(BaseHTTPException):
+    error_code = 'no_audio_uploaded'
+    description = "Please upload your audio."
+    code = 400
+
+
+class AudioTooLargeError(BaseHTTPException):
+    error_code = 'audio_too_large'
+    description = "Audio size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedAudioTypeError(BaseHTTPException):
+    error_code = 'unsupported_audio_type'
+    description = "Audio type not allowed."
+    code = 415
+
+
+class ProviderNotSupportSpeechToTextError(BaseHTTPException):
+    error_code = 'provider_not_support_speech_to_text'
+    description = "Provider not support speech to text."
+    code = 400
+

api/controllers/service_api/dataset/document.py

@@ -69,12 +69,16 @@ class DocumentListApi(DatasetApiResource):
         document_data = {
             'data_source': {
                 'type': 'upload_file',
-                'info': upload_file.id
+                'info': [
+                    {
+                        'upload_file_id': upload_file.id
+                    }
+                ]
             }
         }
 
         try:
-            document = DocumentService.save_document_with_dataset_id(
+            documents, batch = DocumentService.save_document_with_dataset_id(
                 dataset=dataset,
                 document_data=document_data,
                 account=dataset.created_by_account,
@@ -83,7 +87,7 @@ class DocumentListApi(DatasetApiResource):
             )
         except ProviderTokenNotInitError:
             raise ProviderNotInitializeError()
-
+        document = documents[0]
         if doc_type and doc_metadata:
             metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
 

api/controllers/web/__init__.py

@@ -7,4 +7,4 @@ bp = Blueprint('web', __name__, url_prefix='/api')
 api = ExternalApi(bp)
 
 
-from . import completion, app, conversation, message, site, saved_message
+from . import completion, app, conversation, message, site, saved_message, audio, passport

api/controllers/web/app.py

@@ -21,6 +21,7 @@ class AppParameterApi(WebApiResource):
         'opening_statement': fields.String,
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
+        'speech_to_text': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -34,6 +35,7 @@ class AppParameterApi(WebApiResource):
             'opening_statement': app_model_config.opening_statement,
             'suggested_questions': app_model_config.suggested_questions_list,
             'suggested_questions_after_answer': app_model_config.suggested_questions_after_answer_dict,
+            'speech_to_text': app_model_config.speech_to_text_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/web/audio.py

@@ -0,0 +1,63 @@
+# -*- coding:utf-8 -*-
+import logging
+
+from flask import request
+from werkzeug.exceptions import InternalServerError
+
+import services
+from controllers.web import api
+from controllers.web.error import AppUnavailableError, ProviderNotInitializeError, CompletionRequestError, \
+    ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, NoAudioUploadedError, AudioTooLargeError, \
+    UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
+from controllers.web.wraps import WebApiResource
+from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+    LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from services.audio_service import AudioService
+from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \
+    UnsupportedAudioTypeServiceError, ProviderNotSupportSpeechToTextServiceError
+from models.model import App, AppModelConfig
+
+
+class AudioApi(WebApiResource):
+    def post(self, app_model: App, end_user):
+        app_model_config: AppModelConfig = app_model.app_model_config
+
+        if not app_model_config.speech_to_text_dict['enabled']:
+            raise AppUnavailableError()
+
+        file = request.files['file']
+
+        try:
+            response = AudioService.transcript(
+                tenant_id=app_model.tenant_id,
+                file=file,
+            )
+
+            return response
+        except services.errors.app_model_config.AppModelConfigBrokenError:
+            logging.exception("App model config broken.")
+            raise AppUnavailableError()
+        except NoAudioUploadedServiceError:
+            raise NoAudioUploadedError()
+        except AudioTooLargeServiceError as e:
+            raise AudioTooLargeError(str(e))
+        except UnsupportedAudioTypeServiceError:
+            raise UnsupportedAudioTypeError()
+        except ProviderNotSupportSpeechToTextServiceError:
+            raise ProviderNotSupportSpeechToTextError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except ValueError as e:
+            raise e
+        except Exception as e:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+
+api.add_resource(AudioApi, '/audio-to-text')

api/controllers/web/error.py

@@ -62,3 +62,27 @@ class AppSuggestedQuestionsAfterAnswerDisabledError(BaseHTTPException):
     error_code = 'app_suggested_questions_after_answer_disabled'
     description = "The 'Suggested Questions After Answer' feature is disabled. Please refresh your page."
     code = 403
+
+
+class NoAudioUploadedError(BaseHTTPException):
+    error_code = 'no_audio_uploaded'
+    description = "Please upload your audio."
+    code = 400
+
+
+class AudioTooLargeError(BaseHTTPException):
+    error_code = 'audio_too_large'
+    description = "Audio size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedAudioTypeError(BaseHTTPException):
+    error_code = 'unsupported_audio_type'
+    description = "Audio type not allowed."
+    code = 415
+
+
+class ProviderNotSupportSpeechToTextError(BaseHTTPException):
+    error_code = 'provider_not_support_speech_to_text'
+    description = "Provider not support speech to text."
+    code = 400

api/controllers/web/passport.py

@@ -0,0 +1,64 @@
+# -*- coding:utf-8 -*-
+import uuid
+from controllers.web import api
+from flask_restful import Resource
+from flask import request
+from werkzeug.exceptions import Unauthorized, NotFound
+from models.model import Site, EndUser, App
+from extensions.ext_database import db
+from libs.passport import PassportService
+
+class PassportResource(Resource):
+    """Base resource for passport."""
+    def get(self):
+        app_id = request.headers.get('X-App-Code')
+        if app_id is None:
+            raise Unauthorized('X-App-Code header is missing.')
+
+        # get site from db and check if it is normal
+        site = db.session.query(Site).filter(
+            Site.code == app_id,
+            Site.status == 'normal'
+        ).first()
+        if not site:
+            raise NotFound()
+        # get app from db and check if it is normal and enable_site
+        app_model = db.session.query(App).filter(App.id == site.app_id).first()
+        if not app_model or app_model.status != 'normal' or not app_model.enable_site:
+            raise NotFound()
+        
+        end_user = EndUser(
+            tenant_id=app_model.tenant_id,
+            app_id=app_model.id,
+            type='browser',
+            is_anonymous=True,
+            session_id=generate_session_id(),
+        )
+        db.session.add(end_user)
+        db.session.commit()
+
+        payload = {
+            "iss": site.app_id,
+            'sub': 'Web API Passport',
+            'app_id': site.app_id,
+            'end_user_id': end_user.id,
+        }
+
+        tk = PassportService().issue(payload)
+
+        return {
+            'access_token': tk,
+        }
+
+api.add_resource(PassportResource, '/passport')
+
+def generate_session_id():
+    """
+    Generate a unique session ID.
+    """
+    while True:
+        session_id = str(uuid.uuid4())
+        existing_count = db.session.query(EndUser) \
+            .filter(EndUser.session_id == session_id).count()
+        if existing_count == 0:
+            return session_id

api/controllers/web/wraps.py

@@ -1,110 +1,48 @@
 # -*- coding:utf-8 -*-
-import uuid
 from functools import wraps
 
-from flask import request, session
+from flask import request
 from flask_restful import Resource
 from werkzeug.exceptions import NotFound, Unauthorized
 
 from extensions.ext_database import db
-from models.model import App, Site, EndUser
+from models.model import App, EndUser
+from libs.passport import PassportService
 
-
-def validate_token(view=None):
+def validate_jwt_token(view=None):
     def decorator(view):
         @wraps(view)
         def decorated(*args, **kwargs):
-            site = validate_and_get_site()
-
-            app_model = db.session.query(App).filter(App.id == site.app_id).first()
-            if not app_model:
-                raise NotFound()
-
-            if app_model.status != 'normal':
-                raise NotFound()
-
-            if not app_model.enable_site:
-                raise NotFound()
-
-            end_user = create_or_update_end_user_for_session(app_model)
+            app_model, end_user = decode_jwt_token()
 
             return view(app_model, end_user, *args, **kwargs)
         return decorated
-
     if view:
         return decorator(view)
     return decorator
 
-
-def validate_and_get_site():
-    """
-    Validate and get API token.
-    """
+def decode_jwt_token():
     auth_header = request.headers.get('Authorization')
     if auth_header is None:
         raise Unauthorized('Authorization header is missing.')
 
     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, tk = 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.')
-
-    site = db.session.query(Site).filter(
-        Site.code == auth_token,
-        Site.status == 'normal'
-    ).first()
-
-    if not site:
+    decoded = PassportService().verify(tk)
+    app_model = db.session.query(App).filter(App.id == decoded['app_id']).first()
+    if not app_model:
+        raise NotFound()
+    end_user = db.session.query(EndUser).filter(EndUser.id == decoded['end_user_id']).first()
+    if not end_user:
         raise NotFound()
 
-    return site
-
-
-def create_or_update_end_user_for_session(app_model):
-    """
-    Create or update session terminal based on session ID.
-    """
-    if 'session_id' not in session:
-        session['session_id'] = generate_session_id()
-
-    session_id = session.get('session_id')
-    end_user = db.session.query(EndUser) \
-        .filter(
-        EndUser.session_id == session_id,
-        EndUser.type == 'browser'
-    ).first()
-
-    if end_user is None:
-        end_user = EndUser(
-            tenant_id=app_model.tenant_id,
-            app_id=app_model.id,
-            type='browser',
-            is_anonymous=True,
-            session_id=session_id
-        )
-        db.session.add(end_user)
-        db.session.commit()
-
-    return end_user
-
-
-def generate_session_id():
-    """
-    Generate a unique session ID.
-    """
-    count = 1
-    session_id = ''
-    while count != 0:
-        session_id = str(uuid.uuid4())
-        count = db.session.query(EndUser) \
-            .filter(EndUser.session_id == session_id).count()
-
-    return session_id
-
+    return app_model, end_user
 
 class WebApiResource(Resource):
-    method_decorators = [validate_token]
+    method_decorators = [validate_jwt_token]

api/core/__init__.py

@@ -3,19 +3,10 @@ from typing import Optional
 
 import langchain
 from flask import Flask
-from jieba.analyse import default_tfidf
-from langchain import set_handler
-from langchain.prompts.base import DEFAULT_FORMATTER_MAPPING
-from llama_index import IndexStructType, QueryMode
-from llama_index.indices.registry import INDEX_STRUT_TYPE_TO_QUERY_MAP
 from pydantic import BaseModel
 
 from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.index.keyword_table.jieba_keyword_table import GPTJIEBAKeywordTableIndex
-from core.index.keyword_table.stopwords import STOPWORDS
 from core.prompt.prompt_template import OneLineFormatter
-from core.vector_store.vector_store import VectorStore
-from core.vector_store.vector_store_index_query import EnhanceGPTVectorStoreIndexQuery
 
 
 class HostedOpenAICredential(BaseModel):
@@ -30,23 +21,8 @@ hosted_llm_credentials = HostedLLMCredentials()
 
 
 def init_app(app: Flask):
-    formatter = OneLineFormatter()
-    DEFAULT_FORMATTER_MAPPING['f-string'] = formatter.format
-    INDEX_STRUT_TYPE_TO_QUERY_MAP[IndexStructType.KEYWORD_TABLE] = GPTJIEBAKeywordTableIndex.get_query_map()
-    INDEX_STRUT_TYPE_TO_QUERY_MAP[IndexStructType.WEAVIATE] = {
-        QueryMode.DEFAULT: EnhanceGPTVectorStoreIndexQuery,
-        QueryMode.EMBEDDING: EnhanceGPTVectorStoreIndexQuery,
-    }
-    INDEX_STRUT_TYPE_TO_QUERY_MAP[IndexStructType.QDRANT] = {
-        QueryMode.DEFAULT: EnhanceGPTVectorStoreIndexQuery,
-        QueryMode.EMBEDDING: EnhanceGPTVectorStoreIndexQuery,
-    }
-
-    default_tfidf.stop_words = STOPWORDS
-
     if os.environ.get("DEBUG") and os.environ.get("DEBUG").lower() == 'true':
         langchain.verbose = True
-        set_handler(DifyStdOutCallbackHandler())
 
     if app.config.get("OPENAI_API_KEY"):
         hosted_llm_credentials.openai = HostedOpenAICredential(api_key=app.config.get("OPENAI_API_KEY"))

api/core/agent/agent_builder.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from langchain import LLMChain
 from langchain.agents import ZeroShotAgent, AgentExecutor, ConversationalAgent
-from langchain.callbacks import CallbackManager
+from langchain.callbacks.manager import CallbackManager
 from langchain.memory.chat_memory import BaseChatMemory
 
 from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
@@ -16,23 +16,20 @@ class AgentBuilder:
     def to_agent_chain(cls, tenant_id: str, tools, memory: Optional[BaseChatMemory],
                        dataset_tool_callback_handler: DatasetToolCallbackHandler,
                        agent_loop_gather_callback_handler: AgentLoopGatherCallbackHandler):
-        llm_callback_manager = CallbackManager([agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()])
         llm = LLMBuilder.to_llm(
             tenant_id=tenant_id,
             model_name=agent_loop_gather_callback_handler.model_name,
             temperature=0,
             max_tokens=1024,
-            callback_manager=llm_callback_manager
+            callbacks=[agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()]
         )
 
-        tool_callback_manager = CallbackManager([
-            agent_loop_gather_callback_handler,
-            dataset_tool_callback_handler,
-            DifyStdOutCallbackHandler()
-        ])
-
         for tool in tools:
-            tool.callback_manager = tool_callback_manager
+            tool.callbacks = [
+                agent_loop_gather_callback_handler,
+                dataset_tool_callback_handler,
+                DifyStdOutCallbackHandler()
+            ]
 
         prompt = cls.build_agent_prompt_template(
             tools=tools,
@@ -54,7 +51,7 @@ class AgentBuilder:
             tools=tools,
             agent=agent,
             memory=memory,
-            callback_manager=agent_callback_manager,
+            callbacks=agent_callback_manager,
             max_iterations=6,
             early_stopping_method="generate",
             # `generate` will continue to complete the last inference after reaching the iteration limit or request time limit

api/core/callback_handler/agent_loop_gather_callback_handler.py

@@ -12,6 +12,7 @@ from core.conversation_message_task import ConversationMessageTask
 
 class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
     """Callback Handler that prints to std out."""
+    raise_error: bool = True
 
     def __init__(self, model_name, conversation_message_task: ConversationMessageTask) -> None:
         """Initialize callback handler."""
@@ -64,10 +65,6 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             self._current_loop.completion = response.generations[0][0].text
             self._current_loop.completion_tokens = response.llm_output['token_usage']['completion_tokens']
 
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Do nothing."""
-        pass
-
     def on_llm_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
@@ -75,21 +72,6 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
         self._agent_loops = []
         self._current_loop = None
 
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        """Print out that we are entering a chain."""
-        pass
-
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        """Print out that we finished a chain."""
-        pass
-
-    def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        logging.error(error)
-
     def on_tool_start(
         self,
         serialized: Dict[str, Any],
@@ -151,16 +133,6 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
         self._agent_loops = []
         self._current_loop = None
 
-    def on_text(
-        self,
-        text: str,
-        color: Optional[str] = None,
-        end: str = "",
-        **kwargs: Optional[str],
-    ) -> None:
-        """Run on additional input from chains and agents."""
-        pass
-
     def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
         """Run on agent end."""
         # Final Answer

api/core/callback_handler/dataset_tool_callback_handler.py

@@ -3,7 +3,6 @@ import logging
 from typing import Any, Dict, List, Union, Optional
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult
 
 from core.callback_handler.entity.dataset_query import DatasetQueryObj
 from core.conversation_message_task import ConversationMessageTask
@@ -11,6 +10,7 @@ from core.conversation_message_task import ConversationMessageTask
 
 class DatasetToolCallbackHandler(BaseCallbackHandler):
     """Callback Handler that prints to std out."""
+    raise_error: bool = True
 
     def __init__(self, conversation_message_task: ConversationMessageTask) -> None:
         """Initialize callback handler."""
@@ -66,52 +66,3 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
     ) -> None:
         """Do nothing."""
         logging.error(error)
-
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        pass
-
-    def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        pass
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Do nothing."""
-        pass
-
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        logging.error(error)
-
-    def on_agent_action(
-        self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
-    ) -> Any:
-        pass
-
-    def on_text(
-        self,
-        text: str,
-        color: Optional[str] = None,
-        end: str = "",
-        **kwargs: Optional[str],
-    ) -> None:
-        """Run on additional input from chains and agents."""
-        pass
-
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
-        """Run on agent end."""
-        pass

api/core/callback_handler/index_tool_callback_handler.py

@@ -1,39 +1,26 @@
-from llama_index import Response
+from typing import List
+
+from langchain.schema import Document
 
 from extensions.ext_database import db
 from models.dataset import DocumentSegment
 
 
-class IndexToolCallbackHandler:
-
-    def __init__(self) -> None:
-        self._response = None
-
-    @property
-    def response(self) -> Response:
-        return self._response
-
-    def on_tool_end(self, response: Response) -> None:
-        """Handle tool end."""
-        self._response = response
-
-
-class DatasetIndexToolCallbackHandler(IndexToolCallbackHandler):
+class DatasetIndexToolCallbackHandler:
     """Callback handler for dataset tool."""
 
     def __init__(self, dataset_id: str) -> None:
-        super().__init__()
         self.dataset_id = dataset_id
 
-    def on_tool_end(self, response: Response) -> None:
+    def on_tool_end(self, documents: List[Document]) -> None:
         """Handle tool end."""
-        for node in response.source_nodes:
-            index_node_id = node.node.doc_id
+        for document in documents:
+            doc_id = document.metadata['doc_id']
 
             # add hit count to document segment
             db.session.query(DocumentSegment).filter(
                 DocumentSegment.dataset_id == self.dataset_id,
-                DocumentSegment.index_node_id == index_node_id
+                DocumentSegment.index_node_id == doc_id
             ).update(
                 {DocumentSegment.hit_count: DocumentSegment.hit_count + 1},
                 synchronize_session=False

api/core/callback_handler/llm_callback_handler.py

@@ -3,7 +3,7 @@ import time
 from typing import Any, Dict, List, Union, Optional
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult, HumanMessage, AIMessage, SystemMessage
+from langchain.schema import AgentAction, AgentFinish, LLMResult, HumanMessage, AIMessage, SystemMessage, BaseMessage
 
 from core.callback_handler.entity.llm_message import LLMMessage
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
@@ -12,6 +12,7 @@ from core.llm.streamable_open_ai import StreamableOpenAI
 
 
 class LLMCallbackHandler(BaseCallbackHandler):
+    raise_error: bool = True
 
     def __init__(self, llm: Union[StreamableOpenAI, StreamableChatOpenAI],
                  conversation_message_task: ConversationMessageTask):
@@ -25,41 +26,41 @@ class LLMCallbackHandler(BaseCallbackHandler):
         """Whether to call verbose callbacks even if verbose is False."""
         return True
 
+    def on_chat_model_start(
+            self,
+            serialized: Dict[str, Any],
+            messages: List[List[BaseMessage]],
+            **kwargs: Any
+    ) -> Any:
+        self.start_at = time.perf_counter()
+        real_prompts = []
+        for message in messages[0]:
+            if message.type == 'human':
+                role = 'user'
+            elif message.type == 'ai':
+                role = 'assistant'
+            else:
+                role = 'system'
+
+            real_prompts.append({
+                "role": role,
+                "text": message.content
+            })
+
+        self.llm_message.prompt = real_prompts
+        self.llm_message.prompt_tokens = self.llm.get_messages_tokens(messages[0])
+
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
         self.start_at = time.perf_counter()
 
-        if 'Chat' in serialized['name']:
-            real_prompts = []
-            messages = []
-            for prompt in prompts:
-                role, content = prompt.split(': ', maxsplit=1)
-                if role == 'human':
-                    role = 'user'
-                    message = HumanMessage(content=content)
-                elif role == 'ai':
-                    role = 'assistant'
-                    message = AIMessage(content=content)
-                else:
-                    message = SystemMessage(content=content)
+        self.llm_message.prompt = [{
+            "role": 'user',
+            "text": prompts[0]
+        }]
 
-                real_prompt = {
-                    "role": role,
-                    "text": content
-                }
-                real_prompts.append(real_prompt)
-                messages.append(message)
-
-            self.llm_message.prompt = real_prompts
-            self.llm_message.prompt_tokens = self.llm.get_messages_tokens(messages)
-        else:
-            self.llm_message.prompt = [{
-                "role": 'user',
-                "text": prompts[0]
-            }]
-
-            self.llm_message.prompt_tokens = self.llm.get_num_tokens(prompts[0])
+        self.llm_message.prompt_tokens = self.llm.get_num_tokens(prompts[0])
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         end_at = time.perf_counter()
@@ -95,58 +96,3 @@ class LLMCallbackHandler(BaseCallbackHandler):
                 self.conversation_message_task.save_message(llm_message=self.llm_message, by_stopped=True)
         else:
             logging.error(error)
-
-    def on_chain_start(
-            self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        pass
-
-    def on_chain_error(
-            self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_tool_start(
-            self,
-            serialized: Dict[str, Any],
-            input_str: str,
-            **kwargs: Any,
-    ) -> None:
-        pass
-
-    def on_agent_action(
-            self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
-    ) -> Any:
-        pass
-
-    def on_tool_end(
-            self,
-            output: str,
-            color: Optional[str] = None,
-            observation_prefix: Optional[str] = None,
-            llm_prefix: Optional[str] = None,
-            **kwargs: Any,
-    ) -> None:
-        pass
-
-    def on_tool_error(
-            self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_text(
-            self,
-            text: str,
-            color: Optional[str] = None,
-            end: str = "",
-            **kwargs: Optional[str],
-    ) -> None:
-        pass
-
-    def on_agent_finish(
-            self, finish: AgentFinish, color: Optional[str] = None, **kwargs: Any
-    ) -> None:
-        pass

api/core/callback_handler/main_chain_gather_callback_handler.py

@@ -1,10 +1,9 @@
 import logging
 import time
 
-from typing import Any, Dict, List, Union, Optional
+from typing import Any, Dict, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult
 
 from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
 from core.callback_handler.entity.chain_result import ChainResult
@@ -14,6 +13,7 @@ from core.conversation_message_task import ConversationMessageTask
 
 class MainChainGatherCallbackHandler(BaseCallbackHandler):
     """Callback Handler that prints to std out."""
+    raise_error: bool = True
 
     def __init__(self, conversation_message_task: ConversationMessageTask) -> None:
         """Initialize callback handler."""
@@ -50,13 +50,15 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
     ) -> None:
         """Print out that we are entering a chain."""
         if not self._current_chain_result:
-            self._current_chain_result = ChainResult(
-                type=serialized['name'],
-                prompt=inputs,
-                started_at=time.perf_counter()
-            )
-            self._current_chain_message = self.conversation_message_task.init_chain(self._current_chain_result)
-            self.agent_loop_gather_callback_handler.current_chain = self._current_chain_message
+            chain_type = serialized['id'][-1]
+            if chain_type:
+                self._current_chain_result = ChainResult(
+                    type=chain_type,
+                    prompt=inputs,
+                    started_at=time.perf_counter()
+                )
+                self._current_chain_message = self.conversation_message_task.init_chain(self._current_chain_result)
+                self.agent_loop_gather_callback_handler.current_chain = self._current_chain_message
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
@@ -74,64 +76,4 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
         logging.error(error)
-        self.clear_chain_results()
-
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        pass
-
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        pass
-
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Do nothing."""
-        pass
-
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        logging.error(error)
-
-    def on_tool_start(
-        self,
-        serialized: Dict[str, Any],
-        input_str: str,
-        **kwargs: Any,
-    ) -> None:
-        pass
-
-    def on_agent_action(
-        self, action: AgentAction, color: Optional[str] = None, **kwargs: Any
-    ) -> Any:
-        pass
-
-    def on_tool_end(
-        self,
-        output: str,
-        color: Optional[str] = None,
-        observation_prefix: Optional[str] = None,
-        llm_prefix: Optional[str] = None,
-        **kwargs: Any,
-    ) -> None:
-        pass
-
-    def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Do nothing."""
-        logging.error(error)
-
-    def on_text(
-        self,
-        text: str,
-        color: Optional[str] = None,
-        end: str = "",
-        **kwargs: Optional[str],
-    ) -> None:
-        """Run on additional input from chains and agents."""
-        pass
-
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
-        """Run on agent end."""
-        pass
+        self.clear_chain_results()

api/core/callback_handler/std_out_callback_handler.py

@@ -1,9 +1,10 @@
+import os
 import sys
 from typing import Any, Dict, List, Optional, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.input import print_text
-from langchain.schema import AgentAction, AgentFinish, LLMResult
+from langchain.schema import AgentAction, AgentFinish, LLMResult, BaseMessage
 
 
 class DifyStdOutCallbackHandler(BaseCallbackHandler):
@@ -13,17 +14,23 @@ class DifyStdOutCallbackHandler(BaseCallbackHandler):
         """Initialize callback handler."""
         self.color = color
 
+    def on_chat_model_start(
+            self,
+            serialized: Dict[str, Any],
+            messages: List[List[BaseMessage]],
+            **kwargs: Any
+    ) -> Any:
+        print_text("\n[on_chat_model_start]\n", color='blue')
+        for sub_messages in messages:
+            for sub_message in sub_messages:
+                print_text(str(sub_message) + "\n", color='blue')
+
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
         """Print out the prompts."""
         print_text("\n[on_llm_start]\n", color='blue')
-
-        if 'Chat' in serialized['name']:
-            for prompt in prompts:
-                print_text(prompt + "\n", color='blue')
-        else:
-            print_text(prompts[0] + "\n", color='blue')
+        print_text(prompts[0] + "\n", color='blue')
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Do nothing."""
@@ -44,8 +51,8 @@ class DifyStdOutCallbackHandler(BaseCallbackHandler):
         self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
     ) -> None:
         """Print out that we are entering a chain."""
-        class_name = serialized["name"]
-        print_text("\n[on_chain_start]\nChain: " + class_name + "\nInputs: " + str(inputs) + "\n", color='pink')
+        chain_type = serialized['id'][-1]
+        print_text("\n[on_chain_start]\nChain: " + chain_type + "\nInputs: " + str(inputs) + "\n", color='pink')
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
@@ -117,6 +124,26 @@ class DifyStdOutCallbackHandler(BaseCallbackHandler):
         """Run on agent end."""
         print_text("[on_agent_finish] " + finish.return_values['output'] + "\n", color='green', end="\n")
 
+    @property
+    def ignore_llm(self) -> bool:
+        """Whether to ignore LLM callbacks."""
+        return not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true'
+
+    @property
+    def ignore_chain(self) -> bool:
+        """Whether to ignore chain callbacks."""
+        return not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true'
+
+    @property
+    def ignore_agent(self) -> bool:
+        """Whether to ignore agent callbacks."""
+        return not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true'
+
+    @property
+    def ignore_chat_model(self) -> bool:
+        """Whether to ignore chat model callbacks."""
+        return not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true'
+
 
 class DifyStreamingStdOutCallbackHandler(DifyStdOutCallbackHandler):
     """Callback handler for streaming. Only works with LLMs that support streaming."""

api/core/chain/chain_builder.py

@@ -1,7 +1,5 @@
 from typing import Optional
 
-from langchain.callbacks import CallbackManager
-
 from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
 from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceChain
 from core.chain.tool_chain import ToolChain
@@ -14,7 +12,7 @@ class ChainBuilder:
             tool=tool,
             input_key=kwargs.get('input_key', 'input'),
             output_key=kwargs.get('output_key', 'tool_output'),
-            callback_manager=CallbackManager([DifyStdOutCallbackHandler()])
+            callbacks=[DifyStdOutCallbackHandler()]
         )
 
     @classmethod
@@ -27,7 +25,7 @@ class ChainBuilder:
                 sensitive_words=sensitive_words.split(","),
                 canned_response=tool_config.get("canned_response", ''),
                 output_key="sensitive_word_avoidance_output",
-                callback_manager=CallbackManager([DifyStdOutCallbackHandler()]),
+                callbacks=[DifyStdOutCallbackHandler()],
                 **kwargs
             )
 

api/core/chain/llm_router_chain.py

@@ -1,15 +1,16 @@
 """Base classes for LLM-powered router chains."""
 from __future__ import annotations
 
-import json
 from typing import Any, Dict, List, Optional, Type, cast, NamedTuple
 
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 from pydantic import root_validator
 
 from langchain.chains import LLMChain
 from langchain.prompts import BasePromptTemplate
-from langchain.schema import BaseOutputParser, OutputParserException, BaseLanguageModel
+from langchain.schema import BaseOutputParser, OutputParserException
 
 from libs.json_in_md_parser import parse_and_check_json_markdown
 
@@ -51,8 +52,9 @@ class LLMRouterChain(Chain):
             raise ValueError
 
     def _call(
-        self,
-        inputs: Dict[str, Any]
+            self,
+            inputs: Dict[str, Any],
+            run_manager: Optional[CallbackManagerForChainRun] = None,
     ) -> Dict[str, Any]:
         output = cast(
             Dict[str, Any],

api/core/chain/main_chain_builder.py

@@ -1,11 +1,9 @@
-from typing import Optional, List
+from typing import Optional, List, cast
 
-from langchain.callbacks import SharedCallbackManager, CallbackManager
 from langchain.chains import SequentialChain
 from langchain.chains.base import Chain
 from langchain.memory.chat_memory import BaseChatMemory
 
-from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
 from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
 from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
 from core.chain.chain_builder import ChainBuilder
@@ -18,6 +16,7 @@ from models.dataset import Dataset
 class MainChainBuilder:
     @classmethod
     def to_langchain_components(cls, tenant_id: str, agent_mode: dict, memory: Optional[BaseChatMemory],
+                                rest_tokens: int,
                                 conversation_message_task: ConversationMessageTask):
         first_input_key = "input"
         final_output_key = "output"
@@ -30,6 +29,7 @@ class MainChainBuilder:
         tool_chains, chains_output_key = cls.get_agent_chains(
             tenant_id=tenant_id,
             agent_mode=agent_mode,
+            rest_tokens=rest_tokens,
             memory=memory,
             conversation_message_task=conversation_message_task
         )
@@ -42,9 +42,8 @@ class MainChainBuilder:
             return None
 
         for chain in chains:
-            # do not add handler into singleton callback manager
-            if not isinstance(chain.callback_manager, SharedCallbackManager):
-                chain.callback_manager.add_handler(chain_callback_handler)
+            chain = cast(Chain, chain)
+            chain.callbacks.append(chain_callback_handler)
 
         # build main chain
         overall_chain = SequentialChain(
@@ -57,7 +56,9 @@ class MainChainBuilder:
         return overall_chain
 
     @classmethod
-    def get_agent_chains(cls, tenant_id: str, agent_mode: dict, memory: Optional[BaseChatMemory],
+    def get_agent_chains(cls, tenant_id: str, agent_mode: dict,
+                         rest_tokens: int,
+                         memory: Optional[BaseChatMemory],
                          conversation_message_task: ConversationMessageTask):
         # agent mode
         chains = []
@@ -93,7 +94,8 @@ class MainChainBuilder:
                     tenant_id=tenant_id,
                     datasets=datasets,
                     conversation_message_task=conversation_message_task,
-                    callback_manager=CallbackManager([DifyStdOutCallbackHandler()])
+                    rest_tokens=rest_tokens,
+                    callbacks=[DifyStdOutCallbackHandler()]
                 )
                 chains.append(multi_dataset_router_chain)
 

api/core/chain/multi_dataset_router_chain.py

@@ -1,9 +1,10 @@
+import math
+import re
 from typing import Mapping, List, Dict, Any, Optional
 
-from langchain import LLMChain, PromptTemplate, ConversationChain
-from langchain.callbacks import CallbackManager
+from langchain import PromptTemplate
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
-from langchain.schema import BaseLanguageModel
 from pydantic import Extra
 
 from core.callback_handler.dataset_tool_callback_handler import DatasetToolCallbackHandler
@@ -11,10 +12,11 @@ from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHan
 from core.chain.llm_router_chain import LLMRouterChain, RouterOutputParser
 from core.conversation_message_task import ConversationMessageTask
 from core.llm.llm_builder import LLMBuilder
-from core.tool.dataset_tool_builder import DatasetToolBuilder
-from core.tool.llama_index_tool import EnhanceLlamaIndexTool
-from models.dataset import Dataset
+from core.tool.dataset_index_tool import DatasetTool
+from models.dataset import Dataset, DatasetProcessRule
 
+DEFAULT_K = 2
+CONTEXT_TOKENS_PERCENT = 0.3
 MULTI_PROMPT_ROUTER_TEMPLATE = """
 Given a raw text input to a language model select the model prompt best suited for \
 the input. You will be given the names of the available prompts and a description of \
@@ -52,7 +54,7 @@ class MultiDatasetRouterChain(Chain):
 
     router_chain: LLMRouterChain
     """Chain for deciding a destination chain and the input to it."""
-    dataset_tools: Mapping[str, EnhanceLlamaIndexTool]
+    dataset_tools: Mapping[str, DatasetTool]
     """Map of name to candidate chains that inputs can be routed to."""
 
     class Config:
@@ -79,41 +81,56 @@ class MultiDatasetRouterChain(Chain):
             tenant_id: str,
             datasets: List[Dataset],
             conversation_message_task: ConversationMessageTask,
+            rest_tokens: int,
             **kwargs: Any,
     ):
         """Convenience constructor for instantiating from destination prompts."""
-        llm_callback_manager = CallbackManager([DifyStdOutCallbackHandler()])
         llm = LLMBuilder.to_llm(
             tenant_id=tenant_id,
             model_name='gpt-3.5-turbo',
             temperature=0,
             max_tokens=1024,
-            callback_manager=llm_callback_manager
+            callbacks=[DifyStdOutCallbackHandler()]
         )
 
-        destinations = ["{}: {}".format(d.id, d.description.replace('\n', ' ') if d.description
+        destinations = ["[[{}]]: {}".format(d.id, d.description.replace('\n', ' ') if d.description
                         else ('useful for when you want to answer queries about the ' + d.name))
                         for d in datasets]
         destinations_str = "\n".join(destinations)
         router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
             destinations=destinations_str
         )
+
         router_prompt = PromptTemplate(
             template=router_template,
             input_variables=["input"],
             output_parser=RouterOutputParser(),
         )
+
         router_chain = LLMRouterChain.from_llm(llm, router_prompt)
         dataset_tools = {}
         for dataset in datasets:
-            dataset_tool = DatasetToolBuilder.build_dataset_tool(
+            # fulfill description when it is empty
+            if dataset.available_document_count == 0 or dataset.available_document_count == 0:
+                continue
+
+            description = dataset.description
+            if not description:
+                description = 'useful for when you want to answer queries about the ' + dataset.name
+
+            k = cls._dynamic_calc_retrieve_k(dataset, rest_tokens)
+            if k == 0:
+                continue
+
+            dataset_tool = DatasetTool(
+                name=f"dataset-{dataset.id}",
+                description=description,
+                k=k,
                 dataset=dataset,
-                response_mode='no_synthesizer',  # "compact"
-                callback_handler=DatasetToolCallbackHandler(conversation_message_task)
+                callbacks=[DatasetToolCallbackHandler(conversation_message_task), DifyStdOutCallbackHandler()]
             )
 
-            if dataset_tool:
-                dataset_tools[dataset.id] = dataset_tool
+            dataset_tools[str(dataset.id)] = dataset_tool
 
         return cls(
             router_chain=router_chain,
@@ -121,9 +138,39 @@ class MultiDatasetRouterChain(Chain):
             **kwargs,
         )
 
+    @classmethod
+    def _dynamic_calc_retrieve_k(cls, dataset: Dataset, rest_tokens: int) -> int:
+        processing_rule = dataset.latest_process_rule
+        if not processing_rule:
+            return DEFAULT_K
+
+        if processing_rule.mode == "custom":
+            rules = processing_rule.rules_dict
+            if not rules:
+                return DEFAULT_K
+
+            segmentation = rules["segmentation"]
+            segment_max_tokens = segmentation["max_tokens"]
+        else:
+            segment_max_tokens = DatasetProcessRule.AUTOMATIC_RULES['segmentation']['max_tokens']
+
+        # when rest_tokens is less than default context tokens
+        if rest_tokens < segment_max_tokens * DEFAULT_K:
+            return rest_tokens // segment_max_tokens
+
+        context_limit_tokens = math.floor(rest_tokens * CONTEXT_TOKENS_PERCENT)
+
+        # when context_limit_tokens is less than default context tokens, use default_k
+        if context_limit_tokens <= segment_max_tokens * DEFAULT_K:
+            return DEFAULT_K
+
+        # Expand the k value when there's still some room left in the 30% rest tokens space
+        return context_limit_tokens // segment_max_tokens
+
     def _call(
         self,
-        inputs: Dict[str, Any]
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
     ) -> Dict[str, Any]:
         if len(self.dataset_tools) == 0:
             return {"text": ''}
@@ -132,13 +179,20 @@ class MultiDatasetRouterChain(Chain):
 
         route = self.router_chain.route(inputs)
 
-        if not route.destination:
+        destination = ''
+        if route.destination:
+            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, route.destination, re.IGNORECASE)
+            if match:
+                destination = match.group()
+
+        if not destination:
             return {"text": ''}
-        elif route.destination in self.dataset_tools:
-            return {"text": self.dataset_tools[route.destination].run(
+        elif destination in self.dataset_tools:
+            return {"text": self.dataset_tools[destination].run(
                 route.next_inputs['input']
             )}
         else:
             raise ValueError(
-                f"Received invalid destination chain name '{route.destination}'"
+                f"Received invalid destination chain name '{destination}'"
             )

api/core/chain/sensitive_word_avoidance_chain.py

@@ -1,5 +1,6 @@
-from typing import List, Dict
+from typing import List, Dict, Optional, Any
 
+from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.base import Chain
 
 
@@ -36,7 +37,11 @@ class SensitiveWordAvoidanceChain(Chain):
                 return self.canned_response
         return text
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+            self,
+            inputs: Dict[str, Any],
+            run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         text = inputs[self.input_key]
         output = self._check_sensitive_word(text)
         return {self.output_key: output}

api/core/chain/tool_chain.py

@@ -1,5 +1,6 @@
-from typing import List, Dict
+from typing import List, Dict, Optional, Any
 
+from langchain.callbacks.manager import CallbackManagerForChainRun, AsyncCallbackManagerForChainRun
 from langchain.chains.base import Chain
 from langchain.tools import BaseTool
 
@@ -30,12 +31,20 @@ class ToolChain(Chain):
         """
         return [self.output_key]
 
-    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    def _call(
+            self,
+            inputs: Dict[str, Any],
+            run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         input = inputs[self.input_key]
         output = self.tool.run(input, self.verbose)
         return {self.output_key: output}
 
-    async def _acall(self, inputs: Dict[str, str]) -> Dict[str, str]:
+    async def _acall(
+            self,
+            inputs: Dict[str, Any],
+            run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
         """Run the logic of this chain and return the output."""
         input = inputs[self.input_key]
         output = await self.tool.arun(input, self.verbose)

api/core/completion.py

@@ -1,17 +1,18 @@
 import logging
 from typing import Optional, List, Union, Tuple
 
-from langchain.callbacks import CallbackManager
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chat_models.base import BaseChatModel
 from langchain.llms import BaseLLM
-from langchain.schema import BaseMessage, BaseLanguageModel, HumanMessage
+from langchain.schema import BaseMessage, HumanMessage
 from requests.exceptions import ChunkedEncodingError
 
 from core.constant import llm_constant
 from core.callback_handler.llm_callback_handler import LLMCallbackHandler
 from core.callback_handler.std_out_callback_handler import DifyStreamingStdOutCallbackHandler, \
     DifyStdOutCallbackHandler
-from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException, PubHandler
+from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
 from core.llm.error import LLMBadRequestError
 from core.llm.llm_builder import LLMBuilder
 from core.chain.main_chain_builder import MainChainBuilder
@@ -22,7 +23,7 @@ from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
 from core.memory.read_only_conversation_token_db_string_buffer_shared_memory import \
     ReadOnlyConversationTokenDBStringBufferSharedMemory
 from core.prompt.prompt_builder import PromptBuilder
-from core.prompt.prompt_template import OutLinePromptTemplate
+from core.prompt.prompt_template import JinjaPromptTemplate
 from core.prompt.prompts import MORE_LIKE_THIS_GENERATE_PROMPT
 from models.model import App, AppModelConfig, Account, Conversation, Message
 
@@ -34,7 +35,7 @@ class Completion:
         """
         errors: ProviderTokenNotInitError
         """
-        cls.validate_query_tokens(app.tenant_id, app_model_config, query)
+        query = PromptBuilder.process_template(query)
 
         memory = None
         if conversation:
@@ -48,6 +49,14 @@ class Completion:
 
             inputs = conversation.inputs
 
+        rest_tokens_for_context_and_memory = cls.get_validate_rest_tokens(
+            mode=app.mode,
+            tenant_id=app.tenant_id,
+            app_model_config=app_model_config,
+            query=query,
+            inputs=inputs
+        )
+
         conversation_message_task = ConversationMessageTask(
             task_id=task_id,
             app=app,
@@ -64,6 +73,7 @@ class Completion:
         main_chain = MainChainBuilder.to_langchain_components(
             tenant_id=app.tenant_id,
             agent_mode=app_model_config.agent_mode_dict,
+            rest_tokens=rest_tokens_for_context_and_memory,
             memory=ReadOnlyConversationTokenDBStringBufferSharedMemory(memory=memory) if memory else None,
             conversation_message_task=conversation_message_task
         )
@@ -115,7 +125,7 @@ class Completion:
             memory=memory
         )
 
-        final_llm.callback_manager = cls.get_llm_callback_manager(final_llm, streaming, conversation_message_task)
+        final_llm.callbacks = cls.get_llm_callbacks(final_llm, streaming, conversation_message_task)
 
         cls.recale_llm_max_tokens(
             final_llm=final_llm,
@@ -133,18 +143,17 @@ class Completion:
                             memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]) -> \
             Tuple[Union[str | List[BaseMessage]], Optional[List[str]]]:
         # disable template string in query
-        query_params = OutLinePromptTemplate.from_template(template=query).input_variables
-        if query_params:
-            for query_param in query_params:
-                if query_param not in inputs:
-                    inputs[query_param] = '{' + query_param + '}'
+        # query_params = JinjaPromptTemplate.from_template(template=query).input_variables
+        # if query_params:
+        #     for query_param in query_params:
+        #         if query_param not in inputs:
+        #             inputs[query_param] = '{{' + query_param + '}}'
 
-        pre_prompt = PromptBuilder.process_template(pre_prompt) if pre_prompt else pre_prompt
         if mode == 'completion':
-            prompt_template = OutLinePromptTemplate.from_template(
+            prompt_template = JinjaPromptTemplate.from_template(
                 template=("""Use the following CONTEXT as your learned knowledge:
 [CONTEXT]
-{context}
+{{context}}
 [END CONTEXT]
 
 When answer to user:
@@ -154,16 +163,16 @@ Avoid mentioning that you obtained the information from the context.
 And answer according to the language of the user's question.
 """ if chain_output else "")
                          + (pre_prompt + "\n" if pre_prompt else "")
-                         + "{query}\n"
+                         + "{{query}}\n"
             )
 
             if chain_output:
                 inputs['context'] = chain_output
-                context_params = OutLinePromptTemplate.from_template(template=chain_output).input_variables
-                if context_params:
-                    for context_param in context_params:
-                        if context_param not in inputs:
-                            inputs[context_param] = '{' + context_param + '}'
+                # context_params = JinjaPromptTemplate.from_template(template=chain_output).input_variables
+                # if context_params:
+                #     for context_param in context_params:
+                #         if context_param not in inputs:
+                #             inputs[context_param] = '{{' + context_param + '}}'
 
             prompt_inputs = {k: inputs[k] for k in prompt_template.input_variables if k in inputs}
             prompt_content = prompt_template.format(
@@ -187,7 +196,7 @@ And answer according to the language of the user's question.
 
             if pre_prompt:
                 pre_prompt_inputs = {k: inputs[k] for k in
-                                     OutLinePromptTemplate.from_template(template=pre_prompt).input_variables
+                                     JinjaPromptTemplate.from_template(template=pre_prompt).input_variables
                                      if k in inputs}
 
                 if pre_prompt_inputs:
@@ -197,7 +206,7 @@ And answer according to the language of the user's question.
                 human_inputs['context'] = chain_output
                 human_message_prompt += """Use the following CONTEXT as your learned knowledge.
 [CONTEXT]
-{context}
+{{context}}
 [END CONTEXT]
 
 When answer to user:
@@ -210,7 +219,7 @@ And answer according to the language of the user's question.
             if pre_prompt:
                 human_message_prompt += pre_prompt
 
-            query_prompt = "\nHuman: {query}\nAI: "
+            query_prompt = "\nHuman: {{query}}\nAI: "
 
             if memory:
                 # append chat histories
@@ -226,11 +235,11 @@ And answer according to the language of the user's question.
                 histories = cls.get_history_messages_from_memory(memory, rest_tokens)
 
                 # disable template string in query
-                histories_params = OutLinePromptTemplate.from_template(template=histories).input_variables
-                if histories_params:
-                    for histories_param in histories_params:
-                        if histories_param not in human_inputs:
-                            human_inputs[histories_param] = '{' + histories_param + '}'
+                # histories_params = JinjaPromptTemplate.from_template(template=histories).input_variables
+                # if histories_params:
+                #     for histories_param in histories_params:
+                #         if histories_param not in human_inputs:
+                #             human_inputs[histories_param] = '{{' + histories_param + '}}'
 
                 human_message_prompt += "\n\n" + histories
 
@@ -247,16 +256,14 @@ And answer according to the language of the user's question.
             return messages, ['\nHuman:']
 
     @classmethod
-    def get_llm_callback_manager(cls, llm: Union[StreamableOpenAI, StreamableChatOpenAI],
-                                 streaming: bool,
-                                 conversation_message_task: ConversationMessageTask) -> CallbackManager:
+    def get_llm_callbacks(cls, llm: Union[StreamableOpenAI, StreamableChatOpenAI],
+                          streaming: bool,
+                          conversation_message_task: ConversationMessageTask) -> List[BaseCallbackHandler]:
         llm_callback_handler = LLMCallbackHandler(llm, conversation_message_task)
         if streaming:
-            callback_handlers = [llm_callback_handler, DifyStreamingStdOutCallbackHandler()]
+            return [llm_callback_handler, DifyStreamingStdOutCallbackHandler()]
         else:
-            callback_handlers = [llm_callback_handler, DifyStdOutCallbackHandler()]
-
-        return CallbackManager(callback_handlers)
+            return [llm_callback_handler, DifyStdOutCallbackHandler()]
 
     @classmethod
     def get_history_messages_from_memory(cls, memory: ReadOnlyConversationTokenDBBufferSharedMemory,
@@ -293,7 +300,8 @@ And answer according to the language of the user's question.
         return memory
 
     @classmethod
-    def validate_query_tokens(cls, tenant_id: str, app_model_config: AppModelConfig, query: str):
+    def get_validate_rest_tokens(cls, mode: str, tenant_id: str, app_model_config: AppModelConfig,
+                                 query: str, inputs: dict) -> int:
         llm = LLMBuilder.to_llm_from_model(
             tenant_id=tenant_id,
             model=app_model_config.model_dict
@@ -302,8 +310,26 @@ And answer according to the language of the user's question.
         model_limited_tokens = llm_constant.max_context_token_length[llm.model_name]
         max_tokens = llm.max_tokens
 
-        if model_limited_tokens - max_tokens - llm.get_num_tokens(query) < 0:
-            raise LLMBadRequestError("Query is too long")
+        # get prompt without memory and context
+        prompt, _ = cls.get_main_llm_prompt(
+            mode=mode,
+            llm=llm,
+            pre_prompt=app_model_config.pre_prompt,
+            query=query,
+            inputs=inputs,
+            chain_output=None,
+            memory=None
+        )
+
+        prompt_tokens = llm.get_num_tokens(prompt) if isinstance(prompt, str) \
+            else llm.get_num_tokens_from_messages(prompt)
+
+        rest_tokens = model_limited_tokens - max_tokens - prompt_tokens
+        if rest_tokens < 0:
+            raise LLMBadRequestError("Query or prefix prompt is too long, you can reduce the prefix prompt, "
+                                     "or shrink the max token, or switch to a llm with a larger token limit size.")
+
+        return rest_tokens
 
     @classmethod
     def recale_llm_max_tokens(cls, final_llm: Union[StreamableOpenAI, StreamableChatOpenAI],
@@ -360,7 +386,7 @@ And answer according to the language of the user's question.
             streaming=streaming
         )
 
-        llm.callback_manager = cls.get_llm_callback_manager(llm, streaming, conversation_message_task)
+        llm.callbacks = cls.get_llm_callbacks(llm, streaming, conversation_message_task)
 
         cls.recale_llm_max_tokens(
             final_llm=llm,

api/core/conversation_message_task.py

@@ -10,7 +10,7 @@ from core.constant import llm_constant
 from core.llm.llm_builder import LLMBuilder
 from core.llm.provider.llm_provider_service import LLMProviderService
 from core.prompt.prompt_builder import PromptBuilder
-from core.prompt.prompt_template import OutLinePromptTemplate
+from core.prompt.prompt_template import JinjaPromptTemplate
 from events.message_event import message_was_created
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
@@ -78,7 +78,7 @@ class ConversationMessageTask:
         if self.mode == 'chat':
             introduction = self.app_model_config.opening_statement
             if introduction:
-                prompt_template = OutLinePromptTemplate.from_template(template=PromptBuilder.process_template(introduction))
+                prompt_template = JinjaPromptTemplate.from_template(template=introduction)
                 prompt_inputs = {k: self.inputs[k] for k in prompt_template.input_variables if k in self.inputs}
                 try:
                     introduction = prompt_template.format(**prompt_inputs)
@@ -86,8 +86,7 @@ class ConversationMessageTask:
                     pass
 
             if self.app_model_config.pre_prompt:
-                pre_prompt = PromptBuilder.process_template(self.app_model_config.pre_prompt)
-                system_message = PromptBuilder.to_system_message(pre_prompt, self.inputs)
+                system_message = PromptBuilder.to_system_message(self.app_model_config.pre_prompt, self.inputs)
                 system_instruction = system_message.content
                 llm = LLMBuilder.to_llm(self.tenant_id, self.model_name)
                 system_instruction_tokens = llm.get_messages_tokens([system_message])
@@ -157,7 +156,7 @@ class ConversationMessageTask:
         self.message.message = llm_message.prompt
         self.message.message_tokens = message_tokens
         self.message.message_unit_price = message_unit_price
-        self.message.answer = llm_message.completion.strip() if llm_message.completion else ''
+        self.message.answer = PromptBuilder.process_template(llm_message.completion.strip()) if llm_message.completion else ''
         self.message.answer_tokens = answer_tokens
         self.message.answer_unit_price = answer_unit_price
         self.message.provider_response_latency = llm_message.latency
@@ -293,12 +292,12 @@ class PubHandler:
         if not user:
             raise ValueError("user is required")
 
-        user_str = 'account-' + user.id if isinstance(user, Account) else 'end-user-' + user.id
+        user_str = 'account-' + str(user.id) if isinstance(user, Account) else 'end-user-' + str(user.id)
         return "generate_result:{}-{}".format(user_str, task_id)
 
     @classmethod
     def generate_stopped_cache_key(cls, user: Union[Account | EndUser], task_id: str):
-        user_str = 'account-' + user.id if isinstance(user, Account) else 'end-user-' + user.id
+        user_str = 'account-' + str(user.id) if isinstance(user, Account) else 'end-user-' + str(user.id)
         return "generate_result_stopped:{}-{}".format(user_str, task_id)
 
     def pub_text(self, text: str):
@@ -306,10 +305,10 @@ class PubHandler:
             'event': 'message',
             'data': {
                 'task_id': self._task_id,
-                'message_id': self._message.id,
+                'message_id': str(self._message.id),
                 'text': text,
                 'mode': self._conversation.mode,
-                'conversation_id': self._conversation.id
+                'conversation_id': str(self._conversation.id)
             }
         }
 

api/core/data_loader/file_extractor.py

@@ -0,0 +1,43 @@
+import tempfile
+from pathlib import Path
+from typing import List, Union
+
+from langchain.document_loaders import TextLoader, Docx2txtLoader
+from langchain.schema import Document
+
+from core.data_loader.loader.csv import CSVLoader
+from core.data_loader.loader.excel import ExcelLoader
+from core.data_loader.loader.html import HTMLLoader
+from core.data_loader.loader.markdown import MarkdownLoader
+from core.data_loader.loader.pdf import PdfLoader
+from extensions.ext_storage import storage
+from models.model import UploadFile
+
+
+class FileExtractor:
+    @classmethod
+    def load(cls, upload_file: UploadFile, return_text: bool = False) -> Union[List[Document] | str]:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            suffix = Path(upload_file.key).suffix
+            file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
+            storage.download(upload_file.key, file_path)
+
+            input_file = Path(file_path)
+            delimiter = '\n'
+            if input_file.suffix == '.xlsx':
+                loader = ExcelLoader(file_path)
+            elif input_file.suffix == '.pdf':
+                loader = PdfLoader(file_path, upload_file=upload_file)
+            elif input_file.suffix in ['.md', '.markdown']:
+                loader = MarkdownLoader(file_path, autodetect_encoding=True)
+            elif input_file.suffix in ['.htm', '.html']:
+                loader = HTMLLoader(file_path)
+            elif input_file.suffix == '.docx':
+                loader = Docx2txtLoader(file_path)
+            elif input_file.suffix == '.csv':
+                loader = CSVLoader(file_path, autodetect_encoding=True)
+            else:
+                # txt
+                loader = TextLoader(file_path, autodetect_encoding=True)
+
+            return delimiter.join([document.page_content for document in loader.load()]) if return_text else loader.load()

api/core/data_loader/loader/csv.py

@@ -0,0 +1,67 @@
+import logging
+from typing import Optional, Dict, List
+
+from langchain.document_loaders import CSVLoader as LCCSVLoader
+from langchain.document_loaders.helpers import detect_file_encodings
+
+from models.dataset import Document
+
+logger = logging.getLogger(__name__)
+
+
+class CSVLoader(LCCSVLoader):
+    def __init__(
+            self,
+            file_path: str,
+            source_column: Optional[str] = None,
+            csv_args: Optional[Dict] = None,
+            encoding: Optional[str] = None,
+            autodetect_encoding: bool = True,
+    ):
+        self.file_path = file_path
+        self.source_column = source_column
+        self.encoding = encoding
+        self.csv_args = csv_args or {}
+        self.autodetect_encoding = autodetect_encoding
+
+    def load(self) -> List[Document]:
+        """Load data into document objects."""
+        try:
+            with open(self.file_path, newline="", encoding=self.encoding) as csvfile:
+                docs = self._read_from_file(csvfile)
+        except UnicodeDecodeError as e:
+            if self.autodetect_encoding:
+                detected_encodings = detect_file_encodings(self.file_path)
+                for encoding in detected_encodings:
+                    logger.debug("Trying encoding: ", encoding.encoding)
+                    try:
+                        with open(self.file_path, newline="", encoding=encoding.encoding) as csvfile:
+                            docs = self._read_from_file(csvfile)
+                        break
+                    except UnicodeDecodeError:
+                        continue
+            else:
+                raise RuntimeError(f"Error loading {self.file_path}") from e
+
+        return docs
+
+    def _read_from_file(self, csvfile):
+        docs = []
+        csv_reader = csv.DictReader(csvfile, **self.csv_args)  # type: ignore
+        for i, row in enumerate(csv_reader):
+            content = "\n".join(f"{k.strip()}: {v.strip()}" for k, v in row.items())
+            try:
+                source = (
+                    row[self.source_column]
+                    if self.source_column is not None
+                    else ''
+                )
+            except KeyError:
+                raise ValueError(
+                    f"Source column '{self.source_column}' not found in CSV file."
+                )
+            metadata = {"source": source, "row": i}
+            doc = Document(page_content=content, metadata=metadata)
+            docs.append(doc)
+
+        return docs

api/core/data_loader/loader/excel.py

@@ -0,0 +1,45 @@
+import json
+import logging
+from typing import List
+
+from langchain.document_loaders.base import BaseLoader
+from langchain.schema import Document
+from openpyxl.reader.excel import load_workbook
+
+logger = logging.getLogger(__name__)
+
+
+class ExcelLoader(BaseLoader):
+    """Load xlxs files.
+
+
+    Args:
+        file_path: Path to the file to load.
+    """
+
+    def __init__(
+        self,
+        file_path: str
+    ):
+        """Initialize with file path."""
+        self._file_path = file_path
+
+    def load(self) -> List[Document]:
+        data = []
+        keys = []
+        wb = load_workbook(filename=self._file_path, read_only=True)
+        # loop over all sheets
+        for sheet in wb:
+            for row in sheet.iter_rows(values_only=True):
+                if all(v is None for v in row):
+                    continue
+                if keys == []:
+                    keys = list(map(str, row))
+                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())
+                    document = Document(page_content=item)
+                    data.append(document)
+
+        return data

api/core/data_loader/loader/html.py

@@ -0,0 +1,35 @@
+import logging
+from typing import List
+
+from bs4 import BeautifulSoup
+from langchain.document_loaders.base import BaseLoader
+from langchain.schema import Document
+
+logger = logging.getLogger(__name__)
+
+
+class HTMLLoader(BaseLoader):
+    """Load html files.
+
+
+    Args:
+        file_path: Path to the file to load.
+    """
+
+    def __init__(
+        self,
+        file_path: str
+    ):
+        """Initialize with file path."""
+        self._file_path = file_path
+
+    def load(self) -> List[Document]:
+        return [Document(page_content=self._load_as_text())]
+
+    def _load_as_text(self) -> str:
+        with open(self._file_path, "rb") as fp:
+            soup = BeautifulSoup(fp, 'html.parser')
+            text = soup.get_text()
+            text = text.strip() if text else ''
+
+        return text

api/core/data_loader/loader/markdown.py

@@ -0,0 +1,134 @@
+import logging
+import re
+from typing import Optional, List, Tuple, cast
+
+from langchain.document_loaders.base import BaseLoader
+from langchain.document_loaders.helpers import detect_file_encodings
+from langchain.schema import Document
+
+logger = logging.getLogger(__name__)
+
+
+class MarkdownLoader(BaseLoader):
+    """Load md files.
+
+
+    Args:
+        file_path: Path to the file to load.
+
+        remove_hyperlinks: Whether to remove hyperlinks from the text.
+
+        remove_images: Whether to remove images from the text.
+
+        encoding: File encoding to use. If `None`, the file will be loaded
+        with the default system encoding.
+
+        autodetect_encoding: Whether to try to autodetect the file encoding
+            if the specified encoding fails.
+    """
+
+    def __init__(
+        self,
+        file_path: str,
+        remove_hyperlinks: bool = True,
+        remove_images: bool = True,
+        encoding: Optional[str] = None,
+        autodetect_encoding: bool = True,
+    ):
+        """Initialize with file path."""
+        self._file_path = file_path
+        self._remove_hyperlinks = remove_hyperlinks
+        self._remove_images = remove_images
+        self._encoding = encoding
+        self._autodetect_encoding = autodetect_encoding
+
+    def load(self) -> List[Document]:
+        tups = self.parse_tups(self._file_path)
+        documents = []
+        for header, value in tups:
+            value = value.strip()
+            if header is None:
+                documents.append(Document(page_content=value))
+            else:
+                documents.append(Document(page_content=f"\n\n{header}\n{value}"))
+
+        return documents
+
+    def markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
+        """Convert a markdown file to a dictionary.
+
+        The keys are the headers and the values are the text under each header.
+
+        """
+        markdown_tups: List[Tuple[Optional[str], str]] = []
+        lines = markdown_text.split("\n")
+
+        current_header = None
+        current_text = ""
+
+        for line in lines:
+            header_match = re.match(r"^#+\s", line)
+            if header_match:
+                if current_header is not None:
+                    markdown_tups.append((current_header, current_text))
+
+                current_header = line
+                current_text = ""
+            else:
+                current_text += line + "\n"
+        markdown_tups.append((current_header, current_text))
+
+        if current_header is not None:
+            # pass linting, assert keys are defined
+            markdown_tups = [
+                (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
+                for key, value in markdown_tups
+            ]
+        else:
+            markdown_tups = [
+                (key, re.sub("\n", "", value)) for key, value in markdown_tups
+            ]
+
+        return markdown_tups
+
+    def remove_images(self, content: str) -> str:
+        """Get a dictionary of a markdown file from its path."""
+        pattern = r"!{1}\[\[(.*)\]\]"
+        content = re.sub(pattern, "", content)
+        return content
+
+    def remove_hyperlinks(self, content: str) -> str:
+        """Get a dictionary of a markdown file from its path."""
+        pattern = r"\[(.*?)\]\((.*?)\)"
+        content = re.sub(pattern, r"\1", content)
+        return content
+
+    def parse_tups(self, filepath: str) -> List[Tuple[Optional[str], str]]:
+        """Parse file into tuples."""
+        content = ""
+        try:
+            with open(filepath, "r", encoding=self._encoding) as f:
+                content = f.read()
+        except UnicodeDecodeError as e:
+            if self._autodetect_encoding:
+                detected_encodings = detect_file_encodings(filepath)
+                for encoding in detected_encodings:
+                    logger.debug("Trying encoding: ", encoding.encoding)
+                    try:
+                        with open(filepath, encoding=encoding.encoding) as f:
+                            content = f.read()
+                        break
+                    except UnicodeDecodeError:
+                        continue
+            else:
+                raise RuntimeError(f"Error loading {filepath}") from e
+        except Exception as e:
+            raise RuntimeError(f"Error loading {filepath}") from e
+
+        if self._remove_hyperlinks:
+            content = self.remove_hyperlinks(content)
+
+        if self._remove_images:
+            content = self.remove_images(content)
+
+        return self.markdown_to_tups(content)

api/core/data_loader/loader/notion.py

@@ -0,0 +1,379 @@
+import json
+import logging
+from typing import List, Dict, Any, Optional
+
+import requests
+from flask import current_app
+from langchain.document_loaders.base import BaseLoader
+from langchain.schema import Document
+
+from extensions.ext_database import db
+from models.dataset import Document as DocumentModel
+from models.source import DataSourceBinding
+
+logger = logging.getLogger(__name__)
+
+BLOCK_CHILD_URL_TMPL = "https://api.notion.com/v1/blocks/{block_id}/children"
+DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}/query"
+SEARCH_URL = "https://api.notion.com/v1/search"
+RETRIEVE_PAGE_URL_TMPL = "https://api.notion.com/v1/pages/{page_id}"
+RETRIEVE_DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}"
+HEADING_TYPE = ['heading_1', 'heading_2', 'heading_3']
+
+
+class NotionLoader(BaseLoader):
+    def __init__(
+            self,
+            notion_access_token: str,
+            notion_workspace_id: str,
+            notion_obj_id: str,
+            notion_page_type: str,
+            document_model: Optional[DocumentModel] = None
+    ):
+        self._document_model = document_model
+        self._notion_workspace_id = notion_workspace_id
+        self._notion_obj_id = notion_obj_id
+        self._notion_page_type = notion_page_type
+        self._notion_access_token = notion_access_token
+
+        if not self._notion_access_token:
+            integration_token = current_app.config.get('NOTION_INTEGRATION_TOKEN')
+            if integration_token is None:
+                raise ValueError(
+                    "Must specify `integration_token` or set environment "
+                    "variable `NOTION_INTEGRATION_TOKEN`."
+                )
+
+            self._notion_access_token = integration_token
+
+    @classmethod
+    def from_document(cls, document_model: DocumentModel):
+        data_source_info = document_model.data_source_info_dict
+        if not data_source_info or 'notion_page_id' not in data_source_info \
+                or 'notion_workspace_id' not in data_source_info:
+            raise ValueError("no notion page found")
+
+        notion_workspace_id = data_source_info['notion_workspace_id']
+        notion_obj_id = data_source_info['notion_page_id']
+        notion_page_type = data_source_info['type']
+        notion_access_token = cls._get_access_token(document_model.tenant_id, notion_workspace_id)
+
+        return cls(
+            notion_access_token=notion_access_token,
+            notion_workspace_id=notion_workspace_id,
+            notion_obj_id=notion_obj_id,
+            notion_page_type=notion_page_type,
+            document_model=document_model
+        )
+
+    def load(self) -> List[Document]:
+        self.update_last_edited_time(
+            self._document_model
+        )
+
+        text_docs = self._load_data_as_documents(self._notion_obj_id, self._notion_page_type)
+
+        return text_docs
+
+    def _load_data_as_documents(
+            self, notion_obj_id: str, notion_page_type: str
+    ) -> List[Document]:
+        docs = []
+        if notion_page_type == 'database':
+            # get all the pages in the database
+            page_text_documents = self._get_notion_database_data(notion_obj_id)
+            docs.extend(page_text_documents)
+        elif notion_page_type == 'page':
+            page_text_list = self._get_notion_block_data(notion_obj_id)
+            for page_text in page_text_list:
+                docs.append(Document(page_content=page_text))
+        else:
+            raise ValueError("notion page type not supported")
+
+        return docs
+
+    def _get_notion_database_data(
+            self, database_id: str, query_dict: Dict[str, Any] = {}
+    ) -> List[Document]:
+        """Get all the pages from a Notion database."""
+        res = requests.post(
+            DATABASE_URL_TMPL.format(database_id=database_id),
+            headers={
+                "Authorization": "Bearer " + self._notion_access_token,
+                "Content-Type": "application/json",
+                "Notion-Version": "2022-06-28",
+            },
+            json=query_dict,
+        )
+
+        data = res.json()
+
+        database_content_list = []
+        if 'results' not in data or data["results"] is None:
+            return []
+        for result in data["results"]:
+            properties = result['properties']
+            data = {}
+            for property_name, property_value in properties.items():
+                type = property_value['type']
+                if type == 'multi_select':
+                    value = []
+                    multi_select_list = property_value[type]
+                    for multi_select in multi_select_list:
+                        value.append(multi_select['name'])
+                elif type == 'rich_text' or type == 'title':
+                    if len(property_value[type]) > 0:
+                        value = property_value[type][0]['plain_text']
+                    else:
+                        value = ''
+                elif type == 'select' or type == 'status':
+                    if property_value[type]:
+                        value = property_value[type]['name']
+                    else:
+                        value = ''
+                else:
+                    value = property_value[type]
+                data[property_name] = value
+            row_dict = {k: v for k, v in data.items() if v}
+            row_content = ''
+            for key, value in row_dict.items():
+                if isinstance(value, dict):
+                    value_dict = {k: v for k, v in value.items() if v}
+                    value_content = ''.join(f'{k}:{v} ' for k, v in value_dict.items())
+                    row_content = row_content + f'{key}:{value_content}\n'
+                else:
+                    row_content = row_content + f'{key}:{value}\n'
+            document = Document(page_content=row_content)
+            database_content_list.append(document)
+
+        return database_content_list
+
+    def _get_notion_block_data(self, page_id: str) -> List[str]:
+        result_lines_arr = []
+        cur_block_id = page_id
+        while True:
+            block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id)
+            query_dict: Dict[str, Any] = {}
+
+            res = requests.request(
+                "GET",
+                block_url,
+                headers={
+                    "Authorization": "Bearer " + self._notion_access_token,
+                    "Content-Type": "application/json",
+                    "Notion-Version": "2022-06-28",
+                },
+                json=query_dict
+            )
+            data = res.json()
+            # current block's heading
+            heading = ''
+            for result in data["results"]:
+                result_type = result["type"]
+                result_obj = result[result_type]
+                cur_result_text_arr = []
+                if result_type == 'table':
+                    result_block_id = result["id"]
+                    text = self._read_table_rows(result_block_id)
+                    text += "\n\n"
+                    result_lines_arr.append(text)
+                else:
+                    if "rich_text" in result_obj:
+                        for rich_text in result_obj["rich_text"]:
+                            # skip if doesn't have text object
+                            if "text" in rich_text:
+                                text = rich_text["text"]["content"]
+                                cur_result_text_arr.append(text)
+                                if result_type in HEADING_TYPE:
+                                    heading = text
+
+                    result_block_id = result["id"]
+                    has_children = result["has_children"]
+                    block_type = result["type"]
+                    if has_children and block_type != 'child_page':
+                        children_text = self._read_block(
+                            result_block_id, num_tabs=1
+                        )
+                        cur_result_text_arr.append(children_text)
+
+                    cur_result_text = "\n".join(cur_result_text_arr)
+                    cur_result_text += "\n\n"
+                    if result_type in HEADING_TYPE:
+                        result_lines_arr.append(cur_result_text)
+                    else:
+                        result_lines_arr.append(f'{heading}\n{cur_result_text}')
+
+            if data["next_cursor"] is None:
+                break
+            else:
+                cur_block_id = data["next_cursor"]
+        return result_lines_arr
+
+    def _read_block(self, block_id: str, num_tabs: int = 0) -> str:
+        """Read a block."""
+        result_lines_arr = []
+        cur_block_id = block_id
+        while True:
+            block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id)
+            query_dict: Dict[str, Any] = {}
+
+            res = requests.request(
+                "GET",
+                block_url,
+                headers={
+                    "Authorization": "Bearer " + self._notion_access_token,
+                    "Content-Type": "application/json",
+                    "Notion-Version": "2022-06-28",
+                },
+                json=query_dict
+            )
+            data = res.json()
+            if 'results' not in data or data["results"] is None:
+                break
+            heading = ''
+            for result in data["results"]:
+                result_type = result["type"]
+                result_obj = result[result_type]
+                cur_result_text_arr = []
+                if result_type == 'table':
+                    result_block_id = result["id"]
+                    text = self._read_table_rows(result_block_id)
+                    result_lines_arr.append(text)
+                else:
+                    if "rich_text" in result_obj:
+                        for rich_text in result_obj["rich_text"]:
+                            # skip if doesn't have text object
+                            if "text" in rich_text:
+                                text = rich_text["text"]["content"]
+                                prefix = "\t" * num_tabs
+                                cur_result_text_arr.append(prefix + text)
+                                if result_type in HEADING_TYPE:
+                                    heading = text
+                    result_block_id = result["id"]
+                    has_children = result["has_children"]
+                    block_type = result["type"]
+                    if has_children and block_type != 'child_page':
+                        children_text = self._read_block(
+                            result_block_id, num_tabs=num_tabs + 1
+                        )
+                        cur_result_text_arr.append(children_text)
+
+                    cur_result_text = "\n".join(cur_result_text_arr)
+                    if result_type in HEADING_TYPE:
+                        result_lines_arr.append(cur_result_text)
+                    else:
+                        result_lines_arr.append(f'{heading}\n{cur_result_text}')
+
+            if data["next_cursor"] is None:
+                break
+            else:
+                cur_block_id = data["next_cursor"]
+
+        result_lines = "\n".join(result_lines_arr)
+        return result_lines
+
+    def _read_table_rows(self, block_id: str) -> str:
+        """Read table rows."""
+        done = False
+        result_lines_arr = []
+        cur_block_id = block_id
+        while not done:
+            block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id)
+            query_dict: Dict[str, Any] = {}
+
+            res = requests.request(
+                "GET",
+                block_url,
+                headers={
+                    "Authorization": "Bearer " + self._notion_access_token,
+                    "Content-Type": "application/json",
+                    "Notion-Version": "2022-06-28",
+                },
+                json=query_dict
+            )
+            data = res.json()
+            # get table headers text
+            table_header_cell_texts = []
+            tabel_header_cells = data["results"][0]['table_row']['cells']
+            for tabel_header_cell in tabel_header_cells:
+                if tabel_header_cell:
+                    for table_header_cell_text in tabel_header_cell:
+                        text = table_header_cell_text["text"]["content"]
+                        table_header_cell_texts.append(text)
+            # get table columns text and format
+            results = data["results"]
+            for i in range(len(results) - 1):
+                column_texts = []
+                tabel_column_cells = data["results"][i + 1]['table_row']['cells']
+                for j in range(len(tabel_column_cells)):
+                    if tabel_column_cells[j]:
+                        for table_column_cell_text in tabel_column_cells[j]:
+                            column_text = table_column_cell_text["text"]["content"]
+                            column_texts.append(f'{table_header_cell_texts[j]}:{column_text}')
+
+                cur_result_text = "\n".join(column_texts)
+                result_lines_arr.append(cur_result_text)
+
+            if data["next_cursor"] is None:
+                done = True
+                break
+            else:
+                cur_block_id = data["next_cursor"]
+
+        result_lines = "\n".join(result_lines_arr)
+        return result_lines
+
+    def update_last_edited_time(self, document_model: DocumentModel):
+        if not document_model:
+            return
+
+        last_edited_time = self.get_notion_last_edited_time()
+        data_source_info = document_model.data_source_info_dict
+        data_source_info['last_edited_time'] = last_edited_time
+        update_params = {
+            DocumentModel.data_source_info: json.dumps(data_source_info)
+        }
+
+        DocumentModel.query.filter_by(id=document_model.id).update(update_params)
+        db.session.commit()
+
+    def get_notion_last_edited_time(self) -> str:
+        obj_id = self._notion_obj_id
+        page_type = self._notion_page_type
+        if page_type == 'database':
+            retrieve_page_url = RETRIEVE_DATABASE_URL_TMPL.format(database_id=obj_id)
+        else:
+            retrieve_page_url = RETRIEVE_PAGE_URL_TMPL.format(page_id=obj_id)
+
+        query_dict: Dict[str, Any] = {}
+
+        res = requests.request(
+            "GET",
+            retrieve_page_url,
+            headers={
+                "Authorization": "Bearer " + self._notion_access_token,
+                "Content-Type": "application/json",
+                "Notion-Version": "2022-06-28",
+            },
+            json=query_dict
+        )
+
+        data = res.json()
+        return data["last_edited_time"]
+
+    @classmethod
+    def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str:
+        data_source_binding = DataSourceBinding.query.filter(
+            db.and_(
+                DataSourceBinding.tenant_id == tenant_id,
+                DataSourceBinding.provider == 'notion',
+                DataSourceBinding.disabled == False,
+                DataSourceBinding.source_info['workspace_id'] == f'"{notion_workspace_id}"'
+            )
+        ).first()
+
+        if not data_source_binding:
+            raise Exception(f'No notion data source binding found for tenant {tenant_id} '
+                            f'and notion workspace {notion_workspace_id}')
+
+        return data_source_binding.access_token

api/core/data_loader/loader/pdf.py

@@ -0,0 +1,55 @@
+import logging
+from typing import List, Optional
+
+from langchain.document_loaders import PyPDFium2Loader
+from langchain.document_loaders.base import BaseLoader
+from langchain.schema import Document
+
+from extensions.ext_storage import storage
+from models.model import UploadFile
+
+logger = logging.getLogger(__name__)
+
+
+class PdfLoader(BaseLoader):
+    """Load pdf files.
+
+
+    Args:
+        file_path: Path to the file to load.
+    """
+
+    def __init__(
+        self,
+        file_path: str,
+        upload_file: Optional[UploadFile] = None
+    ):
+        """Initialize with file path."""
+        self._file_path = file_path
+        self._upload_file = upload_file
+
+    def load(self) -> List[Document]:
+        plaintext_file_key = ''
+        plaintext_file_exists = False
+        if self._upload_file:
+            if self._upload_file.hash:
+                plaintext_file_key = 'upload_files/' + self._upload_file.tenant_id + '/' \
+                                     + self._upload_file.hash + '.0625.plaintext'
+                try:
+                    text = storage.load(plaintext_file_key).decode('utf-8')
+                    plaintext_file_exists = True
+                    return [Document(page_content=text)]
+                except FileNotFoundError:
+                    pass
+        documents = PyPDFium2Loader(file_path=self._file_path).load()
+        text_list = []
+        for document in documents:
+            text_list.append(document.page_content)
+        text = "\n\n".join(text_list)
+
+        # save plaintext file for caching
+        if not plaintext_file_exists and plaintext_file_key:
+            storage.save(plaintext_file_key, text.encode('utf-8'))
+
+        return documents
+

api/core/docstore/dataset_docstore.py

@@ -1,10 +1,6 @@
 from typing import Any, Dict, Optional, Sequence
 
-import tiktoken
-from llama_index.data_structs import Node
-from llama_index.docstore.types import BaseDocumentStore
-from llama_index.docstore.utils import json_to_doc
-from llama_index.schema import BaseDocument
+from langchain.schema import Document
 from sqlalchemy import func
 
 from core.llm.token_calculator import TokenCalculator
@@ -12,7 +8,7 @@ from extensions.ext_database import db
 from models.dataset import Dataset, DocumentSegment
 
 
-class DatesetDocumentStore(BaseDocumentStore):
+class DatesetDocumentStore:
     def __init__(
         self,
         dataset: Dataset,
@@ -48,7 +44,7 @@ class DatesetDocumentStore(BaseDocumentStore):
         return self._embedding_model_name
 
     @property
-    def docs(self) -> Dict[str, BaseDocument]:
+    def docs(self) -> Dict[str, Document]:
         document_segments = db.session.query(DocumentSegment).filter(
             DocumentSegment.dataset_id == self._dataset.id
         ).all()
@@ -56,13 +52,20 @@ class DatesetDocumentStore(BaseDocumentStore):
         output = {}
         for document_segment in document_segments:
             doc_id = document_segment.index_node_id
-            result = self.segment_to_dict(document_segment)
-            output[doc_id] = json_to_doc(result)
+            output[doc_id] = Document(
+                page_content=document_segment.content,
+                metadata={
+                    "doc_id": document_segment.index_node_id,
+                    "doc_hash": document_segment.index_node_hash,
+                    "document_id": document_segment.document_id,
+                    "dataset_id": document_segment.dataset_id,
+                }
+            )
 
         return output
 
     def add_documents(
-        self, docs: Sequence[BaseDocument], allow_update: bool = True
+        self, docs: Sequence[Document], allow_update: bool = True
     ) -> None:
         max_position = db.session.query(func.max(DocumentSegment.position)).filter(
             DocumentSegment.document == self._document_id
@@ -72,23 +75,20 @@ class DatesetDocumentStore(BaseDocumentStore):
             max_position = 0
 
         for doc in docs:
-            if doc.is_doc_id_none:
-                raise ValueError("doc_id not set")
+            if not isinstance(doc, Document):
+                raise ValueError("doc must be a Document")
 
-            if not isinstance(doc, Node):
-                raise ValueError("doc must be a Node")
-
-            segment_document = self.get_document(doc_id=doc.get_doc_id(), raise_error=False)
+            segment_document = self.get_document(doc_id=doc.metadata['doc_id'], raise_error=False)
 
             # NOTE: doc could already exist in the store, but we overwrite it
             if not allow_update and segment_document:
                 raise ValueError(
-                    f"doc_id {doc.get_doc_id()} already exists. "
+                    f"doc_id {doc.metadata['doc_id']} already exists. "
                     "Set allow_update to True to overwrite."
                 )
 
             # calc embedding use tokens
-            tokens = TokenCalculator.get_num_tokens(self._embedding_model_name, doc.get_text())
+            tokens = TokenCalculator.get_num_tokens(self._embedding_model_name, doc.page_content)
 
             if not segment_document:
                 max_position += 1
@@ -97,19 +97,19 @@ class DatesetDocumentStore(BaseDocumentStore):
                     tenant_id=self._dataset.tenant_id,
                     dataset_id=self._dataset.id,
                     document_id=self._document_id,
-                    index_node_id=doc.get_doc_id(),
-                    index_node_hash=doc.get_doc_hash(),
+                    index_node_id=doc.metadata['doc_id'],
+                    index_node_hash=doc.metadata['doc_hash'],
                     position=max_position,
-                    content=doc.get_text(),
-                    word_count=len(doc.get_text()),
+                    content=doc.page_content,
+                    word_count=len(doc.page_content),
                     tokens=tokens,
                     created_by=self._user_id,
                 )
                 db.session.add(segment_document)
             else:
-                segment_document.content = doc.get_text()
-                segment_document.index_node_hash = doc.get_doc_hash()
-                segment_document.word_count = len(doc.get_text())
+                segment_document.content = doc.page_content
+                segment_document.index_node_hash = doc.metadata['doc_hash']
+                segment_document.word_count = len(doc.page_content)
                 segment_document.tokens = tokens
 
             db.session.commit()
@@ -121,7 +121,7 @@ class DatesetDocumentStore(BaseDocumentStore):
 
     def get_document(
         self, doc_id: str, raise_error: bool = True
-    ) -> Optional[BaseDocument]:
+    ) -> Optional[Document]:
         document_segment = self.get_document_segment(doc_id)
 
         if document_segment is None:
@@ -130,8 +130,15 @@ class DatesetDocumentStore(BaseDocumentStore):
             else:
                 return None
 
-        result = self.segment_to_dict(document_segment)
-        return json_to_doc(result)
+        return Document(
+            page_content=document_segment.content,
+            metadata={
+                "doc_id": document_segment.index_node_id,
+                "doc_hash": document_segment.index_node_hash,
+                "document_id": document_segment.document_id,
+                "dataset_id": document_segment.dataset_id,
+            }
+        )
 
     def delete_document(self, doc_id: str, raise_error: bool = True) -> None:
         document_segment = self.get_document_segment(doc_id)
@@ -164,15 +171,6 @@ class DatesetDocumentStore(BaseDocumentStore):
 
         return document_segment.index_node_hash
 
-    def update_docstore(self, other: "BaseDocumentStore") -> None:
-        """Update docstore.
-
-        Args:
-            other (BaseDocumentStore): docstore to update from
-
-        """
-        self.add_documents(list(other.docs.values()))
-
     def get_document_segment(self, doc_id: str) -> DocumentSegment:
         document_segment = db.session.query(DocumentSegment).filter(
             DocumentSegment.dataset_id == self._dataset.id,
@@ -180,11 +178,3 @@ class DatesetDocumentStore(BaseDocumentStore):
         ).first()
 
         return document_segment
-
-    def segment_to_dict(self, segment: DocumentSegment) -> Dict[str, Any]:
-        return {
-            "doc_id": segment.index_node_id,
-            "doc_hash": segment.index_node_hash,
-            "text": segment.content,
-            "__type__": Node.get_type()
-        }

api/core/docstore/empty_docstore.py

@@ -1,51 +0,0 @@
-from typing import Any, Dict, Optional, Sequence
-from llama_index.docstore.types import BaseDocumentStore
-from llama_index.schema import BaseDocument
-
-
-class EmptyDocumentStore(BaseDocumentStore):
-    @classmethod
-    def from_dict(cls, config_dict: Dict[str, Any]) -> "EmptyDocumentStore":
-        return cls()
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Serialize to dict."""
-        return {}
-
-    @property
-    def docs(self) -> Dict[str, BaseDocument]:
-        return {}
-
-    def add_documents(
-        self, docs: Sequence[BaseDocument], allow_update: bool = True
-    ) -> None:
-        pass
-
-    def document_exists(self, doc_id: str) -> bool:
-        """Check if document exists."""
-        return False
-
-    def get_document(
-        self, doc_id: str, raise_error: bool = True
-    ) -> Optional[BaseDocument]:
-        return None
-
-    def delete_document(self, doc_id: str, raise_error: bool = True) -> None:
-        pass
-
-    def set_document_hash(self, doc_id: str, doc_hash: str) -> None:
-        """Set the hash for a given doc_id."""
-        pass
-
-    def get_document_hash(self, doc_id: str) -> Optional[str]:
-        """Get the stored hash for a document, if it exists."""
-        return None
-
-    def update_docstore(self, other: "BaseDocumentStore") -> None:
-        """Update docstore.
-
-        Args:
-            other (BaseDocumentStore): docstore to update from
-
-        """
-        self.add_documents(list(other.docs.values()))

api/core/embedding/cached_embedding.py

@@ -0,0 +1,72 @@
+import logging
+from typing import List
+
+from langchain.embeddings.base import Embeddings
+from sqlalchemy.exc import IntegrityError
+
+from extensions.ext_database import db
+from libs import helper
+from models.dataset import Embedding
+
+
+class CacheEmbedding(Embeddings):
+    def __init__(self, embeddings: Embeddings):
+        self._embeddings = embeddings
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed search docs."""
+        # use doc embedding cache or store if not exists
+        text_embeddings = []
+        embedding_queue_texts = []
+        for text in texts:
+            hash = helper.generate_text_hash(text)
+            embedding = db.session.query(Embedding).filter_by(hash=hash).first()
+            if embedding:
+                text_embeddings.append(embedding.get_embedding())
+            else:
+                embedding_queue_texts.append(text)
+
+        embedding_results = self._embeddings.embed_documents(embedding_queue_texts)
+
+        i = 0
+        for text in embedding_queue_texts:
+            hash = helper.generate_text_hash(text)
+
+            try:
+                embedding = Embedding(hash=hash)
+                embedding.set_embedding(embedding_results[i])
+                db.session.add(embedding)
+                db.session.commit()
+            except IntegrityError:
+                db.session.rollback()
+                continue
+            except:
+                logging.exception('Failed to add embedding to db')
+                continue
+
+            i += 1
+
+        text_embeddings.extend(embedding_results)
+        return text_embeddings
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed query text."""
+        # use doc embedding cache or store if not exists
+        hash = helper.generate_text_hash(text)
+        embedding = db.session.query(Embedding).filter_by(hash=hash).first()
+        if embedding:
+            return embedding.get_embedding()
+
+        embedding_results = self._embeddings.embed_query(text)
+
+        try:
+            embedding = Embedding(hash=hash)
+            embedding.set_embedding(embedding_results)
+            db.session.add(embedding)
+            db.session.commit()
+        except IntegrityError:
+            db.session.rollback()
+        except:
+            logging.exception('Failed to add embedding to db')
+
+        return embedding_results

api/core/embedding/openai_embedding.py

@@ -1,214 +0,0 @@
-from typing import Optional, Any, List
-
-import openai
-from llama_index.embeddings.base import BaseEmbedding
-from llama_index.embeddings.openai import OpenAIEmbeddingMode, OpenAIEmbeddingModelType, _QUERY_MODE_MODEL_DICT, \
-    _TEXT_MODE_MODEL_DICT
-from tenacity import wait_random_exponential, retry, stop_after_attempt
-
-from core.llm.error_handle_wraps import handle_llm_exceptions, handle_llm_exceptions_async
-
-
-@retry(reraise=True, wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
-def get_embedding(
-        text: str,
-        engine: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **kwargs
-) -> List[float]:
-    """Get embedding.
-
-    NOTE: Copied from OpenAI's embedding utils:
-    https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py
-
-    Copied here to avoid importing unnecessary dependencies
-    like matplotlib, plotly, scipy, sklearn.
-
-    """
-    text = text.replace("\n", " ")
-    return openai.Embedding.create(input=[text], engine=engine, api_key=api_key, **kwargs)["data"][0]["embedding"]
-
-
-@retry(reraise=True, wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
-async def aget_embedding(text: str, engine: Optional[str] = None, api_key: Optional[str] = None, **kwargs) -> List[
-    float]:
-    """Asynchronously get embedding.
-
-    NOTE: Copied from OpenAI's embedding utils:
-    https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py
-
-    Copied here to avoid importing unnecessary dependencies
-    like matplotlib, plotly, scipy, sklearn.
-
-    """
-    # replace newlines, which can negatively affect performance.
-    text = text.replace("\n", " ")
-
-    return (await openai.Embedding.acreate(input=[text], engine=engine, api_key=api_key, **kwargs))["data"][0][
-        "embedding"
-    ]
-
-
-@retry(reraise=True, wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
-def get_embeddings(
-        list_of_text: List[str],
-        engine: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **kwargs
-) -> List[List[float]]:
-    """Get embeddings.
-
-    NOTE: Copied from OpenAI's embedding utils:
-    https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py
-
-    Copied here to avoid importing unnecessary dependencies
-    like matplotlib, plotly, scipy, sklearn.
-
-    """
-    assert len(list_of_text) <= 2048, "The batch size should not be larger than 2048."
-
-    # replace newlines, which can negatively affect performance.
-    list_of_text = [text.replace("\n", " ") for text in list_of_text]
-
-    data = openai.Embedding.create(input=list_of_text, engine=engine, api_key=api_key, **kwargs).data
-    data = sorted(data, key=lambda x: x["index"])  # maintain the same order as input.
-    return [d["embedding"] for d in data]
-
-
-@retry(reraise=True, wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
-async def aget_embeddings(
-        list_of_text: List[str], engine: Optional[str] = None, api_key: Optional[str] = None, **kwargs
-) -> List[List[float]]:
-    """Asynchronously get embeddings.
-
-    NOTE: Copied from OpenAI's embedding utils:
-    https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py
-
-    Copied here to avoid importing unnecessary dependencies
-    like matplotlib, plotly, scipy, sklearn.
-
-    """
-    assert len(list_of_text) <= 2048, "The batch size should not be larger than 2048."
-
-    # replace newlines, which can negatively affect performance.
-    list_of_text = [text.replace("\n", " ") for text in list_of_text]
-
-    data = (await openai.Embedding.acreate(input=list_of_text, engine=engine, api_key=api_key, **kwargs)).data
-    data = sorted(data, key=lambda x: x["index"])  # maintain the same order as input.
-    return [d["embedding"] for d in data]
-
-
-class OpenAIEmbedding(BaseEmbedding):
-
-    def __init__(
-            self,
-            mode: str = OpenAIEmbeddingMode.TEXT_SEARCH_MODE,
-            model: str = OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002,
-            deployment_name: Optional[str] = None,
-            openai_api_key: Optional[str] = None,
-            **kwargs: Any,
-    ) -> None:
-        """Init params."""
-        new_kwargs = {}
-
-        if 'embed_batch_size' in kwargs:
-            new_kwargs['embed_batch_size'] = kwargs['embed_batch_size']
-
-        if 'tokenizer' in kwargs:
-            new_kwargs['tokenizer'] = kwargs['tokenizer']
-
-        super().__init__(**new_kwargs)
-        self.mode = OpenAIEmbeddingMode(mode)
-        self.model = OpenAIEmbeddingModelType(model)
-        self.deployment_name = deployment_name
-        self.openai_api_key = openai_api_key
-        self.openai_api_type = kwargs.get('openai_api_type')
-        self.openai_api_version = kwargs.get('openai_api_version')
-        self.openai_api_base = kwargs.get('openai_api_base')
-
-    @handle_llm_exceptions
-    def _get_query_embedding(self, query: str) -> List[float]:
-        """Get query embedding."""
-        if self.deployment_name is not None:
-            engine = self.deployment_name
-        else:
-            key = (self.mode, self.model)
-            if key not in _QUERY_MODE_MODEL_DICT:
-                raise ValueError(f"Invalid mode, model combination: {key}")
-            engine = _QUERY_MODE_MODEL_DICT[key]
-        return get_embedding(query, engine=engine, api_key=self.openai_api_key,
-                             api_type=self.openai_api_type, api_version=self.openai_api_version,
-                             api_base=self.openai_api_base)
-
-    def _get_text_embedding(self, text: str) -> List[float]:
-        """Get text embedding."""
-        if self.deployment_name is not None:
-            engine = self.deployment_name
-        else:
-            key = (self.mode, self.model)
-            if key not in _TEXT_MODE_MODEL_DICT:
-                raise ValueError(f"Invalid mode, model combination: {key}")
-            engine = _TEXT_MODE_MODEL_DICT[key]
-        return get_embedding(text, engine=engine, api_key=self.openai_api_key,
-                             api_type=self.openai_api_type, api_version=self.openai_api_version,
-                             api_base=self.openai_api_base)
-
-    async def _aget_text_embedding(self, text: str) -> List[float]:
-        """Asynchronously get text embedding."""
-        if self.deployment_name is not None:
-            engine = self.deployment_name
-        else:
-            key = (self.mode, self.model)
-            if key not in _TEXT_MODE_MODEL_DICT:
-                raise ValueError(f"Invalid mode, model combination: {key}")
-            engine = _TEXT_MODE_MODEL_DICT[key]
-        return await aget_embedding(text, engine=engine, api_key=self.openai_api_key,
-                                    api_type=self.openai_api_type, api_version=self.openai_api_version,
-                                    api_base=self.openai_api_base)
-
-    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
-        """Get text embeddings.
-
-        By default, this is a wrapper around _get_text_embedding.
-        Can be overriden for batch queries.
-
-        """
-        if self.openai_api_type and self.openai_api_type == 'azure':
-            embeddings = []
-            for text in texts:
-                embeddings.append(self._get_text_embedding(text))
-
-            return embeddings
-
-        if self.deployment_name is not None:
-            engine = self.deployment_name
-        else:
-            key = (self.mode, self.model)
-            if key not in _TEXT_MODE_MODEL_DICT:
-                raise ValueError(f"Invalid mode, model combination: {key}")
-            engine = _TEXT_MODE_MODEL_DICT[key]
-        embeddings = get_embeddings(texts, engine=engine, api_key=self.openai_api_key,
-                                    api_type=self.openai_api_type, api_version=self.openai_api_version,
-                                    api_base=self.openai_api_base)
-        return embeddings
-
-    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
-        """Asynchronously get text embeddings."""
-        if self.openai_api_type and self.openai_api_type == 'azure':
-            embeddings = []
-            for text in texts:
-                embeddings.append(await self._aget_text_embedding(text))
-
-            return embeddings
-
-        if self.deployment_name is not None:
-            engine = self.deployment_name
-        else:
-            key = (self.mode, self.model)
-            if key not in _TEXT_MODE_MODEL_DICT:
-                raise ValueError(f"Invalid mode, model combination: {key}")
-            engine = _TEXT_MODE_MODEL_DICT[key]
-        embeddings = await aget_embeddings(texts, engine=engine, api_key=self.openai_api_key,
-                                           api_type=self.openai_api_type, api_version=self.openai_api_version,
-                                           api_base=self.openai_api_base)
-        return embeddings

api/core/generator/llm_generator.py

@@ -1,7 +1,8 @@
 import logging
 
+from langchain import PromptTemplate
 from langchain.chat_models.base import BaseChatModel
-from langchain.schema import HumanMessage, OutputParserException
+from langchain.schema import HumanMessage, OutputParserException, BaseMessage
 
 from core.constant import llm_constant
 from core.llm.llm_builder import LLMBuilder
@@ -10,7 +11,7 @@ from core.llm.token_calculator import TokenCalculator
 from core.prompt.output_parser.rule_config_generator import RuleConfigGeneratorOutputParser
 
 from core.prompt.output_parser.suggested_questions_after_answer import SuggestedQuestionsAfterAnswerOutputParser
-from core.prompt.prompt_template import OutLinePromptTemplate
+from core.prompt.prompt_template import JinjaPromptTemplate, OutLinePromptTemplate
 from core.prompt.prompts import CONVERSATION_TITLE_PROMPT, CONVERSATION_SUMMARY_PROMPT, INTRODUCTION_GENERATE_PROMPT
 
 
@@ -22,10 +23,10 @@ class LLMGenerator:
     @classmethod
     def generate_conversation_name(cls, tenant_id: str, query, answer):
         prompt = CONVERSATION_TITLE_PROMPT
-        prompt = prompt.format(query=query, answer=answer)
+        prompt = prompt.format(query=query)
         llm: StreamableOpenAI = LLMBuilder.to_llm(
             tenant_id=tenant_id,
-            model_name=generate_base_model,
+            model_name='gpt-3.5-turbo',
             max_tokens=50
         )
 
@@ -39,11 +40,12 @@ class LLMGenerator:
     @classmethod
     def generate_conversation_summary(cls, tenant_id: str, messages):
         max_tokens = 200
+        model = 'gpt-3.5-turbo'
 
         prompt = CONVERSATION_SUMMARY_PROMPT
         prompt_with_empty_context = prompt.format(context='')
-        prompt_tokens = TokenCalculator.get_num_tokens(generate_base_model, prompt_with_empty_context)
-        rest_tokens = llm_constant.max_context_token_length[generate_base_model] - prompt_tokens - max_tokens
+        prompt_tokens = TokenCalculator.get_num_tokens(model, prompt_with_empty_context)
+        rest_tokens = llm_constant.max_context_token_length[model] - prompt_tokens - max_tokens - 1
 
         context = ''
         for message in messages:
@@ -51,14 +53,17 @@ class LLMGenerator:
                 continue
 
             message_qa_text = "Human:" + message.query + "\nAI:" + message.answer + "\n"
-            if rest_tokens - TokenCalculator.get_num_tokens(generate_base_model, context + message_qa_text) > 0:
+            if rest_tokens - TokenCalculator.get_num_tokens(model, context + message_qa_text) > 0:
                 context += message_qa_text
 
+        if not context:
+            return '[message too long, no summary]'
+
         prompt = prompt.format(context=context)
 
         llm: StreamableOpenAI = LLMBuilder.to_llm(
             tenant_id=tenant_id,
-            model_name=generate_base_model,
+            model_name=model,
             max_tokens=max_tokens
         )
 
@@ -91,8 +96,8 @@ class LLMGenerator:
         output_parser = SuggestedQuestionsAfterAnswerOutputParser()
         format_instructions = output_parser.get_format_instructions()
 
-        prompt = OutLinePromptTemplate(
-            template="{histories}\n{format_instructions}\nquestions:\n",
+        prompt = JinjaPromptTemplate(
+            template="{{histories}}\n{{format_instructions}}\nquestions:\n",
             input_variables=["histories"],
             partial_variables={"format_instructions": format_instructions}
         )
@@ -101,7 +106,7 @@ class LLMGenerator:
 
         llm: StreamableOpenAI = LLMBuilder.to_llm(
             tenant_id=tenant_id,
-            model_name=generate_base_model,
+            model_name='gpt-3.5-turbo',
             temperature=0,
             max_tokens=256
         )
@@ -113,6 +118,8 @@ class LLMGenerator:
 
         try:
             output = llm(query)
+            if isinstance(output, BaseMessage):
+                output = output.content
             questions = output_parser.parse(output)
         except Exception:
             logging.exception("Error generating suggested questions after answer")

api/core/index/base.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+from abc import abstractmethod, ABC
+from typing import List, Any
+
+from langchain.schema import Document, BaseRetriever
+
+from models.dataset import Dataset
+
+
+class BaseIndex(ABC):
+
+    def __init__(self, dataset: Dataset):
+        self.dataset = dataset
+
+    @abstractmethod
+    def create(self, texts: list[Document], **kwargs) -> BaseIndex:
+        raise NotImplementedError
+
+    @abstractmethod
+    def add_texts(self, texts: list[Document], **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def text_exists(self, id: str) -> bool:
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_by_ids(self, ids: list[str]) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_by_document_id(self, document_id: str):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_retriever(self, **kwargs: Any) -> BaseRetriever:
+        raise NotImplementedError
+
+    @abstractmethod
+    def search(
+            self, query: str,
+            **kwargs: Any
+    ) -> List[Document]:
+        raise NotImplementedError
+
+    def delete(self) -> None:
+        raise NotImplementedError
+
+    def _filter_duplicate_texts(self, texts: list[Document]) -> list[Document]:
+        for text in texts:
+            doc_id = text.metadata['doc_id']
+            exists_duplicate_node = self.text_exists(doc_id)
+            if exists_duplicate_node:
+                texts.remove(text)
+
+        return texts
+
+    def _get_uuids(self, texts: list[Document]) -> list[str]:
+        return [text.metadata['doc_id'] for text in texts]

api/core/index/index.py

@@ -0,0 +1,41 @@
+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.llm.llm_builder import LLMBuilder
+from models.dataset import Dataset
+
+
+class IndexBuilder:
+    @classmethod
+    def get_index(cls, dataset: Dataset, indexing_technique: str, ignore_high_quality_check: bool = False):
+        if indexing_technique == "high_quality":
+            if not ignore_high_quality_check and dataset.indexing_technique != 'high_quality':
+                return None
+
+            model_credentials = LLMBuilder.get_model_credentials(
+                tenant_id=dataset.tenant_id,
+                model_provider=LLMBuilder.get_default_provider(dataset.tenant_id),
+                model_name='text-embedding-ada-002'
+            )
+
+            embeddings = CacheEmbedding(OpenAIEmbeddings(
+                **model_credentials
+            ))
+
+            return VectorIndex(
+                dataset=dataset,
+                config=current_app.config,
+                embeddings=embeddings
+            )
+        elif indexing_technique == "economy":
+            return KeywordTableIndex(
+                dataset=dataset,
+                config=KeywordTableConfig(
+                    max_keywords_per_chunk=10
+                )
+            )
+        else:
+            raise ValueError('Unknown indexing technique')

api/core/index/index_builder.py

@@ -1,60 +0,0 @@
-from langchain.callbacks import CallbackManager
-from llama_index import ServiceContext, PromptHelper, LLMPredictor
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.embedding.openai_embedding import OpenAIEmbedding
-from core.llm.llm_builder import LLMBuilder
-
-
-class IndexBuilder:
-    @classmethod
-    def get_default_service_context(cls, tenant_id: str) -> ServiceContext:
-        # set number of output tokens
-        num_output = 512
-
-        # only for verbose
-        callback_manager = CallbackManager([DifyStdOutCallbackHandler()])
-
-        llm = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name='text-davinci-003',
-            temperature=0,
-            max_tokens=num_output,
-            callback_manager=callback_manager,
-        )
-
-        llm_predictor = LLMPredictor(llm=llm)
-
-        # These parameters here will affect the logic of segmenting the final synthesized response.
-        # The number of refinement iterations in the synthesis process depends
-        # on whether the length of the segmented output exceeds the max_input_size.
-        prompt_helper = PromptHelper(
-            max_input_size=3500,
-            num_output=num_output,
-            max_chunk_overlap=20
-        )
-
-        provider = LLMBuilder.get_default_provider(tenant_id)
-
-        model_credentials = LLMBuilder.get_model_credentials(
-            tenant_id=tenant_id,
-            model_provider=provider,
-            model_name='text-embedding-ada-002'
-        )
-
-        return ServiceContext.from_defaults(
-            llm_predictor=llm_predictor,
-            prompt_helper=prompt_helper,
-            embed_model=OpenAIEmbedding(**model_credentials),
-        )
-
-    @classmethod
-    def get_fake_llm_service_context(cls, tenant_id: str) -> ServiceContext:
-        llm = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name='fake'
-        )
-
-        return ServiceContext.from_defaults(
-            llm_predictor=LLMPredictor(llm=llm),
-            embed_model=OpenAIEmbedding()
-        )

api/core/index/keyword_table/jieba_keyword_table.py

@@ -1,159 +0,0 @@
-import re
-from typing import (
-    Any,
-    Dict,
-    List,
-    Set,
-    Optional
-)
-
-import jieba.analyse
-
-from core.index.keyword_table.stopwords import STOPWORDS
-from llama_index.indices.query.base import IS
-from llama_index import QueryMode
-from llama_index.indices.base import QueryMap
-from llama_index.indices.keyword_table.base import BaseGPTKeywordTableIndex
-from llama_index.indices.keyword_table.query import BaseGPTKeywordTableQuery
-from llama_index.docstore import BaseDocumentStore
-from llama_index.indices.postprocessor.node import (
-    BaseNodePostprocessor,
-)
-from llama_index.indices.response.response_builder import ResponseMode
-from llama_index.indices.service_context import ServiceContext
-from llama_index.optimization.optimizer import BaseTokenUsageOptimizer
-from llama_index.prompts.prompts import (
-    QuestionAnswerPrompt,
-    RefinePrompt,
-    SimpleInputPrompt,
-)
-
-from core.index.query.synthesizer import EnhanceResponseSynthesizer
-
-
-def jieba_extract_keywords(
-        text_chunk: str,
-        max_keywords: Optional[int] = None,
-        expand_with_subtokens: bool = True,
-) -> Set[str]:
-    """Extract keywords with JIEBA tfidf."""
-    keywords = jieba.analyse.extract_tags(
-        sentence=text_chunk,
-        topK=max_keywords,
-    )
-
-    if expand_with_subtokens:
-        return set(expand_tokens_with_subtokens(keywords))
-    else:
-        return set(keywords)
-
-
-def expand_tokens_with_subtokens(tokens: Set[str]) -> Set[str]:
-    """Get subtokens from a list of tokens., filtering for stopwords."""
-    results = set()
-    for token in tokens:
-        results.add(token)
-        sub_tokens = re.findall(r"\w+", token)
-        if len(sub_tokens) > 1:
-            results.update({w for w in sub_tokens if w not in list(STOPWORDS)})
-
-    return results
-
-
-class GPTJIEBAKeywordTableIndex(BaseGPTKeywordTableIndex):
-    """GPT JIEBA Keyword Table Index.
-
-    This index uses a JIEBA keyword extractor to extract keywords from the text.
-
-    """
-
-    def _extract_keywords(self, text: str) -> Set[str]:
-        """Extract keywords from text."""
-        return jieba_extract_keywords(text, max_keywords=self.max_keywords_per_chunk)
-
-    @classmethod
-    def get_query_map(self) -> QueryMap:
-        """Get query map."""
-        super_map = super().get_query_map()
-        super_map[QueryMode.DEFAULT] = GPTKeywordTableJIEBAQuery
-        return super_map
-
-    def _delete(self, doc_id: str, **delete_kwargs: Any) -> None:
-        """Delete a document."""
-        # get set of ids that correspond to node
-        node_idxs_to_delete = {doc_id}
-
-        # delete node_idxs from keyword to node idxs mapping
-        keywords_to_delete = set()
-        for keyword, node_idxs in self._index_struct.table.items():
-            if node_idxs_to_delete.intersection(node_idxs):
-                self._index_struct.table[keyword] = node_idxs.difference(
-                    node_idxs_to_delete
-                )
-                if not self._index_struct.table[keyword]:
-                    keywords_to_delete.add(keyword)
-
-        for keyword in keywords_to_delete:
-            del self._index_struct.table[keyword]
-
-
-class GPTKeywordTableJIEBAQuery(BaseGPTKeywordTableQuery):
-    """GPT Keyword Table Index JIEBA Query.
-
-    Extracts keywords using JIEBA keyword extractor.
-    Set when `mode="jieba"` in `query` method of `GPTKeywordTableIndex`.
-
-    .. code-block:: python
-
-        response = index.query("<query_str>", mode="jieba")
-
-    See BaseGPTKeywordTableQuery for arguments.
-
-    """
-
-    @classmethod
-    def from_args(
-            cls,
-            index_struct: IS,
-            service_context: ServiceContext,
-            docstore: Optional[BaseDocumentStore] = None,
-            node_postprocessors: Optional[List[BaseNodePostprocessor]] = None,
-            verbose: bool = False,
-            # response synthesizer args
-            response_mode: ResponseMode = ResponseMode.DEFAULT,
-            text_qa_template: Optional[QuestionAnswerPrompt] = None,
-            refine_template: Optional[RefinePrompt] = None,
-            simple_template: Optional[SimpleInputPrompt] = None,
-            response_kwargs: Optional[Dict] = None,
-            use_async: bool = False,
-            streaming: bool = False,
-            optimizer: Optional[BaseTokenUsageOptimizer] = None,
-            # class-specific args
-            **kwargs: Any,
-    ) -> "BaseGPTIndexQuery":
-        response_synthesizer = EnhanceResponseSynthesizer.from_args(
-            service_context=service_context,
-            text_qa_template=text_qa_template,
-            refine_template=refine_template,
-            simple_template=simple_template,
-            response_mode=response_mode,
-            response_kwargs=response_kwargs,
-            use_async=use_async,
-            streaming=streaming,
-            optimizer=optimizer,
-        )
-        return cls(
-            index_struct=index_struct,
-            service_context=service_context,
-            response_synthesizer=response_synthesizer,
-            docstore=docstore,
-            node_postprocessors=node_postprocessors,
-            verbose=verbose,
-            **kwargs,
-        )
-
-    def _get_keywords(self, query_str: str) -> List[str]:
-        """Extract keywords."""
-        return list(
-            jieba_extract_keywords(query_str, max_keywords=self.max_keywords_per_query)
-        )

api/core/index/keyword_table_index.py

@@ -1,135 +0,0 @@
-import json
-from typing import List, Optional
-
-from llama_index import ServiceContext, LLMPredictor, OpenAIEmbedding
-from llama_index.data_structs import KeywordTable, Node
-from llama_index.indices.keyword_table.base import BaseGPTKeywordTableIndex
-from llama_index.indices.registry import load_index_struct_from_dict
-
-from core.docstore.dataset_docstore import DatesetDocumentStore
-from core.docstore.empty_docstore import EmptyDocumentStore
-from core.index.index_builder import IndexBuilder
-from core.index.keyword_table.jieba_keyword_table import GPTJIEBAKeywordTableIndex
-from core.llm.llm_builder import LLMBuilder
-from extensions.ext_database import db
-from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment
-
-
-class KeywordTableIndex:
-
-    def __init__(self, dataset: Dataset):
-        self._dataset = dataset
-
-    def add_nodes(self, nodes: List[Node]):
-        llm = LLMBuilder.to_llm(
-            tenant_id=self._dataset.tenant_id,
-            model_name='fake'
-        )
-
-        service_context = ServiceContext.from_defaults(
-            llm_predictor=LLMPredictor(llm=llm),
-            embed_model=OpenAIEmbedding()
-        )
-
-        dataset_keyword_table = self.get_keyword_table()
-        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
-            index_struct = KeywordTable()
-        else:
-            index_struct_dict = dataset_keyword_table.keyword_table_dict
-            index_struct: KeywordTable = load_index_struct_from_dict(index_struct_dict)
-
-        # create index
-        index = GPTJIEBAKeywordTableIndex(
-            index_struct=index_struct,
-            docstore=EmptyDocumentStore(),
-            service_context=service_context
-        )
-
-        for node in nodes:
-            keywords = index._extract_keywords(node.get_text())
-            self.update_segment_keywords(node.doc_id, list(keywords))
-            index._index_struct.add_node(list(keywords), node)
-
-        index_struct_dict = index.index_struct.to_dict()
-
-        if not dataset_keyword_table:
-            dataset_keyword_table = DatasetKeywordTable(
-                dataset_id=self._dataset.id,
-                keyword_table=json.dumps(index_struct_dict)
-            )
-            db.session.add(dataset_keyword_table)
-        else:
-            dataset_keyword_table.keyword_table = json.dumps(index_struct_dict)
-
-        db.session.commit()
-
-    def del_nodes(self, node_ids: List[str]):
-        llm = LLMBuilder.to_llm(
-            tenant_id=self._dataset.tenant_id,
-            model_name='fake'
-        )
-
-        service_context = ServiceContext.from_defaults(
-            llm_predictor=LLMPredictor(llm=llm),
-            embed_model=OpenAIEmbedding()
-        )
-
-        dataset_keyword_table = self.get_keyword_table()
-        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
-            return
-        else:
-            index_struct_dict = dataset_keyword_table.keyword_table_dict
-            index_struct: KeywordTable = load_index_struct_from_dict(index_struct_dict)
-
-        # create index
-        index = GPTJIEBAKeywordTableIndex(
-            index_struct=index_struct,
-            docstore=EmptyDocumentStore(),
-            service_context=service_context
-        )
-
-        for node_id in node_ids:
-            index.delete(node_id)
-
-        index_struct_dict = index.index_struct.to_dict()
-
-        if not dataset_keyword_table:
-            dataset_keyword_table = DatasetKeywordTable(
-                dataset_id=self._dataset.id,
-                keyword_table=json.dumps(index_struct_dict)
-            )
-            db.session.add(dataset_keyword_table)
-        else:
-            dataset_keyword_table.keyword_table = json.dumps(index_struct_dict)
-
-        db.session.commit()
-
-    @property
-    def query_index(self) -> Optional[BaseGPTKeywordTableIndex]:
-        docstore = DatesetDocumentStore(
-            dataset=self._dataset,
-            user_id=self._dataset.created_by,
-            embedding_model_name="text-embedding-ada-002"
-        )
-
-        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
-
-        dataset_keyword_table = self.get_keyword_table()
-        if not dataset_keyword_table or not dataset_keyword_table.keyword_table_dict:
-            return None
-
-        index_struct: KeywordTable = load_index_struct_from_dict(dataset_keyword_table.keyword_table_dict)
-
-        return GPTJIEBAKeywordTableIndex(index_struct=index_struct, docstore=docstore, service_context=service_context)
-
-    def get_keyword_table(self):
-        dataset_keyword_table = self._dataset.dataset_keyword_table
-        if dataset_keyword_table:
-            return dataset_keyword_table
-        return None
-
-    def update_segment_keywords(self, node_id: str, keywords: List[str]):
-        document_segment = db.session.query(DocumentSegment).filter(DocumentSegment.index_node_id == node_id).first()
-        if document_segment:
-            document_segment.keywords = keywords
-            db.session.commit()

api/core/index/keyword_table_index/jieba_keyword_table_handler.py

@@ -0,0 +1,33 @@
+import re
+from typing import Set
+
+import jieba
+from jieba.analyse import default_tfidf
+
+from core.index.keyword_table_index.stopwords import STOPWORDS
+
+
+class JiebaKeywordTableHandler:
+
+    def __init__(self):
+        default_tfidf.stop_words = STOPWORDS
+
+    def extract_keywords(self, text: str, max_keywords_per_chunk: int = 10) -> Set[str]:
+        """Extract keywords with JIEBA tfidf."""
+        keywords = jieba.analyse.extract_tags(
+            sentence=text,
+            topK=max_keywords_per_chunk,
+        )
+
+        return set(self._expand_tokens_with_subtokens(keywords))
+
+    def _expand_tokens_with_subtokens(self, tokens: Set[str]) -> Set[str]:
+        """Get subtokens from a list of tokens., filtering for stopwords."""
+        results = set()
+        for token in tokens:
+            results.add(token)
+            sub_tokens = re.findall(r"\w+", token)
+            if len(sub_tokens) > 1:
+                results.update({w for w in sub_tokens if w not in list(STOPWORDS)})
+
+        return results

api/core/index/keyword_table_index/keyword_table_index.py

@@ -0,0 +1,238 @@
+import json
+from collections import defaultdict
+from typing import Any, List, Optional, Dict
+
+from langchain.schema import Document, BaseRetriever
+from pydantic import BaseModel, Field, Extra
+
+from core.index.base import BaseIndex
+from core.index.keyword_table_index.jieba_keyword_table_handler import JiebaKeywordTableHandler
+from extensions.ext_database import db
+from models.dataset import Dataset, DocumentSegment, DatasetKeywordTable
+
+
+class KeywordTableConfig(BaseModel):
+    max_keywords_per_chunk: int = 10
+
+
+class KeywordTableIndex(BaseIndex):
+    def __init__(self, dataset: Dataset, config: KeywordTableConfig = KeywordTableConfig()):
+        super().__init__(dataset)
+        self._config = config
+
+    def create(self, texts: list[Document], **kwargs) -> BaseIndex:
+        keyword_table_handler = JiebaKeywordTableHandler()
+        keyword_table = {}
+        for text in texts:
+            keywords = keyword_table_handler.extract_keywords(text.page_content, self._config.max_keywords_per_chunk)
+            self._update_segment_keywords(text.metadata['doc_id'], list(keywords))
+            keyword_table = self._add_text_to_keyword_table(keyword_table, text.metadata['doc_id'], list(keywords))
+
+        dataset_keyword_table = DatasetKeywordTable(
+            dataset_id=self.dataset.id,
+            keyword_table=json.dumps({
+                '__type__': 'keyword_table',
+                '__data__': {
+                    "index_id": self.dataset.id,
+                    "summary": None,
+                    "table": {}
+                }
+            }, cls=SetEncoder)
+        )
+        db.session.add(dataset_keyword_table)
+        db.session.commit()
+
+        self._save_dataset_keyword_table(keyword_table)
+
+        return self
+
+    def add_texts(self, texts: list[Document], **kwargs):
+        keyword_table_handler = JiebaKeywordTableHandler()
+
+        keyword_table = self._get_dataset_keyword_table()
+        for text in texts:
+            keywords = keyword_table_handler.extract_keywords(text.page_content, self._config.max_keywords_per_chunk)
+            self._update_segment_keywords(text.metadata['doc_id'], list(keywords))
+            keyword_table = self._add_text_to_keyword_table(keyword_table, text.metadata['doc_id'], list(keywords))
+
+        self._save_dataset_keyword_table(keyword_table)
+
+    def text_exists(self, id: str) -> bool:
+        keyword_table = self._get_dataset_keyword_table()
+        return id in set.union(*keyword_table.values())
+
+    def delete_by_ids(self, ids: list[str]) -> None:
+        keyword_table = self._get_dataset_keyword_table()
+        keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids)
+
+        self._save_dataset_keyword_table(keyword_table)
+
+    def delete_by_document_id(self, document_id: str):
+        # get segment ids by document_id
+        segments = db.session.query(DocumentSegment).filter(
+            DocumentSegment.dataset_id == self.dataset.id,
+            DocumentSegment.document_id == document_id
+        ).all()
+
+        ids = [segment.id for segment in segments]
+
+        keyword_table = self._get_dataset_keyword_table()
+        keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids)
+
+        self._save_dataset_keyword_table(keyword_table)
+
+    def get_retriever(self, **kwargs: Any) -> BaseRetriever:
+        return KeywordTableRetriever(index=self, **kwargs)
+
+    def search(
+            self, query: str,
+            **kwargs: Any
+    ) -> List[Document]:
+        keyword_table = self._get_dataset_keyword_table()
+
+        search_kwargs = kwargs.get('search_kwargs') if kwargs.get('search_kwargs') else {}
+        k = search_kwargs.get('k') if search_kwargs.get('k') else 4
+
+        sorted_chunk_indices = self._retrieve_ids_by_query(keyword_table, query, k)
+
+        documents = []
+        for chunk_index in sorted_chunk_indices:
+            segment = db.session.query(DocumentSegment).filter(
+                DocumentSegment.dataset_id == self.dataset.id,
+                DocumentSegment.index_node_id == chunk_index
+            ).first()
+
+            if segment:
+                documents.append(Document(
+                    page_content=segment.content,
+                    metadata={
+                        "doc_id": chunk_index,
+                        "document_id": segment.document_id,
+                        "dataset_id": segment.dataset_id,
+                    }
+                ))
+
+        return documents
+
+    def delete(self) -> None:
+        dataset_keyword_table = self.dataset.dataset_keyword_table
+        if dataset_keyword_table:
+            db.session.delete(dataset_keyword_table)
+            db.session.commit()
+
+    def _save_dataset_keyword_table(self, keyword_table):
+        keyword_table_dict = {
+            '__type__': 'keyword_table',
+            '__data__': {
+                "index_id": self.dataset.id,
+                "summary": None,
+                "table": keyword_table
+            }
+        }
+        self.dataset.dataset_keyword_table.keyword_table = json.dumps(keyword_table_dict, cls=SetEncoder)
+        db.session.commit()
+
+    def _get_dataset_keyword_table(self) -> Optional[dict]:
+        dataset_keyword_table = self.dataset.dataset_keyword_table
+        if dataset_keyword_table:
+            if dataset_keyword_table.keyword_table_dict:
+                return dataset_keyword_table.keyword_table_dict['__data__']['table']
+        else:
+            dataset_keyword_table = DatasetKeywordTable(
+                dataset_id=self.dataset.id,
+                keyword_table=json.dumps({
+                    '__type__': 'keyword_table',
+                    '__data__': {
+                        "index_id": self.dataset.id,
+                        "summary": None,
+                        "table": {}
+                    }
+                }, cls=SetEncoder)
+            )
+            db.session.add(dataset_keyword_table)
+            db.session.commit()
+
+        return {}
+
+    def _add_text_to_keyword_table(self, keyword_table: dict, id: str, keywords: list[str]) -> dict:
+        for keyword in keywords:
+            if keyword not in keyword_table:
+                keyword_table[keyword] = set()
+            keyword_table[keyword].add(id)
+        return keyword_table
+
+    def _delete_ids_from_keyword_table(self, keyword_table: dict, ids: list[str]) -> dict:
+        # get set of ids that correspond to node
+        node_idxs_to_delete = set(ids)
+
+        # delete node_idxs from keyword to node idxs mapping
+        keywords_to_delete = set()
+        for keyword, node_idxs in keyword_table.items():
+            if node_idxs_to_delete.intersection(node_idxs):
+                keyword_table[keyword] = node_idxs.difference(
+                    node_idxs_to_delete
+                )
+                if not keyword_table[keyword]:
+                    keywords_to_delete.add(keyword)
+
+        for keyword in keywords_to_delete:
+            del keyword_table[keyword]
+
+        return keyword_table
+
+    def _retrieve_ids_by_query(self, keyword_table: dict, query: str, k: int = 4):
+        keyword_table_handler = JiebaKeywordTableHandler()
+        keywords = keyword_table_handler.extract_keywords(query)
+
+        # go through text chunks in order of most matching keywords
+        chunk_indices_count: Dict[str, int] = defaultdict(int)
+        keywords = [keyword for keyword in keywords if keyword in set(keyword_table.keys())]
+        for keyword in keywords:
+            for node_id in keyword_table[keyword]:
+                chunk_indices_count[node_id] += 1
+
+        sorted_chunk_indices = sorted(
+            list(chunk_indices_count.keys()),
+            key=lambda x: chunk_indices_count[x],
+            reverse=True,
+        )
+
+        return sorted_chunk_indices[: k]
+
+    def _update_segment_keywords(self, node_id: str, keywords: List[str]):
+        document_segment = db.session.query(DocumentSegment).filter(DocumentSegment.index_node_id == node_id).first()
+        if document_segment:
+            document_segment.keywords = keywords
+            db.session.commit()
+
+
+class KeywordTableRetriever(BaseRetriever, BaseModel):
+    index: KeywordTableIndex
+    search_kwargs: dict = Field(default_factory=dict)
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    def get_relevant_documents(self, query: str) -> List[Document]:
+        """Get documents relevant for a query.
+
+        Args:
+            query: string to find relevant documents for
+
+        Returns:
+            List of relevant documents
+        """
+        return self.index.search(query, **self.search_kwargs)
+
+    async def aget_relevant_documents(self, query: str) -> List[Document]:
+        raise NotImplementedError("KeywordTableRetriever does not support async")
+
+
+class SetEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, set):
+            return list(obj)
+        return super().default(obj)

api/core/index/query/synthesizer.py

@@ -1,79 +0,0 @@
-from typing import (
-    Any,
-    Dict,
-    Optional, Sequence,
-)
-
-from llama_index.indices.response.response_synthesis import ResponseSynthesizer
-from llama_index.indices.response.response_builder import ResponseMode, BaseResponseBuilder, get_response_builder
-from llama_index.indices.service_context import ServiceContext
-from llama_index.optimization.optimizer import BaseTokenUsageOptimizer
-from llama_index.prompts.prompts import (
-    QuestionAnswerPrompt,
-    RefinePrompt,
-    SimpleInputPrompt,
-)
-from llama_index.types import RESPONSE_TEXT_TYPE
-
-
-class EnhanceResponseSynthesizer(ResponseSynthesizer):
-    @classmethod
-    def from_args(
-            cls,
-            service_context: ServiceContext,
-            streaming: bool = False,
-            use_async: bool = False,
-            text_qa_template: Optional[QuestionAnswerPrompt] = None,
-            refine_template: Optional[RefinePrompt] = None,
-            simple_template: Optional[SimpleInputPrompt] = None,
-            response_mode: ResponseMode = ResponseMode.DEFAULT,
-            response_kwargs: Optional[Dict] = None,
-            optimizer: Optional[BaseTokenUsageOptimizer] = None,
-    ) -> "ResponseSynthesizer":
-        response_builder: Optional[BaseResponseBuilder] = None
-        if response_mode != ResponseMode.NO_TEXT:
-            if response_mode == 'no_synthesizer':
-                response_builder = NoSynthesizer(
-                    service_context=service_context,
-                    simple_template=simple_template,
-                    streaming=streaming,
-                )
-            else:
-                response_builder = get_response_builder(
-                    service_context,
-                    text_qa_template,
-                    refine_template,
-                    simple_template,
-                    response_mode,
-                    use_async=use_async,
-                    streaming=streaming,
-                )
-        return cls(response_builder, response_mode, response_kwargs, optimizer)
-
-
-class NoSynthesizer(BaseResponseBuilder):
-    def __init__(
-            self,
-            service_context: ServiceContext,
-            simple_template: Optional[SimpleInputPrompt] = None,
-            streaming: bool = False,
-    ) -> None:
-        super().__init__(service_context, streaming)
-
-    async def aget_response(
-            self,
-            query_str: str,
-            text_chunks: Sequence[str],
-            prev_response: Optional[str] = None,
-            **response_kwargs: Any,
-    ) -> RESPONSE_TEXT_TYPE:
-        return "\n".join(text_chunks)
-
-    def get_response(
-            self,
-            query_str: str,
-            text_chunks: Sequence[str],
-            prev_response: Optional[str] = None,
-            **response_kwargs: Any,
-    ) -> RESPONSE_TEXT_TYPE:
-        return "\n".join(text_chunks)

api/core/index/readers/html_parser.py

@@ -1,22 +0,0 @@
-from pathlib import Path
-from typing import Dict
-
-from bs4 import BeautifulSoup
-from llama_index.readers.file.base_parser import BaseParser
-
-
-class HTMLParser(BaseParser):
-    """HTML parser."""
-
-    def _init_parser(self) -> Dict:
-        """Init parser."""
-        return {}
-
-    def parse_file(self, file: Path, errors: str = "ignore") -> str:
-        """Parse file."""
-        with open(file, "rb") as fp:
-            soup = BeautifulSoup(fp, 'html.parser')
-            text = soup.get_text()
-            text = text.strip() if text else ''
-
-        return text

api/core/index/readers/markdown_parser.py

@@ -1,111 +0,0 @@
-"""Markdown parser.
-
-Contains parser for md files.
-
-"""
-import re
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
-
-from llama_index.readers.file.base_parser import BaseParser
-
-
-class MarkdownParser(BaseParser):
-    """Markdown parser.
-
-    Extract text from markdown files.
-    Returns dictionary with keys as headers and values as the text between headers.
-
-    """
-
-    def __init__(
-        self,
-        *args: Any,
-        remove_hyperlinks: bool = True,
-        remove_images: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        """Init params."""
-        super().__init__(*args, **kwargs)
-        self._remove_hyperlinks = remove_hyperlinks
-        self._remove_images = remove_images
-
-    def markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
-        """Convert a markdown file to a dictionary.
-
-        The keys are the headers and the values are the text under each header.
-
-        """
-        markdown_tups: List[Tuple[Optional[str], str]] = []
-        lines = markdown_text.split("\n")
-
-        current_header = None
-        current_text = ""
-
-        for line in lines:
-            header_match = re.match(r"^#+\s", line)
-            if header_match:
-                if current_header is not None:
-                    markdown_tups.append((current_header, current_text))
-
-                current_header = line
-                current_text = ""
-            else:
-                current_text += line + "\n"
-        markdown_tups.append((current_header, current_text))
-
-        if current_header is not None:
-            # pass linting, assert keys are defined
-            markdown_tups = [
-                (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
-                for key, value in markdown_tups
-            ]
-        else:
-            markdown_tups = [
-                (key, re.sub("\n", "", value)) for key, value in markdown_tups
-            ]
-
-        return markdown_tups
-
-    def remove_images(self, content: str) -> str:
-        """Get a dictionary of a markdown file from its path."""
-        pattern = r"!{1}\[\[(.*)\]\]"
-        content = re.sub(pattern, "", content)
-        return content
-
-    def remove_hyperlinks(self, content: str) -> str:
-        """Get a dictionary of a markdown file from its path."""
-        pattern = r"\[(.*?)\]\((.*?)\)"
-        content = re.sub(pattern, r"\1", content)
-        return content
-
-    def _init_parser(self) -> Dict:
-        """Initialize the parser with the config."""
-        return {}
-
-    def parse_tups(
-        self, filepath: Path, errors: str = "ignore"
-    ) -> List[Tuple[Optional[str], str]]:
-        """Parse file into tuples."""
-        with open(filepath, "r", encoding="utf-8") as f:
-            content = f.read()
-        if self._remove_hyperlinks:
-            content = self.remove_hyperlinks(content)
-        if self._remove_images:
-            content = self.remove_images(content)
-        markdown_tups = self.markdown_to_tups(content)
-        return markdown_tups
-
-    def parse_file(
-        self, filepath: Path, errors: str = "ignore"
-    ) -> Union[str, List[str]]:
-        """Parse file into string."""
-        tups = self.parse_tups(filepath, errors=errors)
-        results = []
-        # TODO: don't include headers right now
-        for header, value in tups:
-            if header is None:
-                results.append(value)
-            else:
-                results.append(f"\n\n{header}\n{value}")
-        return results

api/core/index/readers/pdf_parser.py

@@ -1,56 +0,0 @@
-from pathlib import Path
-from typing import Dict
-
-from flask import current_app
-from llama_index.readers.file.base_parser import BaseParser
-from pypdf import PdfReader
-
-from extensions.ext_storage import storage
-from models.model import UploadFile
-
-
-class PDFParser(BaseParser):
-    """PDF parser."""
-
-    def _init_parser(self) -> Dict:
-        """Init parser."""
-        return {}
-
-    def parse_file(self, file: Path, errors: str = "ignore") -> str:
-        """Parse file."""
-        if not current_app.config.get('PDF_PREVIEW', True):
-            return ''
-
-        plaintext_file_key = ''
-        plaintext_file_exists = False
-        if self._parser_config and 'upload_file' in self._parser_config and self._parser_config['upload_file']:
-            upload_file: UploadFile = self._parser_config['upload_file']
-            if upload_file.hash:
-                plaintext_file_key = 'upload_files/' + upload_file.tenant_id + '/' + upload_file.hash + '.plaintext'
-                try:
-                    text = storage.load(plaintext_file_key).decode('utf-8')
-                    plaintext_file_exists = True
-                    return text
-                except FileNotFoundError:
-                    pass
-
-        text_list = []
-        with open(file, "rb") as fp:
-            # Create a PDF object
-            pdf = PdfReader(fp)
-
-            # Get the number of pages in the PDF document
-            num_pages = len(pdf.pages)
-
-            # Iterate over every page
-            for page in range(num_pages):
-                # Extract the text from the page
-                page_text = pdf.pages[page].extract_text()
-                text_list.append(page_text)
-        text = "\n".join(text_list)
-
-        # save plaintext file for caching
-        if not plaintext_file_exists and plaintext_file_key:
-            storage.save(plaintext_file_key, text.encode('utf-8'))
-
-        return text

api/core/index/readers/xlsx_parser.py

@@ -1,31 +0,0 @@
-from pathlib import Path
-import json
-from typing import Dict
-from openpyxl import load_workbook
-
-from llama_index.readers.file.base_parser import BaseParser
-from flask import current_app
-
-
-class XLSXParser(BaseParser):
-    """XLSX parser."""
-
-    def _init_parser(self) -> Dict:
-        """Init parser"""
-        return {}
-
-    def parse_file(self, file: Path, errors: str = "ignore") -> str:
-        data = []
-        keys = []
-        with open(file, "r") as fp:
-            wb = load_workbook(filename=file, read_only=True)
-            # loop over all sheets
-            for sheet in wb:
-                for row in sheet.iter_rows(values_only=True):
-                    if all(v is None for v in row):
-                        continue
-                    if keys == []:
-                        keys = list(map(str, row))
-                    else:
-                        data.append(json.dumps(dict(zip(keys, list(map(str, row)))), ensure_ascii=False))
-        return '\n\n'.join(data)

api/core/index/vector_index.py

@@ -1,136 +0,0 @@
-import json
-import logging
-from typing import List, Optional
-
-from llama_index.data_structs import Node
-from requests import ReadTimeout
-from sqlalchemy.exc import IntegrityError
-from tenacity import retry, stop_after_attempt, retry_if_exception_type
-
-from core.index.index_builder import IndexBuilder
-from core.vector_store.base import BaseGPTVectorStoreIndex
-from extensions.ext_vector_store import vector_store
-from extensions.ext_database import db
-from models.dataset import Dataset, Embedding
-
-
-class VectorIndex:
-
-    def __init__(self, dataset: Dataset):
-        self._dataset = dataset
-
-    def add_nodes(self, nodes: List[Node], duplicate_check: bool = False):
-        if not self._dataset.index_struct_dict:
-            index_id = "Vector_index_" + self._dataset.id.replace("-", "_")
-            self._dataset.index_struct = json.dumps(vector_store.to_index_struct(index_id))
-            db.session.commit()
-
-        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
-
-        index = vector_store.get_index(
-            service_context=service_context,
-            index_struct=self._dataset.index_struct_dict
-        )
-
-        if duplicate_check:
-            nodes = self._filter_duplicate_nodes(index, nodes)
-
-        embedding_queue_nodes = []
-        embedded_nodes = []
-        for node in nodes:
-            node_hash = node.doc_hash
-
-            # if node hash in cached embedding tables, use cached embedding
-            embedding = db.session.query(Embedding).filter_by(hash=node_hash).first()
-            if embedding:
-                node.embedding = embedding.get_embedding()
-                embedded_nodes.append(node)
-            else:
-                embedding_queue_nodes.append(node)
-
-        if embedding_queue_nodes:
-            embedding_results = index._get_node_embedding_results(
-                embedding_queue_nodes,
-                set(),
-            )
-
-            # pre embed nodes for cached embedding
-            for embedding_result in embedding_results:
-                node = embedding_result.node
-                node.embedding = embedding_result.embedding
-
-                try:
-                    embedding = Embedding(hash=node.doc_hash)
-                    embedding.set_embedding(node.embedding)
-                    db.session.add(embedding)
-                    db.session.commit()
-                except IntegrityError:
-                    db.session.rollback()
-                    continue
-                except:
-                    logging.exception('Failed to add embedding to db')
-                    continue
-
-                embedded_nodes.append(node)
-
-        self.index_insert_nodes(index, embedded_nodes)
-
-    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
-    def index_insert_nodes(self, index: BaseGPTVectorStoreIndex, nodes: List[Node]):
-        index.insert_nodes(nodes)
-
-    def del_nodes(self, node_ids: List[str]):
-        if not self._dataset.index_struct_dict:
-            return
-
-        service_context = IndexBuilder.get_fake_llm_service_context(tenant_id=self._dataset.tenant_id)
-
-        index = vector_store.get_index(
-            service_context=service_context,
-            index_struct=self._dataset.index_struct_dict
-        )
-
-        for node_id in node_ids:
-            self.index_delete_node(index, node_id)
-
-    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
-    def index_delete_node(self, index: BaseGPTVectorStoreIndex, node_id: str):
-        index.delete_node(node_id)
-
-    def del_doc(self, doc_id: str):
-        if not self._dataset.index_struct_dict:
-            return
-
-        service_context = IndexBuilder.get_fake_llm_service_context(tenant_id=self._dataset.tenant_id)
-
-        index = vector_store.get_index(
-            service_context=service_context,
-            index_struct=self._dataset.index_struct_dict
-        )
-
-        self.index_delete_doc(index, doc_id)
-
-    @retry(reraise=True, retry=retry_if_exception_type(ReadTimeout), stop=stop_after_attempt(3))
-    def index_delete_doc(self, index: BaseGPTVectorStoreIndex, doc_id: str):
-        index.delete(doc_id)
-
-    @property
-    def query_index(self) -> Optional[BaseGPTVectorStoreIndex]:
-        if not self._dataset.index_struct_dict:
-            return None
-
-        service_context = IndexBuilder.get_default_service_context(tenant_id=self._dataset.tenant_id)
-
-        return vector_store.get_index(
-            service_context=service_context,
-            index_struct=self._dataset.index_struct_dict
-        )
-
-    def _filter_duplicate_nodes(self, index: BaseGPTVectorStoreIndex, nodes: List[Node]) -> List[Node]:
-        for node in nodes:
-            node_id = node.doc_id
-            exists_duplicate_node = index.exists_by_node_id(node_id)
-            if exists_duplicate_node:
-                nodes.remove(node)
-
-        return nodes

api/core/index/vector_index/base.py

@@ -0,0 +1,175 @@
+import json
+import logging
+from abc import abstractmethod
+from typing import List, Any, cast
+
+from langchain.embeddings.base import Embeddings
+from langchain.schema import Document, BaseRetriever
+from langchain.vectorstores import VectorStore
+from weaviate import UnexpectedStatusCodeException
+
+from core.index.base import BaseIndex
+from extensions.ext_database import db
+from models.dataset import Dataset, DocumentSegment
+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
+
+    @abstractmethod
+    def get_index_name(self, dataset: Dataset) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_index_struct(self) -> dict:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_vector_store(self) -> VectorStore:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_vector_store_class(self) -> type:
+        raise NotImplementedError
+
+    def search(
+            self, query: str,
+            **kwargs: Any
+    ) -> List[Document]:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        search_type = kwargs.get('search_type') if kwargs.get('search_type') else 'similarity'
+        search_kwargs = kwargs.get('search_kwargs') if kwargs.get('search_kwargs') else {}
+
+        if search_type == 'similarity_score_threshold':
+            score_threshold = search_kwargs.get("score_threshold")
+            if (score_threshold is None) or (not isinstance(score_threshold, float)):
+                search_kwargs['score_threshold'] = .0
+
+            docs_with_similarity = vector_store.similarity_search_with_relevance_scores(
+                query, **search_kwargs
+            )
+
+            docs = []
+            for doc, similarity in docs_with_similarity:
+                doc.metadata['score'] = similarity
+                docs.append(doc)
+
+            return docs
+
+        # similarity k
+        # mmr k, fetch_k, lambda_mult
+        # similarity_score_threshold k
+        return vector_store.as_retriever(
+            search_type=search_type,
+            search_kwargs=search_kwargs
+        ).get_relevant_documents(query)
+
+    def get_retriever(self, **kwargs: Any) -> BaseRetriever:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        return vector_store.as_retriever(**kwargs)
+
+    def add_texts(self, texts: list[Document], **kwargs):
+        if self._is_origin():
+            self.recreate_dataset(self.dataset)
+
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        if kwargs.get('duplicate_check', False):
+            texts = self._filter_duplicate_texts(texts)
+
+        uuids = self._get_uuids(texts)
+        vector_store.add_documents(texts, uuids=uuids)
+
+    def text_exists(self, id: str) -> bool:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        return vector_store.text_exists(id)
+
+    def delete_by_ids(self, ids: list[str]) -> None:
+        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)
+
+        for node_id in ids:
+            vector_store.del_text(node_id)
+
+    def delete(self) -> None:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        vector_store.delete()
+
+    def _is_origin(self):
+        return False
+
+    def recreate_dataset(self, dataset: Dataset):
+        logging.info(f"Recreating 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)
+
+        origin_index_struct = self.dataset.index_struct[:]
+        self.dataset.index_struct = None
+
+        if documents:
+            try:
+                self.create(documents)
+            except Exception as e:
+                self.dataset.index_struct = origin_index_struct
+                raise e
+
+            dataset.index_struct = json.dumps(self.to_index_struct())
+
+        db.session.commit()
+
+        self.dataset = dataset
+        logging.info(f"Dataset {dataset.id} recreate successfully.")

api/core/index/vector_index/qdrant_vector_index.py

@@ -0,0 +1,116 @@
+import os
+from typing import Optional, Any, List, cast
+
+import qdrant_client
+from langchain.embeddings.base import Embeddings
+from langchain.schema import Document, BaseRetriever
+from langchain.vectorstores import VectorStore
+from pydantic import BaseModel
+
+from core.index.base import BaseIndex
+from core.index.vector_index.base import BaseVectorIndex
+from core.vector_store.qdrant_vector_store import QdrantVectorStore
+from models.dataset import Dataset
+
+
+class QdrantConfig(BaseModel):
+    endpoint: str
+    api_key: Optional[str]
+    root_path: Optional[str]
+    
+    def to_qdrant_params(self):
+        if self.endpoint and self.endpoint.startswith('path:'):
+            path = self.endpoint.replace('path:', '')
+            if not os.path.isabs(path):
+                path = os.path.join(self.root_path, path)
+
+            return {
+                'path': path
+            }
+        else:
+            return {
+                'url': self.endpoint,
+                'api_key': self.api_key,
+            }
+
+
+class QdrantVectorIndex(BaseVectorIndex):
+    def __init__(self, dataset: Dataset, config: QdrantConfig, embeddings: Embeddings):
+        super().__init__(dataset, embeddings)
+        self._client_config = config
+
+    def get_type(self) -> str:
+        return 'qdrant'
+
+    def get_index_name(self, dataset: Dataset) -> str:
+        if self.dataset.index_struct_dict:
+            return self.dataset.index_struct_dict['vector_store']['collection_name']
+
+        dataset_id = dataset.id
+        return "Index_" + dataset_id.replace("-", "_")
+
+    def to_index_struct(self) -> dict:
+        return {
+            "type": self.get_type(),
+            "vector_store": {"collection_name": self.get_index_name(self.dataset)}
+        }
+
+    def create(self, texts: list[Document], **kwargs) -> BaseIndex:
+        uuids = self._get_uuids(texts)
+        self._vector_store = QdrantVectorStore.from_documents(
+            texts,
+            self._embeddings,
+            collection_name=self.get_index_name(self.dataset),
+            ids=uuids,
+            content_payload_key='text',
+            **self._client_config.to_qdrant_params()
+        )
+
+        return self
+
+    def _get_vector_store(self) -> VectorStore:
+        """Only for created index."""
+        if self._vector_store:
+            return self._vector_store
+        
+        client = qdrant_client.QdrantClient(
+            **self._client_config.to_qdrant_params()
+        )
+
+        return QdrantVectorStore(
+            client=client,
+            collection_name=self.get_index_name(self.dataset),
+            embeddings=self._embeddings,
+            content_payload_key='text'
+        )
+
+    def _get_vector_store_class(self) -> type:
+        return QdrantVectorStore
+
+    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)
+
+        from qdrant_client.http import models
+
+        vector_store.del_texts(models.Filter(
+            must=[
+                models.FieldCondition(
+                    key="metadata.document_id",
+                    match=models.MatchValue(value=document_id),
+                ),
+            ],
+        ))
+
+    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_'):
+                # original class_prefix
+                return True
+
+        return False

api/core/index/vector_index/vector_index.py

@@ -0,0 +1,69 @@
+import json
+
+from flask import current_app
+from langchain.embeddings.base import Embeddings
+
+from core.index.vector_index.base import BaseVectorIndex
+from extensions.ext_database import db
+from models.dataset import Dataset, Document
+
+
+class VectorIndex:
+    def __init__(self, dataset: Dataset, config: dict, embeddings: Embeddings):
+        self._dataset = dataset
+        self._embeddings = embeddings
+        self._vector_index = self._init_vector_index(dataset, config, embeddings)
+
+    def _init_vector_index(self, dataset: Dataset, config: dict, embeddings: Embeddings) -> BaseVectorIndex:
+        vector_type = config.get('VECTOR_STORE')
+
+        if self._dataset.index_struct_dict:
+            vector_type = self._dataset.index_struct_dict['type']
+
+        if not vector_type:
+            raise ValueError(f"Vector store must be specified.")
+
+        if vector_type == "weaviate":
+            from core.index.vector_index.weaviate_vector_index import WeaviateVectorIndex, WeaviateConfig
+
+            return WeaviateVectorIndex(
+                dataset=dataset,
+                config=WeaviateConfig(
+                    endpoint=config.get('WEAVIATE_ENDPOINT'),
+                    api_key=config.get('WEAVIATE_API_KEY'),
+                    batch_size=int(config.get('WEAVIATE_BATCH_SIZE'))
+                ),
+                embeddings=embeddings
+            )
+        elif vector_type == "qdrant":
+            from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+            return QdrantVectorIndex(
+                dataset=dataset,
+                config=QdrantConfig(
+                    endpoint=config.get('QDRANT_URL'),
+                    api_key=config.get('QDRANT_API_KEY'),
+                    root_path=current_app.root_path
+                ),
+                embeddings=embeddings
+            )
+        else:
+            raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.")
+
+    def add_texts(self, texts: list[Document], **kwargs):
+        if not self._dataset.index_struct_dict:
+            self._vector_index.create(texts, **kwargs)
+            self._dataset.index_struct = json.dumps(self._vector_index.to_index_struct())
+            db.session.commit()
+            return
+
+        self._vector_index.add_texts(texts, **kwargs)
+
+    def __getattr__(self, name):
+        if self._vector_index is not None:
+            method = getattr(self._vector_index, name)
+            if callable(method):
+                return method
+
+        raise AttributeError(f"'VectorIndex' object has no attribute '{name}'")
+

api/core/index/vector_index/weaviate_vector_index.py

@@ -0,0 +1,136 @@
+from typing import Optional, cast
+
+import requests
+import weaviate
+from langchain.embeddings.base import Embeddings
+from langchain.schema import Document, BaseRetriever
+from langchain.vectorstores import VectorStore
+from pydantic import BaseModel, root_validator
+
+from core.index.base import BaseIndex
+from core.index.vector_index.base import BaseVectorIndex
+from core.vector_store.weaviate_vector_store import WeaviateVectorStore
+from models.dataset import Dataset
+
+
+class WeaviateConfig(BaseModel):
+    endpoint: str
+    api_key: Optional[str]
+    batch_size: int = 100
+
+    @root_validator()
+    def validate_config(cls, values: dict) -> dict:
+        if not values['endpoint']:
+            raise ValueError("config WEAVIATE_ENDPOINT is required")
+        return values
+
+
+class WeaviateVectorIndex(BaseVectorIndex):
+    def __init__(self, dataset: Dataset, config: WeaviateConfig, embeddings: Embeddings):
+        super().__init__(dataset, embeddings)
+        self._client = self._init_client(config)
+
+    def _init_client(self, config: WeaviateConfig) -> weaviate.Client:
+        auth_config = weaviate.auth.AuthApiKey(api_key=config.api_key)
+
+        weaviate.connect.connection.has_grpc = False
+
+        try:
+            client = weaviate.Client(
+                url=config.endpoint,
+                auth_client_secret=auth_config,
+                timeout_config=(5, 60),
+                startup_period=None
+            )
+        except requests.exceptions.ConnectionError:
+            raise ConnectionError("Vector database connection error")
+
+        client.batch.configure(
+            # `batch_size` takes an `int` value to enable auto-batching
+            # (`None` is used for manual batching)
+            batch_size=config.batch_size,
+            # dynamically update the `batch_size` based on import speed
+            dynamic=True,
+            # `timeout_retries` takes an `int` value to retry on time outs
+            timeout_retries=3,
+        )
+
+        return client
+
+    def get_type(self) -> str:
+        return 'weaviate'
+
+    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 WeaviateVectorStore
+
+    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/indexing_runner.py

@@ -1,32 +1,36 @@
 import datetime
 import json
+import logging
 import re
-import tempfile
 import time
-from pathlib import Path
-from typing import Optional, List
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+import uuid
+from typing import Optional, List, cast
 
-from llama_index import SimpleDirectoryReader
-from llama_index.data_structs import Node
-from llama_index.data_structs.node_v2 import DocumentRelationship
-from llama_index.node_parser import SimpleNodeParser, NodeParser
-from llama_index.readers.file.base import DEFAULT_FILE_EXTRACTOR
-from llama_index.readers.file.markdown_parser import MarkdownParser
-from core.index.readers.xlsx_parser import XLSXParser
+from flask import current_app
+from flask_login import current_user
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.schema import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter, TextSplitter
+
+from core.data_loader.file_extractor import FileExtractor
+from core.data_loader.loader.notion import NotionLoader
 from core.docstore.dataset_docstore import DatesetDocumentStore
-from core.index.keyword_table_index import KeywordTableIndex
-from core.index.readers.html_parser import HTMLParser
-from core.index.readers.markdown_parser import MarkdownParser
-from core.index.readers.pdf_parser import PDFParser
-from core.index.spiltter.fixed_text_splitter import FixedRecursiveCharacterTextSplitter
-from core.index.vector_index import VectorIndex
+from core.embedding.cached_embedding import CacheEmbedding
+from core.index.index import IndexBuilder
+from core.index.keyword_table_index.keyword_table_index import KeywordTableIndex, KeywordTableConfig
+from core.index.vector_index.vector_index import VectorIndex
+from core.llm.error import ProviderTokenNotInitError
+from core.llm.llm_builder import LLMBuilder
+from core.spiltter.fixed_text_splitter import FixedRecursiveCharacterTextSplitter
 from core.llm.token_calculator import TokenCalculator
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from extensions.ext_storage import storage
-from models.dataset import Document, Dataset, DocumentSegment, DatasetProcessRule
+from libs import helper
+from models.dataset import Document as DatasetDocument
+from models.dataset import Dataset, DocumentSegment, DatasetProcessRule
 from models.model import UploadFile
+from models.source import DataSourceBinding
 
 
 class IndexingRunner:
@@ -35,231 +39,319 @@ class IndexingRunner:
         self.storage = storage
         self.embedding_model_name = embedding_model_name
 
-    def run(self, document: Document):
+    def run(self, dataset_documents: List[DatasetDocument]):
         """Run the indexing process."""
-        # get dataset
-        dataset = Dataset.query.filter_by(
-            id=document.dataset_id
-        ).first()
+        for dataset_document in dataset_documents:
+            try:
+                # get dataset
+                dataset = Dataset.query.filter_by(
+                    id=dataset_document.dataset_id
+                ).first()
 
-        if not dataset:
-            raise ValueError("no dataset found")
+                if not dataset:
+                    raise ValueError("no dataset found")
 
-        # load file
-        text_docs = self._load_data(document)
+                # load file
+                text_docs = self._load_data(dataset_document)
 
-        # get the process rule
-        processing_rule = db.session.query(DatasetProcessRule). \
-            filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
-            first()
+                # get the process rule
+                processing_rule = db.session.query(DatasetProcessRule). \
+                    filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id). \
+                    first()
 
-        # get node parser for splitting
-        node_parser = self._get_node_parser(processing_rule)
+                # get splitter
+                splitter = self._get_splitter(processing_rule)
 
-        # split to nodes
-        nodes = self._step_split(
-            text_docs=text_docs,
-            node_parser=node_parser,
-            dataset=dataset,
-            document=document,
-            processing_rule=processing_rule
-        )
+                # split to documents
+                documents = self._step_split(
+                    text_docs=text_docs,
+                    splitter=splitter,
+                    dataset=dataset,
+                    dataset_document=dataset_document,
+                    processing_rule=processing_rule
+                )
 
-        # build index
-        self._build_index(
-            dataset=dataset,
-            document=document,
-            nodes=nodes
-        )
+                # build index
+                self._build_index(
+                    dataset=dataset,
+                    dataset_document=dataset_document,
+                    documents=documents
+                )
+            except DocumentIsPausedException:
+                raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
+            except ProviderTokenNotInitError as e:
+                dataset_document.indexing_status = 'error'
+                dataset_document.error = str(e.description)
+                dataset_document.stopped_at = datetime.datetime.utcnow()
+                db.session.commit()
+            except Exception as e:
+                logging.exception("consume document failed")
+                dataset_document.indexing_status = 'error'
+                dataset_document.error = str(e)
+                dataset_document.stopped_at = datetime.datetime.utcnow()
+                db.session.commit()
 
-    def run_in_splitting_status(self, document: Document):
+    def run_in_splitting_status(self, dataset_document: DatasetDocument):
         """Run the indexing process when the index_status is splitting."""
-        # get dataset
-        dataset = Dataset.query.filter_by(
-            id=document.dataset_id
-        ).first()
+        try:
+            # get dataset
+            dataset = Dataset.query.filter_by(
+                id=dataset_document.dataset_id
+            ).first()
 
-        if not dataset:
-            raise ValueError("no dataset found")
+            if not dataset:
+                raise ValueError("no dataset found")
 
-        # get exist document_segment list and delete
-        document_segments = DocumentSegment.query.filter_by(
-            dataset_id=dataset.id,
-            document_id=document.id
-        ).all()
-        db.session.delete(document_segments)
-        db.session.commit()
-        # load file
-        text_docs = self._load_data(document)
+            # get exist document_segment list and delete
+            document_segments = DocumentSegment.query.filter_by(
+                dataset_id=dataset.id,
+                document_id=dataset_document.id
+            ).all()
 
-        # get the process rule
-        processing_rule = db.session.query(DatasetProcessRule). \
-            filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
-            first()
+            db.session.delete(document_segments)
+            db.session.commit()
 
-        # get node parser for splitting
-        node_parser = self._get_node_parser(processing_rule)
+            # load file
+            text_docs = self._load_data(dataset_document)
 
-        # split to nodes
-        nodes = self._step_split(
-            text_docs=text_docs,
-            node_parser=node_parser,
-            dataset=dataset,
-            document=document,
-            processing_rule=processing_rule
-        )
+            # get the process rule
+            processing_rule = db.session.query(DatasetProcessRule). \
+                filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id). \
+                first()
 
-        # build index
-        self._build_index(
-            dataset=dataset,
-            document=document,
-            nodes=nodes
-        )
+            # get splitter
+            splitter = self._get_splitter(processing_rule)
 
-    def run_in_indexing_status(self, document: Document):
+            # split to documents
+            documents = self._step_split(
+                text_docs=text_docs,
+                splitter=splitter,
+                dataset=dataset,
+                dataset_document=dataset_document,
+                processing_rule=processing_rule
+            )
+
+            # build index
+            self._build_index(
+                dataset=dataset,
+                dataset_document=dataset_document,
+                documents=documents
+            )
+        except DocumentIsPausedException:
+            raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
+        except ProviderTokenNotInitError as e:
+            dataset_document.indexing_status = 'error'
+            dataset_document.error = str(e.description)
+            dataset_document.stopped_at = datetime.datetime.utcnow()
+            db.session.commit()
+        except Exception as e:
+            logging.exception("consume document failed")
+            dataset_document.indexing_status = 'error'
+            dataset_document.error = str(e)
+            dataset_document.stopped_at = datetime.datetime.utcnow()
+            db.session.commit()
+
+    def run_in_indexing_status(self, dataset_document: DatasetDocument):
         """Run the indexing process when the index_status is indexing."""
-        # get dataset
-        dataset = Dataset.query.filter_by(
-            id=document.dataset_id
-        ).first()
+        try:
+            # get dataset
+            dataset = Dataset.query.filter_by(
+                id=dataset_document.dataset_id
+            ).first()
 
-        if not dataset:
-            raise ValueError("no dataset found")
+            if not dataset:
+                raise ValueError("no dataset found")
 
-        # get exist document_segment list and delete
-        document_segments = DocumentSegment.query.filter_by(
-            dataset_id=dataset.id,
-            document_id=document.id
-        ).all()
-        nodes = []
-        if document_segments:
-            for document_segment in document_segments:
-                # transform segment to node
-                if document_segment.status != "completed":
-                    relationships = {
-                        DocumentRelationship.SOURCE: document_segment.document_id,
-                    }
+            # get exist document_segment list and delete
+            document_segments = DocumentSegment.query.filter_by(
+                dataset_id=dataset.id,
+                document_id=dataset_document.id
+            ).all()
 
-                    previous_segment = document_segment.previous_segment
-                    if previous_segment:
-                        relationships[DocumentRelationship.PREVIOUS] = previous_segment.index_node_id
+            documents = []
+            if document_segments:
+                for document_segment in document_segments:
+                    # transform segment to node
+                    if document_segment.status != "completed":
+                        document = Document(
+                            page_content=document_segment.content,
+                            metadata={
+                                "doc_id": document_segment.index_node_id,
+                                "doc_hash": document_segment.index_node_hash,
+                                "document_id": document_segment.document_id,
+                                "dataset_id": document_segment.dataset_id,
+                            }
+                        )
 
-                    next_segment = document_segment.next_segment
-                    if next_segment:
-                        relationships[DocumentRelationship.NEXT] = next_segment.index_node_id
-                    node = Node(
-                        doc_id=document_segment.index_node_id,
-                        doc_hash=document_segment.index_node_hash,
-                        text=document_segment.content,
-                        extra_info=None,
-                        node_info=None,
-                        relationships=relationships
-                    )
-                    nodes.append(node)
+                        documents.append(document)
 
-        # build index
-        self._build_index(
-            dataset=dataset,
-            document=document,
-            nodes=nodes
-        )
+            # build index
+            self._build_index(
+                dataset=dataset,
+                dataset_document=dataset_document,
+                documents=documents
+            )
+        except DocumentIsPausedException:
+            raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
+        except ProviderTokenNotInitError as e:
+            dataset_document.indexing_status = 'error'
+            dataset_document.error = str(e.description)
+            dataset_document.stopped_at = datetime.datetime.utcnow()
+            db.session.commit()
+        except Exception as e:
+            logging.exception("consume document failed")
+            dataset_document.indexing_status = 'error'
+            dataset_document.error = str(e)
+            dataset_document.stopped_at = datetime.datetime.utcnow()
+            db.session.commit()
 
-    def indexing_estimate(self, file_detail: UploadFile, tmp_processing_rule: dict) -> dict:
+    def file_indexing_estimate(self, file_details: List[UploadFile], tmp_processing_rule: dict) -> dict:
         """
         Estimate the indexing for the document.
         """
-        # load data from file
-        text_docs = self._load_data_from_file(file_detail)
-
-        processing_rule = DatasetProcessRule(
-            mode=tmp_processing_rule["mode"],
-            rules=json.dumps(tmp_processing_rule["rules"])
-        )
-
-        # get node parser for splitting
-        node_parser = self._get_node_parser(processing_rule)
-
-        # split to nodes
-        nodes = self._split_to_nodes(
-            text_docs=text_docs,
-            node_parser=node_parser,
-            processing_rule=processing_rule
-        )
-
         tokens = 0
         preview_texts = []
-        for node in nodes:
-            if len(preview_texts) < 5:
-                preview_texts.append(node.get_text())
+        total_segments = 0
+        for file_detail in file_details:
+            # load data from file
+            text_docs = FileExtractor.load(file_detail)
 
-            tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
+            processing_rule = DatasetProcessRule(
+                mode=tmp_processing_rule["mode"],
+                rules=json.dumps(tmp_processing_rule["rules"])
+            )
+
+            # get splitter
+            splitter = self._get_splitter(processing_rule)
+
+            # split to documents
+            documents = self._split_to_documents(
+                text_docs=text_docs,
+                splitter=splitter,
+                processing_rule=processing_rule
+            )
+            total_segments += len(documents)
+            for document in documents:
+                if len(preview_texts) < 5:
+                    preview_texts.append(document.page_content)
+
+                tokens += TokenCalculator.get_num_tokens(self.embedding_model_name,
+                                                         self.filter_string(document.page_content))
 
         return {
-            "total_segments": len(nodes),
+            "total_segments": total_segments,
             "tokens": tokens,
             "total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
             "currency": TokenCalculator.get_currency(self.embedding_model_name),
             "preview": preview_texts
         }
 
-    def _load_data(self, document: Document) -> List[Document]:
+    def notion_indexing_estimate(self, notion_info_list: list, tmp_processing_rule: dict) -> dict:
+        """
+        Estimate the indexing for the document.
+        """
+        # load data from notion
+        tokens = 0
+        preview_texts = []
+        total_segments = 0
+        for notion_info in notion_info_list:
+            workspace_id = notion_info['workspace_id']
+            data_source_binding = DataSourceBinding.query.filter(
+                db.and_(
+                    DataSourceBinding.tenant_id == current_user.current_tenant_id,
+                    DataSourceBinding.provider == 'notion',
+                    DataSourceBinding.disabled == False,
+                    DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
+                )
+            ).first()
+            if not data_source_binding:
+                raise ValueError('Data source binding not found.')
+
+            for page in notion_info['pages']:
+                loader = NotionLoader(
+                    notion_access_token=data_source_binding.access_token,
+                    notion_workspace_id=workspace_id,
+                    notion_obj_id=page['page_id'],
+                    notion_page_type=page['type']
+                )
+                documents = loader.load()
+
+                processing_rule = DatasetProcessRule(
+                    mode=tmp_processing_rule["mode"],
+                    rules=json.dumps(tmp_processing_rule["rules"])
+                )
+
+                # get splitter
+                splitter = self._get_splitter(processing_rule)
+
+                # split to documents
+                documents = self._split_to_documents(
+                    text_docs=documents,
+                    splitter=splitter,
+                    processing_rule=processing_rule
+                )
+                total_segments += len(documents)
+                for document in documents:
+                    if len(preview_texts) < 5:
+                        preview_texts.append(document.page_content)
+
+                    tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, document.page_content)
+
+        return {
+            "total_segments": total_segments,
+            "tokens": tokens,
+            "total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
+            "currency": TokenCalculator.get_currency(self.embedding_model_name),
+            "preview": preview_texts
+        }
+
+    def _load_data(self, dataset_document: DatasetDocument) -> List[Document]:
         # load file
-        if document.data_source_type != "upload_file":
+        if dataset_document.data_source_type not in ["upload_file", "notion_import"]:
             return []
 
-        data_source_info = document.data_source_info_dict
-        if not data_source_info or 'upload_file_id' not in data_source_info:
-            raise ValueError("no upload file found")
+        data_source_info = dataset_document.data_source_info_dict
+        text_docs = []
+        if dataset_document.data_source_type == 'upload_file':
+            if not data_source_info or 'upload_file_id' not in data_source_info:
+                raise ValueError("no upload file found")
 
-        file_detail = db.session.query(UploadFile). \
-            filter(UploadFile.id == data_source_info['upload_file_id']). \
-            one_or_none()
+            file_detail = db.session.query(UploadFile). \
+                filter(UploadFile.id == data_source_info['upload_file_id']). \
+                one_or_none()
 
-        text_docs = self._load_data_from_file(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()
 
         # update document status to splitting
         self._update_document_index_status(
-            document_id=document.id,
+            document_id=dataset_document.id,
             after_indexing_status="splitting",
             extra_update_params={
-                Document.file_id: file_detail.id,
-                Document.word_count: sum([len(text_doc.text) for text_doc in text_docs]),
-                Document.parsing_completed_at: datetime.datetime.utcnow()
+                DatasetDocument.word_count: sum([len(text_doc.page_content) for text_doc in text_docs]),
+                DatasetDocument.parsing_completed_at: datetime.datetime.utcnow()
             }
         )
 
         # replace doc id to document model id
+        text_docs = cast(List[Document], text_docs)
         for text_doc in text_docs:
             # remove invalid symbol
-            text_doc.text = self.filter_string(text_doc.get_text())
-            text_doc.doc_id = document.id
+            text_doc.page_content = self.filter_string(text_doc.page_content)
+            text_doc.metadata['document_id'] = dataset_document.id
+            text_doc.metadata['dataset_id'] = dataset_document.dataset_id
 
         return text_docs
 
     def filter_string(self, text):
-        pattern = re.compile('[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\x80-\xFF]')
-        return pattern.sub('', text)
+        text = re.sub(r'<\|', '<', text)
+        text = re.sub(r'\|>', '>', text)
+        text = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\x80-\xFF]', '', text)
+        return text
 
-    def _load_data_from_file(self, upload_file: UploadFile) -> List[Document]:
-        with tempfile.TemporaryDirectory() as temp_dir:
-            suffix = Path(upload_file.key).suffix
-            filepath = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
-            self.storage.download(upload_file.key, filepath)
-
-            file_extractor = DEFAULT_FILE_EXTRACTOR.copy()
-            file_extractor[".markdown"] = MarkdownParser()
-            file_extractor[".md"] = MarkdownParser()
-            file_extractor[".html"] = HTMLParser()
-            file_extractor[".htm"] = HTMLParser()
-            file_extractor[".pdf"] = PDFParser({'upload_file': upload_file})
-            file_extractor[".xlsx"] = XLSXParser()
-
-            loader = SimpleDirectoryReader(input_files=[filepath], file_extractor=file_extractor)
-            text_docs = loader.load_data()
-
-            return text_docs
-
-    def _get_node_parser(self, processing_rule: DatasetProcessRule) -> NodeParser:
+    def _get_splitter(self, processing_rule: DatasetProcessRule) -> TextSplitter:
         """
         Get the NodeParser object according to the processing rule.
         """
@@ -288,68 +380,83 @@ class IndexingRunner:
                 separators=["\n\n", "。", ".", " ", ""]
             )
 
-        return SimpleNodeParser(text_splitter=character_splitter, include_extra_info=True)
+        return character_splitter
 
-    def _step_split(self, text_docs: List[Document], node_parser: NodeParser,
-                    dataset: Dataset, document: Document, processing_rule: DatasetProcessRule) -> List[Node]:
+    def _step_split(self, text_docs: List[Document], splitter: TextSplitter,
+                    dataset: Dataset, dataset_document: DatasetDocument, processing_rule: DatasetProcessRule) \
+            -> List[Document]:
         """
-        Split the text documents into nodes and save them to the document segment.
+        Split the text documents into documents and save them to the document segment.
         """
-        nodes = self._split_to_nodes(
+        documents = self._split_to_documents(
             text_docs=text_docs,
-            node_parser=node_parser,
+            splitter=splitter,
             processing_rule=processing_rule
         )
 
         # save node to document segment
         doc_store = DatesetDocumentStore(
             dataset=dataset,
-            user_id=document.created_by,
+            user_id=dataset_document.created_by,
             embedding_model_name=self.embedding_model_name,
-            document_id=document.id
+            document_id=dataset_document.id
         )
 
-        doc_store.add_documents(nodes)
+        # add document segments
+        doc_store.add_documents(documents)
 
         # update document status to indexing
         cur_time = datetime.datetime.utcnow()
         self._update_document_index_status(
-            document_id=document.id,
+            document_id=dataset_document.id,
             after_indexing_status="indexing",
             extra_update_params={
-                Document.cleaning_completed_at: cur_time,
-                Document.splitting_completed_at: cur_time,
+                DatasetDocument.cleaning_completed_at: cur_time,
+                DatasetDocument.splitting_completed_at: cur_time,
             }
         )
 
         # update segment status to indexing
         self._update_segments_by_document(
-            document_id=document.id,
+            dataset_document_id=dataset_document.id,
             update_params={
                 DocumentSegment.status: "indexing",
                 DocumentSegment.indexing_at: datetime.datetime.utcnow()
             }
         )
 
-        return nodes
+        return documents
 
-    def _split_to_nodes(self, text_docs: List[Document], node_parser: NodeParser,
-                        processing_rule: DatasetProcessRule) -> List[Node]:
+    def _split_to_documents(self, text_docs: List[Document], splitter: TextSplitter,
+                            processing_rule: DatasetProcessRule) -> List[Document]:
         """
         Split the text documents into nodes.
         """
-        all_nodes = []
+        all_documents = []
         for text_doc in text_docs:
             # document clean
-            document_text = self._document_clean(text_doc.get_text(), processing_rule)
-            text_doc.text = document_text
+            document_text = self._document_clean(text_doc.page_content, processing_rule)
+            text_doc.page_content = document_text
 
             # parse document to nodes
-            nodes = node_parser.get_nodes_from_documents([text_doc])
-            nodes = [node for node in nodes if node.text is not None and node.text.strip()]
-            all_nodes.extend(nodes)
+            documents = splitter.split_documents([text_doc])
 
-        return all_nodes
+            split_documents = []
+            for document in documents:
+                if document.page_content is None or not document.page_content.strip():
+                    continue
+
+                doc_id = str(uuid.uuid4())
+                hash = helper.generate_text_hash(document.page_content)
+
+                document.metadata['doc_id'] = doc_id
+                document.metadata['doc_hash'] = hash
+
+                split_documents.append(document)
+
+            all_documents.extend(split_documents)
+
+        return all_documents
 
     def _document_clean(self, text: str, processing_rule: DatasetProcessRule) -> str:
         """
@@ -380,37 +487,38 @@ class IndexingRunner:
 
         return text
 
-    def _build_index(self, dataset: Dataset, document: Document, nodes: List[Node]) -> None:
+    def _build_index(self, dataset: Dataset, dataset_document: DatasetDocument, documents: List[Document]) -> None:
         """
         Build the index for the document.
         """
-        vector_index = VectorIndex(dataset=dataset)
-        keyword_table_index = KeywordTableIndex(dataset=dataset)
+        vector_index = IndexBuilder.get_index(dataset, 'high_quality')
+        keyword_table_index = IndexBuilder.get_index(dataset, 'economy')
 
         # chunk nodes by chunk size
         indexing_start_at = time.perf_counter()
         tokens = 0
         chunk_size = 100
-        for i in range(0, len(nodes), chunk_size):
+        for i in range(0, len(documents), chunk_size):
             # check document is paused
-            self._check_document_paused_status(document.id)
-            chunk_nodes = nodes[i:i + chunk_size]
+            self._check_document_paused_status(dataset_document.id)
+            chunk_documents = documents[i:i + chunk_size]
 
             tokens += sum(
-                TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text()) for node in chunk_nodes
+                TokenCalculator.get_num_tokens(self.embedding_model_name, document.page_content)
+                for document in chunk_documents
             )
 
             # save vector index
-            if dataset.indexing_technique == "high_quality":
-                vector_index.add_nodes(chunk_nodes)
+            if vector_index:
+                vector_index.add_texts(chunk_documents)
 
             # save keyword index
-            keyword_table_index.add_nodes(chunk_nodes)
+            keyword_table_index.add_texts(chunk_documents)
 
-            node_ids = [node.doc_id for node in chunk_nodes]
+            document_ids = [document.metadata['doc_id'] for document in chunk_documents]
             db.session.query(DocumentSegment).filter(
-                DocumentSegment.document_id == document.id,
-                DocumentSegment.index_node_id.in_(node_ids),
+                DocumentSegment.document_id == dataset_document.id,
+                DocumentSegment.index_node_id.in_(document_ids),
                 DocumentSegment.status == "indexing"
             ).update({
                 DocumentSegment.status: "completed",
@@ -423,12 +531,12 @@ class IndexingRunner:
 
         # update document status to completed
         self._update_document_index_status(
-            document_id=document.id,
+            document_id=dataset_document.id,
             after_indexing_status="completed",
             extra_update_params={
-                Document.tokens: tokens,
-                Document.completed_at: datetime.datetime.utcnow(),
-                Document.indexing_latency: indexing_end_at - indexing_start_at,
+                DatasetDocument.tokens: tokens,
+                DatasetDocument.completed_at: datetime.datetime.utcnow(),
+                DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
             }
         )
 
@@ -443,25 +551,25 @@ class IndexingRunner:
         """
         Update the document indexing status.
         """
-        count = Document.query.filter_by(id=document_id, is_paused=True).count()
+        count = DatasetDocument.query.filter_by(id=document_id, is_paused=True).count()
         if count > 0:
             raise DocumentIsPausedException()
 
         update_params = {
-            Document.indexing_status: after_indexing_status
+            DatasetDocument.indexing_status: after_indexing_status
         }
 
         if extra_update_params:
             update_params.update(extra_update_params)
 
-        Document.query.filter_by(id=document_id).update(update_params)
+        DatasetDocument.query.filter_by(id=document_id).update(update_params)
         db.session.commit()
 
-    def _update_segments_by_document(self, document_id: str, update_params: dict) -> None:
+    def _update_segments_by_document(self, dataset_document_id: str, update_params: dict) -> None:
         """
         Update the document segment by document id.
         """
-        DocumentSegment.query.filter_by(document_id=document_id).update(update_params)
+        DocumentSegment.query.filter_by(document_id=dataset_document_id).update(update_params)
         db.session.commit()
 
 

api/core/llm/error_handle_wraps.py

@@ -17,14 +17,16 @@ def handle_llm_exceptions(func):
             raise LLMBadRequestError(str(e))
         except openai.error.APIConnectionError as e:
             logging.exception("Failed to connect to OpenAI API.")
-            raise LLMAPIConnectionError(str(e))
+            raise LLMAPIConnectionError(e.__class__.__name__ + ":" + str(e))
         except (openai.error.APIError, openai.error.ServiceUnavailableError, openai.error.Timeout) as e:
             logging.exception("OpenAI service unavailable.")
-            raise LLMAPIUnavailableError(str(e))
+            raise LLMAPIUnavailableError(e.__class__.__name__ + ":" + str(e))
         except openai.error.RateLimitError as e:
             raise LLMRateLimitError(str(e))
         except openai.error.AuthenticationError as e:
             raise LLMAuthorizationError(str(e))
+        except openai.error.OpenAIError as e:
+            raise LLMBadRequestError(e.__class__.__name__ + ":" + str(e))
 
     return wrapper
 
@@ -39,13 +41,15 @@ def handle_llm_exceptions_async(func):
             raise LLMBadRequestError(str(e))
         except openai.error.APIConnectionError as e:
             logging.exception("Failed to connect to OpenAI API.")
-            raise LLMAPIConnectionError(str(e))
+            raise LLMAPIConnectionError(e.__class__.__name__ + ":" + str(e))
         except (openai.error.APIError, openai.error.ServiceUnavailableError, openai.error.Timeout) as e:
             logging.exception("OpenAI service unavailable.")
-            raise LLMAPIUnavailableError(str(e))
+            raise LLMAPIUnavailableError(e.__class__.__name__ + ":" + str(e))
         except openai.error.RateLimitError as e:
             raise LLMRateLimitError(str(e))
         except openai.error.AuthenticationError as e:
             raise LLMAuthorizationError(str(e))
+        except openai.error.OpenAIError as e:
+            raise LLMBadRequestError(e.__class__.__name__ + ":" + str(e))
 
     return wrapper

api/core/llm/llm_builder.py

@@ -1,7 +1,6 @@
-from typing import Union, Optional
+from typing import Union, Optional, List
 
-from langchain.callbacks import CallbackManager
-from langchain.llms.fake import FakeListLLM
+from langchain.callbacks.base import BaseCallbackHandler
 
 from core.constant import llm_constant
 from core.llm.error import ProviderTokenNotInitError
@@ -32,12 +31,11 @@ class LLMBuilder:
     """
 
     @classmethod
-    def to_llm(cls, tenant_id: str, model_name: str, **kwargs) -> Union[StreamableOpenAI, StreamableChatOpenAI, FakeListLLM]:
-        if model_name == 'fake':
-            return FakeListLLM(responses=[])
-
+    def to_llm(cls, tenant_id: str, model_name: str, **kwargs) -> Union[StreamableOpenAI, StreamableChatOpenAI]:
         provider = cls.get_default_provider(tenant_id)
 
+        model_credentials = cls.get_model_credentials(tenant_id, provider, model_name)
+
         mode = cls.get_mode_by_model(model_name)
         if mode == 'chat':
             if provider == 'openai':
@@ -52,16 +50,21 @@ class LLMBuilder:
         else:
             raise ValueError(f"model name {model_name} is not supported.")
 
-        model_credentials = cls.get_model_credentials(tenant_id, provider, model_name)
+
+        model_kwargs = {
+            'top_p': kwargs.get('top_p', 1),
+            'frequency_penalty': kwargs.get('frequency_penalty', 0),
+            'presence_penalty': kwargs.get('presence_penalty', 0),
+        }
+
+        model_extras_kwargs = model_kwargs if mode == 'completion' else {'model_kwargs': model_kwargs}
 
         return llm_cls(
             model_name=model_name,
             temperature=kwargs.get('temperature', 0),
             max_tokens=kwargs.get('max_tokens', 256),
-            top_p=kwargs.get('top_p', 1),
-            frequency_penalty=kwargs.get('frequency_penalty', 0),
-            presence_penalty=kwargs.get('presence_penalty', 0),
-            callback_manager=kwargs.get('callback_manager', None),
+            **model_extras_kwargs,
+            callbacks=kwargs.get('callbacks', None),
             streaming=kwargs.get('streaming', False),
             # request_timeout=None
             **model_credentials
@@ -69,7 +72,7 @@ class LLMBuilder:
 
     @classmethod
     def to_llm_from_model(cls, tenant_id: str, model: dict, streaming: bool = False,
-                          callback_manager: Optional[CallbackManager] = None) -> Union[StreamableOpenAI, StreamableChatOpenAI]:
+                          callbacks: Optional[List[BaseCallbackHandler]] = None) -> Union[StreamableOpenAI, StreamableChatOpenAI]:
         model_name = model.get("name")
         completion_params = model.get("completion_params", {})
 
@@ -82,7 +85,7 @@ class LLMBuilder:
             frequency_penalty=completion_params.get('frequency_penalty', 0.1),
             presence_penalty=completion_params.get('presence_penalty', 0.1),
             streaming=streaming,
-            callback_manager=callback_manager
+            callbacks=callbacks
         )
 
     @classmethod

api/core/llm/provider/azure_provider.py

@@ -42,7 +42,11 @@ class AzureProvider(BaseProvider):
         """
         config = self.get_provider_api_key(model_id=model_id)
         config['openai_api_type'] = 'azure'
-        config['deployment_name'] = model_id.replace('.', '') if model_id else None
+        if model_id == 'text-embedding-ada-002':
+            config['deployment'] = model_id.replace('.', '') if model_id else None
+            config['chunk_size'] = 1
+        else:
+            config['deployment_name'] = model_id.replace('.', '') if model_id else None
         return config
 
     def get_provider_name(self):

api/core/llm/streamable_azure_chat_open_ai.py

@@ -1,3 +1,4 @@
+from langchain.callbacks.manager import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun, Callbacks
 from langchain.schema import BaseMessage, ChatResult, LLMResult
 from langchain.chat_models import AzureChatOpenAI
 from typing import Optional, List, Dict, Any
@@ -68,60 +69,22 @@ class StreamableAzureChatOpenAI(AzureChatOpenAI):
 
         return message_tokens
 
-    def _generate(
-            self, messages: List[BaseMessage], stop: Optional[List[str]] = None
-    ) -> ChatResult:
-        self.callback_manager.on_llm_start(
-            {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-            verbose=self.verbose
-        )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-        self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    async def _agenerate(
-            self, messages: List[BaseMessage], stop: Optional[List[str]] = None
-    ) -> ChatResult:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-                verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-                verbose=self.verbose
-            )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_end(result, verbose=self.verbose)
-        else:
-            self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
     @handle_llm_exceptions
     def generate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(messages, stop)
+        return super().generate(messages, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(messages, stop)
+        return await super().agenerate(messages, stop, callbacks, **kwargs)

api/core/llm/streamable_azure_open_ai.py

@@ -1,5 +1,4 @@
-import os
-
+from langchain.callbacks.manager import Callbacks
 from langchain.llms import AzureOpenAI
 from langchain.schema import LLMResult
 from typing import Optional, List, Dict, Mapping, Any
@@ -53,12 +52,20 @@ class StreamableAzureOpenAI(AzureOpenAI):
 
     @handle_llm_exceptions
     def generate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(prompts, stop)
+        return super().generate(prompts, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(prompts, stop)
+        return await super().agenerate(prompts, stop, callbacks, **kwargs)

api/core/llm/streamable_chat_open_ai.py

@@ -1,6 +1,7 @@
 import os
 
-from langchain.schema import BaseMessage, ChatResult, LLMResult
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import BaseMessage, LLMResult
 from langchain.chat_models import ChatOpenAI
 from typing import Optional, List, Dict, Any
 
@@ -70,57 +71,22 @@ class StreamableChatOpenAI(ChatOpenAI):
 
         return message_tokens
 
-    def _generate(
-        self, messages: List[BaseMessage], stop: Optional[List[str]] = None
-    ) -> ChatResult:
-        self.callback_manager.on_llm_start(
-            {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-        )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-        self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    async def _agenerate(
-        self, messages: List[BaseMessage], stop: Optional[List[str]] = None
-    ) -> ChatResult:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-            )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_end(result, verbose=self.verbose)
-        else:
-            self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
     @handle_llm_exceptions
     def generate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(messages, stop)
+        return super().generate(messages, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(messages, stop)
+        return await super().agenerate(messages, stop, callbacks, **kwargs)

api/core/llm/streamable_open_ai.py

@@ -1,5 +1,6 @@
 import os
 
+from langchain.callbacks.manager import Callbacks
 from langchain.schema import LLMResult
 from typing import Optional, List, Dict, Any, Mapping
 from langchain import OpenAI
@@ -48,15 +49,22 @@ class StreamableOpenAI(OpenAI):
             "organization": self.openai_organization if self.openai_organization else None,
         }}
 
-
     @handle_llm_exceptions
     def generate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(prompts, stop)
+        return super().generate(prompts, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(prompts, stop)
+        return await super().agenerate(prompts, stop, callbacks, **kwargs)