feat: bump version to 0.3.22 (#1153)

Co-authored-by: jyong <jyong@dify.ai>
Co-authored-by: StyleZhang <jasonapring2015@outlook.com>

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,49 @@
+name: "🕷️ Bug report"
+description: Report errors or unexpected behavior
+labels:
+- bug
+body:
+- type: markdown
+  attributes:
+    value: Please make sure to [search for existing issues](https://github.com/langgenius/dify/issues) before filing a new one!
+- type: input
+  attributes:
+    label: Dify version
+    placeholder: 0.3.21
+    description: See about section in Dify console
+  validations:
+    required: true
+
+- type: dropdown
+  attributes:
+    label: Cloud or Self Hosted
+    description: How / Where was Dify installed from?
+    multiple: true
+    options:
+      - Cloud
+      - Self Hosted
+      - Other (please specify in "Steps to Reproduce")
+  validations:
+    required: true
+
+- type: textarea
+  attributes:
+    label: Steps to reproduce
+    description: We highly suggest including screenshots and a bug report log.
+    placeholder: Having detailed steps helps us reproduce the bug. 
+  validations:
+    required: true
+
+- type: textarea
+  attributes:
+    label: ✔️ Expected Behavior
+    placeholder: What were you expecting?
+  validations:
+    required: false
+
+- type: textarea
+  attributes:
+    label: ❌ Actual Behavior
+    placeholder: What happened instead?
+  validations:
+    required: false

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,8 @@
+blank_issues_enabled: false
+contact_links:
+  - name: "\U0001F4DA Dify user documentation"
+    url: https://docs.dify.ai/getting-started/readme
+    about: Documentation for users of Dify
+  - name: "\U0001F4DA Dify dev documentation"
+    url: https://docs.dify.ai/getting-started/install-self-hosted
+    about: Documentation for people interested in developing and contributing for Dify

.github/ISSUE_TEMPLATE/document_issue.yml

@@ -0,0 +1,11 @@
+name: "📚 Documentation Issue"
+description: Report issues in our documentation
+labels: 
+- ducumentation
+body:
+- type: textarea
+  attributes: 
+    label: Provide a description of requested docs changes
+    placeholder: Briefly describe which document needs to be corrected and why.
+  validations:
+    required: true

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,26 @@
+name: "⭐ Feature or enhancement request"
+description: Propose something new.
+labels:
+- enhancement
+body:
+- type: textarea
+  attributes: 
+    label: Description of the new feature / enhancement
+    placeholder: What is the expected behavior of the proposed feature?
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Scenario when this would be used?
+    placeholder: What is the scenario this would be used? Why is this important to your workflow as a dify user?
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Supporting information
+    placeholder: "Having additional evidence, data, tweets, blog posts, research, ... anything is extremely helpful. This information provides context to the scenario that may otherwise be lost."
+  validations:
+    required: false
+- type: markdown
+  attributes:
+    value: Please limit one request per issue.

.github/ISSUE_TEMPLATE/translation_issue.yml

@@ -0,0 +1,46 @@
+name: "🌐 Localization/Translation issue"
+description: Report incorrect translations.
+labels:
+- translation
+body:
+- type: markdown
+  attributes:
+    value: Please make sure to [search for existing issues](https://github.com/langgenius/dify/issues) before filing a new one!
+- type: input
+  attributes:
+    label: Dify version
+    placeholder: 0.3.21
+    description: Hover over system tray icon or look at Settings
+  validations:
+    required: true
+- type: input
+  attributes:
+    label: Utility with translation issue
+    placeholder: Some area
+    description: Please input here the utility with the translation issue
+  validations:
+    required: true
+- type: input
+  attributes:
+    label: 🌐 Language affected
+    placeholder: "German"
+  validations:
+    required: true
+- type: textarea
+  attributes: 
+    label: ❌ Actual phrase(s)
+    placeholder: What is there? Please include a screenshot as that is extremely helpful.
+  validations:
+    required: true
+- type: textarea
+  attributes: 
+    label: ✔️ Expected phrase(s)
+    placeholder: What was expected?
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: ℹ Why is the current translation wrong
+    placeholder: Why do you feel this is incorrect?
+  validations:
+    required: true

.github/ISSUE_TEMPLATE/🐛-bug-report.md

@@ -1,32 +0,0 @@
----
-name: "\U0001F41B Bug report"
-about: Create a report to help us improve
-title: ''
-labels: bug
-assignees: ''
-
----
-
-<!--
-  Please provide a clear and concise description of what the bug is. Include
-  screenshots if needed. Please test using the latest version of the relevant
-  Dify packages to make sure your issue has not already been fixed.
--->
-
-Dify version: Cloud | Self Host
-
-## Steps To Reproduce
-<!--
-  Your bug will get fixed much faster if we can run your code and it doesn't
-  have dependencies other than Dify. Issues without reproduction steps or
-  code examples may be immediately closed as not actionable.
--->
-
-1.
-2.
-
-
-## The current behavior
-
-
-## The expected behavior

.github/ISSUE_TEMPLATE/🚀-feature-request.md

@@ -1,20 +0,0 @@
----
-name: "\U0001F680 Feature request"
-about: Suggest an idea for this project
-title: ''
-labels: enhancement
-assignees: ''
-
----
-
-**Is your feature request related to a problem? Please describe.**
-A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-
-**Describe the solution you'd like**
-A clear and concise description of what you want to happen.
-
-**Describe alternatives you've considered**
-A clear and concise description of any alternative solutions or features you've considered.
-
-**Additional context**
-Add any other context or screenshots about the feature request here.

.github/ISSUE_TEMPLATE/🤔-questions-and-help.md

@@ -1,10 +0,0 @@
----
-name: "\U0001F914 Questions and Help"
-about: Ask a usage or consultation question
-title: ''
-labels: ''
-assignees: ''
-
----
-
-

.github/workflows/check_no_chinese_comments.py

@@ -20,7 +20,8 @@ def check_file_for_chinese_comments(file_path):
 def main():
     has_chinese = False
     excluded_files = ["model_template.py", 'stopwords.py', 'commands.py',
-                      'indexing_runner.py', 'web_reader_tool.py', 'spark_provider.py']
+                      'indexing_runner.py', 'web_reader_tool.py', 'spark_provider.py',
+                      'prompts.py']
 
     for root, _, files in os.walk("."):
         for file in files:

.gitignore

@@ -149,4 +149,5 @@ sdks/python-client/build
 sdks/python-client/dist
 sdks/python-client/dify_client.egg-info
 
-.vscode/
+.vscode/*
+!.vscode/launch.json

.vscode/launch.json

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

CONTRIBUTING.md

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

CONTRIBUTING_CN.md

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

CONTRIBUTING_JA.md

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

api/Dockerfile

@@ -1,7 +1,18 @@
-FROM python:3.10-slim
+# packages install stage
+FROM python:3.10-slim AS base
 
 LABEL maintainer="takatost@gmail.com"
 
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends gcc g++ python3-dev libc-dev libffi-dev
+
+COPY requirements.txt /requirements.txt
+
+RUN pip install --prefix=/pkg -r requirements.txt
+
+# build stage
+FROM python:3.10-slim AS builder
+
 ENV FLASK_APP app.py
 ENV EDITION SELF_HOSTED
 ENV DEPLOY_ENV PRODUCTION
@@ -15,15 +26,17 @@ EXPOSE 5001
 
 WORKDIR /app/api
 
-RUN apt-get update && \
-    apt-get install -y bash curl wget vim gcc g++ python3-dev libc-dev libffi-dev nodejs
-
-COPY requirements.txt /app/api/requirements.txt
-
-RUN pip install -r requirements.txt
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends bash curl wget vim nodejs \
+    && apt-get autoremove \
+    && rm -rf /var/lib/apt/lists/*
 
+COPY --from=base /pkg /usr/local
 COPY . /app/api/
 
+RUN python -c "from transformers import GPT2TokenizerFast; GPT2TokenizerFast.from_pretrained('gpt2')"
+ENV TRANSFORMERS_OFFLINE true
+
 COPY docker/entrypoint.sh /entrypoint.sh
 RUN chmod +x /entrypoint.sh
 

api/README.md

@@ -52,11 +52,13 @@
    flask run --host 0.0.0.0 --port=5001 --debug
    ```
 7. Setup your application by visiting http://localhost:5001/console/api/setup or other apis...
-8. If you need to debug local async processing, you can run `celery -A app.celery worker -Q dataset,generation,mail`, celery can do dataset importing and other async tasks.
+8. If you need to debug local async processing, you can run `celery -A app.celery worker -P gevent -c 1 --loglevel INFO -Q dataset,generation,mail`, celery can do dataset importing and other async tasks.
 
-8. Start frontend:
+8. Start frontend
+
+   You can start the frontend by running `npm install && npm run dev` in web/ folder, or you can use docker to start the frontend, for example:
 
    ```
-   docker run -it -d --platform linux/amd64 -p 3000:3000 -e EDITION=SELF_HOSTED -e CONSOLE_URL=http://127.0.0.1:5000 --name web-self-hosted langgenius/dify-web:latest
+   docker run -it -d --platform linux/amd64 -p 3000:3000 -e EDITION=SELF_HOSTED -e CONSOLE_URL=http://127.0.0.1:5001 --name web-self-hosted langgenius/dify-web:latest
    ```
    This will start a dify frontend, now you are all set, happy coding!

api/app.py

@@ -1,6 +1,6 @@
 # -*- coding:utf-8 -*-
 import os
-from datetime import datetime
+from datetime import datetime, timedelta
 
 from werkzeug.exceptions import Forbidden
 
@@ -145,8 +145,12 @@ def load_user(user_id):
                     _create_tenant_for_account(account)
                 session['workspace_id'] = account.current_tenant_id
 
-            account.last_active_at = datetime.utcnow()
-            db.session.commit()
+            current_time = datetime.utcnow()
+
+            # update last_active_at when last_active_at is more than 10 minutes ago
+            if current_time - account.last_active_at > timedelta(minutes=10):
+                account.last_active_at = current_time
+                db.session.commit()
 
             # Log in the user with the updated user_id
             flask_login.login_user(account, remember=True)

api/commands.py

@@ -1,22 +1,30 @@
 import datetime
+import json
 import math
 import random
 import string
 import time
 
 import click
+from tqdm import tqdm
 from flask import current_app
+from langchain.embeddings import OpenAIEmbeddings
 from werkzeug.exceptions import NotFound
 
+from core.embedding.cached_embedding import CacheEmbedding
 from core.index.index import IndexBuilder
+from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.embedding.openai_embedding import OpenAIEmbedding
+from core.model_providers.models.entity.model_params import ModelType
 from core.model_providers.providers.hosted import hosted_model_providers
+from core.model_providers.providers.openai_provider import OpenAIProvider
 from libs.password import password_pattern, valid_password, hash_password
 from libs.helper import email as email_validate
 from extensions.ext_database import db
 from libs.rsa import generate_key_pair
-from models.account import InvitationCode, Tenant
+from models.account import InvitationCode, Tenant, TenantAccountJoin
 from models.dataset import Dataset, DatasetQuery, Document
-from models.model import Account
+from models.model import Account, AppModelConfig, App
 import secrets
 import base64
 
@@ -296,6 +304,243 @@ def sync_anthropic_hosted_providers():
     click.echo(click.style('Congratulations! Synced {} anthropic hosted providers.'.format(count), fg='green'))
 
 
+@click.command('create-qdrant-indexes', help='Create qdrant indexes.')
+def create_qdrant_indexes():
+    click.echo(click.style('Start create qdrant indexes.', fg='green'))
+    create_count = 0
+
+    page = 1
+    while True:
+        try:
+            datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality') \
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+        except NotFound:
+            break
+
+        page += 1
+        for dataset in datasets:
+            if dataset.index_struct_dict:
+                if dataset.index_struct_dict['type'] != 'qdrant':
+                    try:
+                        click.echo('Create dataset qdrant index: {}'.format(dataset.id))
+                        try:
+                            embedding_model = ModelFactory.get_embedding_model(
+                                tenant_id=dataset.tenant_id,
+                                model_provider_name=dataset.embedding_model_provider,
+                                model_name=dataset.embedding_model
+                            )
+                        except Exception:
+                            try:
+                                embedding_model = ModelFactory.get_embedding_model(
+                                    tenant_id=dataset.tenant_id
+                                )
+                                dataset.embedding_model = embedding_model.name
+                                dataset.embedding_model_provider = embedding_model.model_provider.provider_name
+                            except Exception:
+                                provider = Provider(
+                                    id='provider_id',
+                                    tenant_id=dataset.tenant_id,
+                                    provider_name='openai',
+                                    provider_type=ProviderType.SYSTEM.value,
+                                    encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
+                                    is_valid=True,
+                                )
+                                model_provider = OpenAIProvider(provider=provider)
+                                embedding_model = OpenAIEmbedding(name="text-embedding-ada-002", model_provider=model_provider)
+                        embeddings = CacheEmbedding(embedding_model)
+
+                        from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+                        index = QdrantVectorIndex(
+                            dataset=dataset,
+                            config=QdrantConfig(
+                                endpoint=current_app.config.get('QDRANT_URL'),
+                                api_key=current_app.config.get('QDRANT_API_KEY'),
+                                root_path=current_app.root_path
+                            ),
+                            embeddings=embeddings
+                        )
+                        if index:
+                            index.create_qdrant_dataset(dataset)
+                            index_struct = {
+                                "type": 'qdrant',
+                                "vector_store": {"class_prefix": dataset.index_struct_dict['vector_store']['class_prefix']}
+                            }
+                            dataset.index_struct = json.dumps(index_struct)
+                            db.session.commit()
+                            create_count += 1
+                        else:
+                            click.echo('passed.')
+                    except Exception as e:
+                        click.echo(
+                            click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                        continue
+
+    click.echo(click.style('Congratulations! Create {} dataset indexes.'.format(create_count), fg='green'))
+
+
+@click.command('update-qdrant-indexes', help='Update qdrant indexes.')
+def update_qdrant_indexes():
+    click.echo(click.style('Start Update qdrant indexes.', fg='green'))
+    create_count = 0
+
+    page = 1
+    while True:
+        try:
+            datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality') \
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+        except NotFound:
+            break
+
+        page += 1
+        for dataset in datasets:
+            if dataset.index_struct_dict:
+                if dataset.index_struct_dict['type'] != 'qdrant':
+                    try:
+                        click.echo('Update dataset qdrant index: {}'.format(dataset.id))
+                        try:
+                            embedding_model = ModelFactory.get_embedding_model(
+                                tenant_id=dataset.tenant_id,
+                                model_provider_name=dataset.embedding_model_provider,
+                                model_name=dataset.embedding_model
+                            )
+                        except Exception:
+                            provider = Provider(
+                                id='provider_id',
+                                tenant_id=dataset.tenant_id,
+                                provider_name='openai',
+                                provider_type=ProviderType.CUSTOM.value,
+                                encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
+                                is_valid=True,
+                            )
+                            model_provider = OpenAIProvider(provider=provider)
+                            embedding_model = OpenAIEmbedding(name="text-embedding-ada-002", model_provider=model_provider)
+                        embeddings = CacheEmbedding(embedding_model)
+
+                        from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+                        index = QdrantVectorIndex(
+                            dataset=dataset,
+                            config=QdrantConfig(
+                                endpoint=current_app.config.get('QDRANT_URL'),
+                                api_key=current_app.config.get('QDRANT_API_KEY'),
+                                root_path=current_app.root_path
+                            ),
+                            embeddings=embeddings
+                        )
+                        if index:
+                            index.update_qdrant_dataset(dataset)
+                            create_count += 1
+                        else:
+                            click.echo('passed.')
+                    except Exception as e:
+                        click.echo(
+                            click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                        continue
+
+    click.echo(click.style('Congratulations! Update {} dataset indexes.'.format(create_count), fg='green'))
+
+@click.command('update_app_model_configs', help='Migrate data to support paragraph variable.')
+@click.option("--batch-size", default=500, help="Number of records to migrate in each batch.")
+def update_app_model_configs(batch_size):
+    pre_prompt_template = '{{default_input}}'
+    user_input_form_template = {
+        "en-US": [
+            {
+                "paragraph": {
+                    "label": "Query",
+                    "variable": "default_input",
+                    "required": False,
+                    "default": ""
+                }
+            }
+        ],
+        "zh-Hans": [
+            {
+                "paragraph": {
+                    "label": "查询内容",
+                    "variable": "default_input",
+                    "required": False,
+                    "default": ""
+                }
+            }
+        ]
+    }
+
+    click.secho("Start migrate old data that the text generator can support paragraph variable.", fg='green')
+
+    total_records = db.session.query(AppModelConfig) \
+        .join(App, App.app_model_config_id == AppModelConfig.id) \
+        .filter(App.mode == 'completion') \
+        .count()
+    
+    if total_records == 0:
+        click.secho("No data to migrate.", fg='green')
+        return
+
+    num_batches = (total_records + batch_size - 1) // batch_size
+
+    with tqdm(total=total_records, desc="Migrating Data") as pbar:
+        for i in range(num_batches):
+            offset = i * batch_size
+            limit = min(batch_size, total_records - offset)
+
+            click.secho(f"Fetching batch {i+1}/{num_batches} from source database...", fg='green')
+            
+            data_batch = db.session.query(AppModelConfig) \
+                .join(App, App.app_model_config_id == AppModelConfig.id) \
+                .filter(App.mode == 'completion') \
+                .order_by(App.created_at) \
+                .offset(offset).limit(limit).all()
+            
+            if not data_batch:
+                click.secho("No more data to migrate.", fg='green')
+                break
+
+            try:
+                click.secho(f"Migrating {len(data_batch)} records...", fg='green')
+                for data in data_batch:
+                    # click.secho(f"Migrating data {data.id}, pre_prompt: {data.pre_prompt}, user_input_form: {data.user_input_form}", fg='green')
+
+                    if data.pre_prompt is None:
+                        data.pre_prompt = pre_prompt_template
+                    else:
+                        if pre_prompt_template in data.pre_prompt:
+                            continue
+                        data.pre_prompt += pre_prompt_template
+
+                    app_data = db.session.query(App) \
+                        .filter(App.id == data.app_id) \
+                        .one()
+                    
+                    account_data = db.session.query(Account) \
+                        .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) \
+                        .filter(TenantAccountJoin.role == 'owner') \
+                        .filter(TenantAccountJoin.tenant_id == app_data.tenant_id) \
+                        .one_or_none()
+
+                    if not account_data:
+                        continue
+
+                    if data.user_input_form is None or data.user_input_form == 'null':
+                        data.user_input_form = json.dumps(user_input_form_template[account_data.interface_language])
+                    else:
+                        raw_json_data = json.loads(data.user_input_form)
+                        raw_json_data.append(user_input_form_template[account_data.interface_language][0])
+                        data.user_input_form = json.dumps(raw_json_data)
+
+                    # click.secho(f"Updated data {data.id}, pre_prompt: {data.pre_prompt}, user_input_form: {data.user_input_form}", fg='green')
+
+                db.session.commit()
+
+            except Exception as e:
+                click.secho(f"Error while migrating data: {e}, app_id: {data.app_id}, app_model_config_id: {data.id}", fg='red')
+                continue
+            
+            click.secho(f"Successfully migrated batch {i+1}/{num_batches}.", fg='green')
+            
+            pbar.update(len(data_batch))
+
 def register_commands(app):
     app.cli.add_command(reset_password)
     app.cli.add_command(reset_email)
@@ -304,3 +549,6 @@ def register_commands(app):
     app.cli.add_command(recreate_all_dataset_indexes)
     app.cli.add_command(sync_anthropic_hosted_providers)
     app.cli.add_command(clean_unused_dataset_indexes)
+    app.cli.add_command(create_qdrant_indexes)
+    app.cli.add_command(update_qdrant_indexes)
+    app.cli.add_command(update_app_model_configs)

api/config.py

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

api/constants/model_template.py

@@ -38,7 +38,18 @@ model_templates = {
                     "presence_penalty": 0,
                     "frequency_penalty": 0
                 }
-            })
+            }),
+            'user_input_form': json.dumps([
+                {
+                    "paragraph": {
+                        "label": "Query",
+                        "variable": "query",
+                        "required": True,
+                        "default": ""
+                    }
+                }
+            ]),
+            'pre_prompt': '{{query}}'
         }
     },
 

api/controllers/console/app/app.py

@@ -29,6 +29,7 @@ model_config_fields = {
     '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'),
+    'retriever_resource': fields.Raw(attribute='retriever_resource_dict'),
     'more_like_this': fields.Raw(attribute='more_like_this_dict'),
     'sensitive_word_avoidance': fields.Raw(attribute='sensitive_word_avoidance_dict'),
     'model': fields.Raw(attribute='model_dict'),

api/controllers/console/app/completion.py

@@ -39,9 +39,10 @@ class CompletionMessageApi(Resource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('model_config', type=dict, required=True, location='json')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='dev', location='json')
         args = parser.parse_args()
 
         streaming = args['response_mode'] != 'blocking'
@@ -115,6 +116,7 @@ class ChatMessageApi(Resource):
         parser.add_argument('model_config', type=dict, required=True, location='json')
         parser.add_argument('conversation_id', type=uuid_value, location='json')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='dev', location='json')
         args = parser.parse_args()
 
         streaming = args['response_mode'] != 'blocking'

api/controllers/console/auth/activate.py

@@ -16,26 +16,25 @@ 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('workspace_id', type=str, required=False, nullable=True, location='args')
+        parser.add_argument('email', type=email, required=False, nullable=True, 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'])
+        workspaceId = args['workspace_id']
+        reg_email = args['email']
+        token = args['token']
 
-        tenant = db.session.query(Tenant).filter(
-            Tenant.id == args['workspace_id'],
-            Tenant.status == 'normal'
-        ).first()
+        invitation = RegisterService.get_invitation_if_token_valid(workspaceId, reg_email, token)
 
-        return {'is_valid': account is not None, 'workspace_name': tenant.name}
+        return {'is_valid': invitation is not None, 'workspace_name': invitation['tenant'].name if invitation else None}
 
 
 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('workspace_id', type=str, required=False, nullable=True, location='json')
+        parser.add_argument('email', type=email, required=False, nullable=True, 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')
@@ -44,12 +43,13 @@ class ActivateApi(Resource):
         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:
+        invitation = RegisterService.get_invitation_if_token_valid(args['workspace_id'], args['email'], args['token'])
+        if invitation is None:
             raise AlreadyActivateError()
 
         RegisterService.revoke_token(args['workspace_id'], args['email'], args['token'])
 
+        account = invitation['account']
         account.name = args['name']
 
         # generate password salt

api/controllers/console/datasets/datasets.py

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

api/controllers/console/datasets/datasets_document.py

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

api/controllers/console/datasets/datasets_segments.py

@@ -149,7 +149,8 @@ class DatasetDocumentSegmentApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
-
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # The role of the current user in the ta table must be admin or owner
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
@@ -158,20 +159,20 @@ class DatasetDocumentSegmentApi(Resource):
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
-
-        # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
+        if dataset.indexing_technique == 'high_quality':
+            # check embedding model setting
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
 
         segment = DocumentSegment.query.filter(
             DocumentSegment.id == str(segment_id),
@@ -244,18 +245,19 @@ class DatasetDocumentSegmentAddApi(Resource):
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
         # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
+        if dataset.indexing_technique == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
         try:
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
@@ -284,25 +286,28 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # check document
         document_id = str(document_id)
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound('Document not found.')
-        # check embedding model setting
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
-        # check segment
+        if dataset.indexing_technique == 'high_quality':
+            # check embedding model setting
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+            # check segment
         segment_id = str(segment_id)
         segment = DocumentSegment.query.filter(
             DocumentSegment.id == str(segment_id),
@@ -339,6 +344,8 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
         # check document
         document_id = str(document_id)
         document = DocumentService.get_document(dataset_id, document_id)
@@ -378,18 +385,6 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound('Document not found.')
-        try:
-            ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id,
-                model_provider_name=dataset.embedding_model_provider,
-                model_name=dataset.embedding_model
-            )
-        except LLMBadRequestError:
-            raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
-        except ProviderTokenNotInitError as ex:
-            raise ProviderNotInitializeError(ex.description)
         # get file from request
         file = request.files['file']
         # check file

api/controllers/console/datasets/file.py

@@ -26,7 +26,7 @@ from models.model import UploadFile
 
 cache = TTLCache(maxsize=None, ttl=30)
 
-ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx']
+ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'docx', 'csv']
 PREVIEW_WORDS_LIMIT = 3000
 
 
@@ -83,7 +83,7 @@ class FileApi(Resource):
             raise FileTooLargeError(message)
 
         extension = file.filename.split('.')[-1]
-        if extension not in ALLOWED_EXTENSIONS:
+        if extension.lower() not in ALLOWED_EXTENSIONS:
             raise UnsupportedFileTypeError()
 
         # user uuid as file name
@@ -136,7 +136,7 @@ class FilePreviewApi(Resource):
 
         # extract text from file
         extension = upload_file.extension
-        if extension not in ALLOWED_EXTENSIONS:
+        if extension.lower() not in ALLOWED_EXTENSIONS:
             raise UnsupportedFileTypeError()
 
         text = FileExtractor.load(upload_file, return_text=True)

api/controllers/console/explore/completion.py

@@ -31,8 +31,9 @@ class CompletionApi(InstalledAppResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='explore_app', location='json')
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'
@@ -92,6 +93,7 @@ class ChatApi(InstalledAppResource):
         parser.add_argument('query', type=str, required=True, location='json')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         parser.add_argument('conversation_id', type=uuid_value, location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='explore_app', location='json')
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'

api/controllers/console/explore/message.py

@@ -30,6 +30,25 @@ class MessageListApi(InstalledAppResource):
         'rating': fields.String
     }
 
+    retriever_resource_fields = {
+        'id': fields.String,
+        'message_id': fields.String,
+        'position': fields.Integer,
+        'dataset_id': fields.String,
+        'dataset_name': fields.String,
+        'document_id': fields.String,
+        'document_name': fields.String,
+        'data_source_type': fields.String,
+        'segment_id': fields.String,
+        'score': fields.Float,
+        'hit_count': fields.Integer,
+        'word_count': fields.Integer,
+        'segment_position': fields.Integer,
+        'index_node_hash': fields.String,
+        'content': fields.String,
+        'created_at': TimestampField
+    }
+
     message_fields = {
         'id': fields.String,
         'conversation_id': fields.String,
@@ -37,6 +56,7 @@ class MessageListApi(InstalledAppResource):
         'query': fields.String,
         'answer': fields.String,
         'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
+        'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
         'created_at': TimestampField
     }
 

api/controllers/console/explore/parameter.py

@@ -24,6 +24,7 @@ class AppParameterApi(InstalledAppResource):
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
         'speech_to_text': fields.Raw,
+        'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -39,6 +40,7 @@ class AppParameterApi(InstalledAppResource):
             '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,
+            'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/console/universal_chat/chat.py

@@ -29,9 +29,11 @@ class UniversalChatApi(UniversalChatResource):
         parser.add_argument('provider', type=str, required=True, location='json')
         parser.add_argument('model', type=str, required=True, location='json')
         parser.add_argument('tools', type=list, required=True, location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='universal_app', location='json')
         args = parser.parse_args()
 
         app_model_config = app_model.app_model_config
+        app_model_config
 
         # update app model config
         args['model_config'] = app_model_config.to_dict()

api/controllers/console/universal_chat/message.py

@@ -36,6 +36,25 @@ class UniversalChatMessageListApi(UniversalChatResource):
         'created_at': TimestampField
     }
 
+    retriever_resource_fields = {
+        'id': fields.String,
+        'message_id': fields.String,
+        'position': fields.Integer,
+        'dataset_id': fields.String,
+        'dataset_name': fields.String,
+        'document_id': fields.String,
+        'document_name': fields.String,
+        'data_source_type': fields.String,
+        'segment_id': fields.String,
+        'score': fields.Float,
+        'hit_count': fields.Integer,
+        'word_count': fields.Integer,
+        'segment_position': fields.Integer,
+        'index_node_hash': fields.String,
+        'content': fields.String,
+        'created_at': TimestampField
+    }
+
     message_fields = {
         'id': fields.String,
         'conversation_id': fields.String,
@@ -43,6 +62,7 @@ class UniversalChatMessageListApi(UniversalChatResource):
         'query': fields.String,
         'answer': fields.String,
         'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
+        'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
         'created_at': TimestampField,
         'agent_thoughts': fields.List(fields.Nested(agent_thought_fields))
     }

api/controllers/console/universal_chat/parameter.py

@@ -1,4 +1,6 @@
 # -*- coding:utf-8 -*-
+import json
+
 from flask_restful import marshal_with, fields
 
 from controllers.console import api
@@ -14,6 +16,7 @@ class UniversalChatParameterApi(UniversalChatResource):
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
         'speech_to_text': fields.Raw,
+        'retriever_resource': fields.Raw,
     }
 
     @marshal_with(parameters_fields)
@@ -21,12 +24,14 @@ class UniversalChatParameterApi(UniversalChatResource):
         """Retrieve app parameters."""
         app_model = universal_app
         app_model_config = app_model.app_model_config
+        app_model_config.retriever_resource = json.dumps({'enabled': True})
 
         return {
             '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,
+            'retriever_resource': app_model_config.retriever_resource_dict,
         }
 
 

api/controllers/console/universal_chat/wraps.py

@@ -47,6 +47,7 @@ def universal_chat_app_required(view=None):
                     suggested_questions=json.dumps([]),
                     suggested_questions_after_answer=json.dumps({'enabled': True}),
                     speech_to_text=json.dumps({'enabled': True}),
+                    retriever_resource=json.dumps({'enabled': True}),
                     more_like_this=None,
                     sensitive_word_avoidance=None,
                     model=json.dumps({

api/controllers/console/workspace/members.py

@@ -49,46 +49,43 @@ class MemberInviteEmailApi(Resource):
     @account_initialization_required
     def post(self):
         parser = reqparse.RequestParser()
-        parser.add_argument('email', type=str, required=True, location='json')
+        parser.add_argument('emails', type=str, required=True, location='json', action='append')
         parser.add_argument('role', type=str, required=True, default='admin', location='json')
         args = parser.parse_args()
 
-        invitee_email = args['email']
+        invitee_emails = args['emails']
         invitee_role = args['role']
         if invitee_role not in ['admin', 'normal']:
             return {'code': 'invalid-role', 'message': 'Invalid role'}, 400
 
         inviter = current_user
-
-        try:
-            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()
-            account, role = account
-            account = marshal(account, account_fields)
-            account['role'] = role
-        except services.errors.account.CannotOperateSelfError as e:
-            return {'code': 'cannot-operate-self', 'message': str(e)}, 400
-        except services.errors.account.NoPermissionError as e:
-            return {'code': 'forbidden', 'message': str(e)}, 403
-        except services.errors.account.AccountAlreadyInTenantError as e:
-            return {'code': 'email-taken', 'message': str(e)}, 409
-        except Exception as e:
-            return {'code': 'unexpected-error', 'message': str(e)}, 500
-
-        # todo:413
+        invitation_results = []
+        console_web_url = current_app.config.get("CONSOLE_WEB_URL")
+        for invitee_email in invitee_emails:
+            try:
+                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 == invitee_email).first()
+                account, role = account
+                invitation_results.append({
+                    'status': 'success',
+                    'email': invitee_email,
+                    'url': f'{console_web_url}/activate?email={invitee_email}&token={token}'
+                })
+                account = marshal(account, account_fields)
+                account['role'] = role
+            except Exception as e:
+                invitation_results.append({
+                    'status': 'failed',
+                    'email': invitee_email,
+                    'message': str(e)
+                })
 
         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
-            )
+            'invitation_results': invitation_results,
         }, 201
 
 

api/controllers/service_api/app/app.py

@@ -25,6 +25,7 @@ class AppParameterApi(AppApiResource):
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
         'speech_to_text': fields.Raw,
+        'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -39,6 +40,7 @@ class AppParameterApi(AppApiResource):
             '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,
+            'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/service_api/app/completion.py

@@ -27,9 +27,11 @@ class CompletionApi(AppApiResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         parser.add_argument('user', type=str, location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='dev', location='json')
+
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'
@@ -91,6 +93,8 @@ class ChatApi(AppApiResource):
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         parser.add_argument('conversation_id', type=uuid_value, location='json')
         parser.add_argument('user', type=str, location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='dev', location='json')
+
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'

api/controllers/service_api/app/message.py

@@ -16,6 +16,24 @@ class MessageListApi(AppApiResource):
     feedback_fields = {
         'rating': fields.String
     }
+    retriever_resource_fields = {
+        'id': fields.String,
+        'message_id': fields.String,
+        'position': fields.Integer,
+        'dataset_id': fields.String,
+        'dataset_name': fields.String,
+        'document_id': fields.String,
+        'document_name': fields.String,
+        'data_source_type': fields.String,
+        'segment_id': fields.String,
+        'score': fields.Float,
+        'hit_count': fields.Integer,
+        'word_count': fields.Integer,
+        'segment_position': fields.Integer,
+        'index_node_hash': fields.String,
+        'content': fields.String,
+        'created_at': TimestampField
+    }
 
     message_fields = {
         'id': fields.String,
@@ -24,6 +42,7 @@ class MessageListApi(AppApiResource):
         'query': fields.String,
         'answer': fields.String,
         'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
+        'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
         'created_at': TimestampField
     }
 

api/controllers/web/app.py

@@ -24,6 +24,7 @@ class AppParameterApi(WebApiResource):
         'suggested_questions': fields.Raw,
         'suggested_questions_after_answer': fields.Raw,
         'speech_to_text': fields.Raw,
+        'retriever_resource': fields.Raw,
         'more_like_this': fields.Raw,
         'user_input_form': fields.Raw,
     }
@@ -38,6 +39,7 @@ class AppParameterApi(WebApiResource):
             '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,
+            'retriever_resource': app_model_config.retriever_resource_dict,
             'more_like_this': app_model_config.more_like_this_dict,
             'user_input_form': app_model_config.user_input_form_list
         }

api/controllers/web/completion.py

@@ -29,8 +29,10 @@ class CompletionApi(WebApiResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='web_app', location='json')
+
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'
@@ -88,6 +90,8 @@ class ChatApi(WebApiResource):
         parser.add_argument('query', type=str, required=True, location='json')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         parser.add_argument('conversation_id', type=uuid_value, location='json')
+        parser.add_argument('retriever_from', type=str, required=False, default='web_app', location='json')
+
         args = parser.parse_args()
 
         streaming = args['response_mode'] == 'streaming'

api/controllers/web/message.py

@@ -29,6 +29,25 @@ class MessageListApi(WebApiResource):
         'rating': fields.String
     }
 
+    retriever_resource_fields = {
+        'id': fields.String,
+        'message_id': fields.String,
+        'position': fields.Integer,
+        'dataset_id': fields.String,
+        'dataset_name': fields.String,
+        'document_id': fields.String,
+        'document_name': fields.String,
+        'data_source_type': fields.String,
+        'segment_id': fields.String,
+        'score': fields.Float,
+        'hit_count': fields.Integer,
+        'word_count': fields.Integer,
+        'segment_position': fields.Integer,
+        'index_node_hash': fields.String,
+        'content': fields.String,
+        'created_at': TimestampField
+    }
+
     message_fields = {
         'id': fields.String,
         'conversation_id': fields.String,
@@ -36,6 +55,7 @@ class MessageListApi(WebApiResource):
         'query': fields.String,
         'answer': fields.String,
         'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
+        'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
         'created_at': TimestampField
     }
 

api/core/agent/agent/multi_dataset_router_agent.py

@@ -1,3 +1,4 @@
+import json
 from typing import Tuple, List, Any, Union, Sequence, Optional, cast
 
 from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
@@ -52,7 +53,11 @@ class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
         elif len(self.tools) == 1:
             tool = next(iter(self.tools))
             tool = cast(DatasetRetrieverTool, tool)
-            rst = tool.run(tool_input={'dataset_id': tool.dataset_id, 'query': kwargs['input']})
+            rst = tool.run(tool_input={'query': kwargs['input']})
+            # output = ''
+            # rst_json = json.loads(rst)
+            # for item in rst_json:
+            #     output += f'{item["content"]}\n'
             return AgentFinish(return_values={"output": rst}, log=rst)
 
         if intermediate_steps:
@@ -60,7 +65,13 @@ class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
             return AgentFinish(return_values={"output": observation}, log=observation)
 
         try:
-            return super().plan(intermediate_steps, callbacks, **kwargs)
+            agent_decision = super().plan(intermediate_steps, callbacks, **kwargs)
+            if isinstance(agent_decision, AgentAction):
+                tool_inputs = agent_decision.tool_input
+                if isinstance(tool_inputs, dict) and 'query' in tool_inputs:
+                    tool_inputs['query'] = kwargs['input']
+                    agent_decision.tool_input = tool_inputs
+            return agent_decision
         except Exception as e:
             new_exception = self.model_instance.handle_exceptions(e)
             raise new_exception

api/core/agent/agent/openai_function_call.py

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

api/core/agent/agent/structed_multi_dataset_router_agent.py

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

api/core/agent/agent/structured_chat.py

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

api/core/callback_handler/dataset_tool_callback_handler.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from json import JSONDecodeError
 
 from typing import Any, Dict, List, Union, Optional
 
@@ -44,10 +45,15 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
         input_str: str,
         **kwargs: Any,
     ) -> None:
-        # tool_name = serialized.get('name')
-        input_dict = json.loads(input_str.replace("'", "\""))
-        dataset_id = input_dict.get('dataset_id')
-        query = input_dict.get('query')
+        tool_name: str = serialized.get('name')
+        dataset_id = tool_name.removeprefix('dataset-')
+
+        try:
+            input_dict = json.loads(input_str.replace("'", "\""))
+            query = input_dict.get('query')
+        except JSONDecodeError:
+            query = input_str
+
         self.conversation_message_task.on_dataset_query_end(DatasetQueryObj(dataset_id=dataset_id, query=query))
 
     def on_tool_end(
@@ -58,12 +64,9 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
         llm_prefix: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        # kwargs={'name': 'Search'}
-        # llm_prefix='Thought:'
-        # observation_prefix='Observation: '
-        # output='53 years'
         pass
 
+
     def on_tool_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:

api/core/callback_handler/index_tool_callback_handler.py

@@ -2,6 +2,7 @@ from typing import List
 
 from langchain.schema import Document
 
+from core.conversation_message_task import ConversationMessageTask
 from extensions.ext_database import db
 from models.dataset import DocumentSegment
 
@@ -9,8 +10,9 @@ from models.dataset import DocumentSegment
 class DatasetIndexToolCallbackHandler:
     """Callback handler for dataset tool."""
 
-    def __init__(self, dataset_id: str) -> None:
+    def __init__(self, dataset_id: str, conversation_message_task: ConversationMessageTask) -> None:
         self.dataset_id = dataset_id
+        self.conversation_message_task = conversation_message_task
 
     def on_tool_end(self, documents: List[Document]) -> None:
         """Handle tool end."""
@@ -27,3 +29,7 @@ class DatasetIndexToolCallbackHandler:
             )
 
             db.session.commit()
+
+    def return_retriever_resource_info(self, resource: List):
+        """Handle return_retriever_resource_info."""
+        self.conversation_message_task.on_dataset_query_finish(resource)

api/core/completion.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import re
 from typing import Optional, List, Union, Tuple
@@ -19,13 +20,15 @@ from core.orchestrator_rule_parser import OrchestratorRuleParser
 from core.prompt.prompt_builder import PromptBuilder
 from core.prompt.prompt_template import JinjaPromptTemplate
 from core.prompt.prompts import MORE_LIKE_THIS_GENERATE_PROMPT
+from models.dataset import DocumentSegment, Dataset, Document
 from models.model import App, AppModelConfig, Account, Conversation, Message, EndUser
 
 
 class Completion:
     @classmethod
     def generate(cls, task_id: str, app: App, app_model_config: AppModelConfig, query: str, inputs: dict,
-                 user: Union[Account, EndUser], conversation: Optional[Conversation], streaming: bool, is_override: bool = False):
+                 user: Union[Account, EndUser], conversation: Optional[Conversation], streaming: bool,
+                 is_override: bool = False, retriever_from: str = 'dev'):
         """
         errors: ProviderTokenNotInitError
         """
@@ -96,7 +99,6 @@ class Completion:
             should_use_agent = agent_executor.should_use_agent(query)
             if should_use_agent:
                 agent_execute_result = agent_executor.run(query)
-
         # run the final llm
         try:
             cls.run_final_llm(
@@ -118,7 +120,8 @@ class Completion:
             return
 
     @classmethod
-    def run_final_llm(cls, model_instance: BaseLLM, mode: str, app_model_config: AppModelConfig, query: str, inputs: dict,
+    def run_final_llm(cls, model_instance: BaseLLM, mode: str, app_model_config: AppModelConfig, query: str,
+                      inputs: dict,
                       agent_execute_result: Optional[AgentExecuteResult],
                       conversation_message_task: ConversationMessageTask,
                       memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]):
@@ -150,7 +153,6 @@ class Completion:
             callbacks=[LLMCallbackHandler(model_instance, conversation_message_task)],
             fake_response=fake_response
         )
-
         return response
 
     @classmethod

api/core/conversation_message_task.py

@@ -1,6 +1,6 @@
 import decimal
 import json
-from typing import Optional, Union
+from typing import Optional, Union, List
 
 from core.callback_handler.entity.agent_loop import AgentLoop
 from core.callback_handler.entity.dataset_query import DatasetQueryObj
@@ -15,7 +15,8 @@ from events.message_event import message_was_created
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import DatasetQuery
-from models.model import AppModelConfig, Conversation, Account, Message, EndUser, App, MessageAgentThought, MessageChain
+from models.model import AppModelConfig, Conversation, Account, Message, EndUser, App, MessageAgentThought, \
+    MessageChain, DatasetRetrieverResource
 
 
 class ConversationMessageTask:
@@ -41,6 +42,8 @@ class ConversationMessageTask:
 
         self.message = None
 
+        self.retriever_resource = None
+
         self.model_dict = self.app_model_config.model_dict
         self.provider_name = self.model_dict.get('provider')
         self.model_name = self.model_dict.get('name')
@@ -137,7 +140,8 @@ class ConversationMessageTask:
         db.session.flush()
 
     def append_message_text(self, text: str):
-        self._pub_handler.pub_text(text)
+        if text is not None:
+            self._pub_handler.pub_text(text)
 
     def save_message(self, llm_message: LLMMessage, by_stopped: bool = False):
         message_tokens = llm_message.prompt_tokens
@@ -156,7 +160,8 @@ class ConversationMessageTask:
         self.message.message_tokens = message_tokens
         self.message.message_unit_price = message_unit_price
         self.message.message_price_unit = message_price_unit
-        self.message.answer = PromptBuilder.process_template(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.answer_price_unit = answer_price_unit
@@ -255,7 +260,36 @@ class ConversationMessageTask:
 
         db.session.add(dataset_query)
 
+    def on_dataset_query_finish(self, resource: List):
+        if resource and len(resource) > 0:
+            for item in resource:
+                dataset_retriever_resource = DatasetRetrieverResource(
+                    message_id=self.message.id,
+                    position=item.get('position'),
+                    dataset_id=item.get('dataset_id'),
+                    dataset_name=item.get('dataset_name'),
+                    document_id=item.get('document_id'),
+                    document_name=item.get('document_name'),
+                    data_source_type=item.get('data_source_type'),
+                    segment_id=item.get('segment_id'),
+                    score=item.get('score') if 'score' in item else None,
+                    hit_count=item.get('hit_count') if 'hit_count' else None,
+                    word_count=item.get('word_count') if 'word_count' in item else None,
+                    segment_position=item.get('segment_position') if 'segment_position' in item else None,
+                    index_node_hash=item.get('index_node_hash') if 'index_node_hash' in item else None,
+                    content=item.get('content'),
+                    retriever_from=item.get('retriever_from'),
+                    created_by=self.user.id
+                )
+                db.session.add(dataset_retriever_resource)
+                db.session.flush()
+            self.retriever_resource = resource
+
+    def message_end(self):
+        self._pub_handler.pub_message_end(self.retriever_resource)
+
     def end(self):
+        self._pub_handler.pub_message_end(self.retriever_resource)
         self._pub_handler.pub_end()
 
 
@@ -349,6 +383,23 @@ class PubHandler:
             self.pub_end()
             raise ConversationTaskStoppedException()
 
+    def pub_message_end(self, retriever_resource: List):
+        content = {
+            'event': 'message_end',
+            'data': {
+                'task_id': self._task_id,
+                'message_id': self._message.id,
+                'mode': self._conversation.mode,
+                'conversation_id': self._conversation.id
+            }
+        }
+        if retriever_resource:
+            content['data']['retriever_resources'] = retriever_resource
+        redis_client.publish(self._channel, json.dumps(content))
+
+        if self._is_stopped():
+            self.pub_end()
+            raise ConversationTaskStoppedException()
 
     def pub_end(self):
         content = {

api/core/data_loader/file_extractor.py

@@ -6,7 +6,7 @@ import requests
 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.csv_loader 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
@@ -47,17 +47,18 @@ class FileExtractor:
                        upload_file: Optional[UploadFile] = None) -> Union[List[Document] | str]:
         input_file = Path(file_path)
         delimiter = '\n'
-        if input_file.suffix == '.xlsx':
+        file_extension = input_file.suffix.lower()
+        if file_extension == '.xlsx':
             loader = ExcelLoader(file_path)
-        elif input_file.suffix == '.pdf':
+        elif file_extension == '.pdf':
             loader = PdfLoader(file_path, upload_file=upload_file)
-        elif input_file.suffix in ['.md', '.markdown']:
+        elif file_extension in ['.md', '.markdown']:
             loader = MarkdownLoader(file_path, autodetect_encoding=True)
-        elif input_file.suffix in ['.htm', '.html']:
+        elif file_extension in ['.htm', '.html']:
             loader = HTMLLoader(file_path)
-        elif input_file.suffix == '.docx':
+        elif file_extension == '.docx':
             loader = Docx2txtLoader(file_path)
-        elif input_file.suffix == '.csv':
+        elif file_extension == '.csv':
             loader = CSVLoader(file_path, autodetect_encoding=True)
         else:
             # txt

api/core/data_loader/loader/csv_loader.py

@@ -1,10 +1,10 @@
 import logging
+import csv
 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
+from langchain.schema import Document
 
 logger = logging.getLogger(__name__)
 

api/core/data_loader/loader/excel.py

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

api/core/docstore/dataset_docstore.py

@@ -67,12 +67,13 @@ class DatesetDocumentStore:
 
         if max_position is None:
             max_position = 0
-
-        embedding_model = ModelFactory.get_embedding_model(
-            tenant_id=self._dataset.tenant_id,
-            model_provider_name=self._dataset.embedding_model_provider,
-            model_name=self._dataset.embedding_model
-        )
+        embedding_model = None
+        if self._dataset.indexing_technique == 'high_quality':
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=self._dataset.tenant_id,
+                model_provider_name=self._dataset.embedding_model_provider,
+                model_name=self._dataset.embedding_model
+            )
 
         for doc in docs:
             if not isinstance(doc, Document):
@@ -88,7 +89,7 @@ class DatesetDocumentStore:
                 )
 
             # calc embedding use tokens
-            tokens = embedding_model.get_num_tokens(doc.page_content)
+            tokens = embedding_model.get_num_tokens(doc.page_content) if embedding_model else 0
 
             if not segment_document:
                 max_position += 1

api/core/generator/llm_generator.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from langchain.schema import OutputParserException
@@ -22,18 +23,25 @@ class LLMGenerator:
         if len(query) > 2000:
             query = query[:300] + "...[TRUNCATED]..." + query[-300:]
 
-        prompt = prompt.format(query=query)
+        query = query.replace("\n", " ")
+
+        prompt += query + "\n"
 
         model_instance = ModelFactory.get_text_generation_model(
             tenant_id=tenant_id,
             model_kwargs=ModelKwargs(
-                max_tokens=50
+                temperature=1,
+                max_tokens=100
             )
         )
 
         prompts = [PromptMessage(content=prompt)]
         response = model_instance.run(prompts)
         answer = response.content
+
+        result_dict = json.loads(answer)
+        answer = result_dict['Your Output']
+
         return answer.strip()
 
     @classmethod

api/core/index/index.py

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

api/core/index/keyword_table_index/keyword_table_index.py

@@ -25,7 +25,7 @@ class KeywordTableIndex(BaseIndex):
         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))
+            self._update_segment_keywords(self.dataset.id, 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(
@@ -52,7 +52,7 @@ class KeywordTableIndex(BaseIndex):
         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))
+            self._update_segment_keywords(self.dataset.id, 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)
@@ -74,7 +74,7 @@ class KeywordTableIndex(BaseIndex):
             DocumentSegment.document_id == document_id
         ).all()
 
-        ids = [segment.id for segment in segments]
+        ids = [segment.index_node_id for segment in segments]
 
         keyword_table = self._get_dataset_keyword_table()
         keyword_table = self._delete_ids_from_keyword_table(keyword_table, ids)
@@ -199,15 +199,18 @@ class KeywordTableIndex(BaseIndex):
 
         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()
+    def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: List[str]):
+        document_segment = db.session.query(DocumentSegment).filter(
+            DocumentSegment.dataset_id == dataset_id,
+            DocumentSegment.index_node_id == node_id
+        ).first()
         if document_segment:
             document_segment.keywords = keywords
             db.session.commit()
 
     def create_segment_keywords(self, node_id: str, keywords: List[str]):
         keyword_table = self._get_dataset_keyword_table()
-        self._update_segment_keywords(node_id, keywords)
+        self._update_segment_keywords(self.dataset.id, node_id, keywords)
         keyword_table = self._add_text_to_keyword_table(keyword_table, node_id, keywords)
         self._save_dataset_keyword_table(keyword_table)
 

api/core/index/vector_index/base.py

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

api/core/index/vector_index/milvus_vector_index.py

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

api/core/index/vector_index/qdrant_vector_index.py

@@ -44,15 +44,20 @@ class QdrantVectorIndex(BaseVectorIndex):
 
     def get_index_name(self, dataset: Dataset) -> str:
         if self.dataset.index_struct_dict:
-            return self.dataset.index_struct_dict['vector_store']['collection_name']
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
+                # original class_prefix
+                class_prefix += '_Node'
+
+            return class_prefix
 
         dataset_id = dataset.id
-        return "Index_" + dataset_id.replace("-", "_")
+        return "Vector_index_" + dataset_id.replace("-", "_") + '_Node'
 
     def to_index_struct(self) -> dict:
         return {
             "type": self.get_type(),
-            "vector_store": {"collection_name": self.get_index_name(self.dataset)}
+            "vector_store": {"class_prefix": self.get_index_name(self.dataset)}
         }
 
     def create(self, texts: list[Document], **kwargs) -> BaseIndex:
@@ -62,7 +67,7 @@ class QdrantVectorIndex(BaseVectorIndex):
             self._embeddings,
             collection_name=self.get_index_name(self.dataset),
             ids=uuids,
-            content_payload_key='text',
+            content_payload_key='page_content',
             **self._client_config.to_qdrant_params()
         )
 
@@ -72,7 +77,9 @@ class QdrantVectorIndex(BaseVectorIndex):
         """Only for created index."""
         if self._vector_store:
             return self._vector_store
-        
+        attributes = ['doc_id', 'dataset_id', 'document_id']
+        if self._is_origin():
+            attributes = ['doc_id']
         client = qdrant_client.QdrantClient(
             **self._client_config.to_qdrant_params()
         )
@@ -81,7 +88,7 @@ class QdrantVectorIndex(BaseVectorIndex):
             client=client,
             collection_name=self.get_index_name(self.dataset),
             embeddings=self._embeddings,
-            content_payload_key='text'
+            content_payload_key='page_content'
         )
 
     def _get_vector_store_class(self) -> type:
@@ -106,10 +113,29 @@ class QdrantVectorIndex(BaseVectorIndex):
             ],
         ))
 
+    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)
+
+        from qdrant_client.http import models
+        for node_id in ids:
+            vector_store.del_texts(models.Filter(
+                must=[
+                    models.FieldCondition(
+                        key="metadata.doc_id",
+                        match=models.MatchValue(value=node_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_'):
+            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
+            if not class_prefix.endswith('_Node'):
                 # original class_prefix
                 return True
 

api/core/indexing_runner.py

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

api/core/model_providers/model_provider_factory.py

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

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

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

api/core/model_providers/models/entity/message.py

@@ -8,6 +8,7 @@ class LLMRunResult(BaseModel):
     content: str
     prompt_tokens: int
     completion_tokens: int
+    source: list = None
 
 
 class MessageType(enum.Enum):

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

@@ -1,11 +1,8 @@
-import decimal
 import logging
-from functools import wraps
 from typing import List, Optional, Any
 
 import anthropic
 from langchain.callbacks.manager import Callbacks
-from langchain.chat_models import ChatAnthropic
 from langchain.schema import LLMResult
 
 from core.model_providers.error import LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError, \
@@ -13,6 +10,7 @@ from core.model_providers.error import LLMBadRequestError, LLMAPIConnectionError
 from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.entity.message import PromptMessage, MessageType
 from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.anthropic_llm import AnthropicLLM
 
 
 class AnthropicModel(BaseLLM):
@@ -20,7 +18,7 @@ class AnthropicModel(BaseLLM):
 
     def _init_client(self) -> Any:
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
-        return ChatAnthropic(
+        return AnthropicLLM(
             model=self.name,
             streaming=self.streaming,
             callbacks=self.callbacks,
@@ -75,7 +73,7 @@ class AnthropicModel(BaseLLM):
         else:
             return ex
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True
 

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

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

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

@@ -138,7 +138,7 @@ class BaseLLM(BaseProviderModel):
                 result = self._run(
                     messages=messages,
                     stop=stop,
-                    callbacks=callbacks if not (self.streaming and not self.support_streaming()) else None,
+                    callbacks=callbacks if not (self.streaming and not self.support_streaming) else None,
                     **kwargs
                 )
             except Exception as ex:
@@ -149,7 +149,7 @@ class BaseLLM(BaseProviderModel):
         else:
             completion_content = result.generations[0][0].text
 
-        if self.streaming and not self.support_streaming():
+        if self.streaming and not self.support_streaming:
             # use FakeLLM to simulate streaming when current model not support streaming but streaming is True
             prompts = self._get_prompt_from_messages(messages, ModelMode.CHAT)
             fake_llm = FakeLLM(
@@ -298,8 +298,8 @@ class BaseLLM(BaseProviderModel):
         else:
             self.client.callbacks.extend(callbacks)
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return False
 
     def get_prompt(self, mode: str,
@@ -342,7 +342,7 @@ class BaseLLM(BaseProviderModel):
             if order == 'context_prompt':
                 prompt += context_prompt_content
             elif order == 'pre_prompt':
-                prompt += (pre_prompt_content + '\n\n') if pre_prompt_content else ''
+                prompt += pre_prompt_content
 
         query_prompt = prompt_rules['query_prompt'] if 'query_prompt' in prompt_rules else '{{query}}'
 

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

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

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

@@ -1,6 +1,5 @@
 from typing import List, Optional, Any
 
-from langchain import HuggingFaceHub
 from langchain.callbacks.manager import Callbacks
 from langchain.schema import LLMResult
 
@@ -9,6 +8,7 @@ from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
 from core.third_party.langchain.llms.huggingface_endpoint_llm import HuggingFaceEndpointLLM
+from core.third_party.langchain.llms.huggingface_hub_llm import HuggingFaceHubLLM
 
 
 class HuggingfaceHubModel(BaseLLM):
@@ -17,15 +17,21 @@ class HuggingfaceHubModel(BaseLLM):
     def _init_client(self) -> Any:
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
         if self.credentials['huggingfacehub_api_type'] == 'inference_endpoints':
+            streaming = self.streaming
+
+            if 'baichuan' in self.name.lower():
+                streaming = False
+
             client = HuggingFaceEndpointLLM(
                 endpoint_url=self.credentials['huggingfacehub_endpoint_url'],
                 task=self.credentials['task_type'],
                 model_kwargs=provider_model_kwargs,
                 huggingfacehub_api_token=self.credentials['huggingfacehub_api_token'],
-                callbacks=self.callbacks
+                callbacks=self.callbacks,
+                streaming=streaming
             )
         else:
-            client = HuggingFaceHub(
+            client = HuggingFaceHubLLM(
                 repo_id=self.name,
                 task=self.credentials['task_type'],
                 model_kwargs=provider_model_kwargs,
@@ -76,7 +82,12 @@ class HuggingfaceHubModel(BaseLLM):
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"Huggingface Hub: {str(ex)}")
 
-    @classmethod
-    def support_streaming(cls):
-        return False
+    @property
+    def support_streaming(self):
+        if self.credentials['huggingfacehub_api_type'] == 'inference_endpoints':
+            if 'baichuan' in self.name.lower():
+                return False
 
+            return True
+        else:
+            return False

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -74,6 +74,6 @@ class XinferenceModel(BaseLLM):
     def handle_exceptions(self, ex: Exception) -> Exception:
         return LLMBadRequestError(f"Xinference: {str(ex)}")
 
-    @classmethod
-    def support_streaming(cls):
+    @property
+    def support_streaming(self):
         return True

api/core/model_providers/providers/anthropic_provider.py

@@ -5,7 +5,6 @@ from typing import Type, Optional
 
 import anthropic
 from flask import current_app
-from langchain.chat_models import ChatAnthropic
 from langchain.schema import HumanMessage
 
 from core.helper import encrypter
@@ -16,6 +15,7 @@ from core.model_providers.models.llm.anthropic_model import AnthropicModel
 from core.model_providers.models.llm.base import ModelType
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 from core.model_providers.providers.hosted import hosted_model_providers
+from core.third_party.langchain.llms.anthropic_llm import AnthropicLLM
 from models.provider import ProviderType
 
 
@@ -92,7 +92,7 @@ class AnthropicProvider(BaseModelProvider):
             if 'anthropic_api_url' in credentials:
                 credential_kwargs['anthropic_api_url'] = credentials['anthropic_api_url']
 
-            chat_llm = ChatAnthropic(
+            chat_llm = AnthropicLLM(
                 model='claude-instant-1',
                 max_tokens_to_sample=10,
                 temperature=0,

api/core/model_providers/providers/huggingface_hub_provider.py

@@ -89,7 +89,8 @@ class HuggingfaceHubProvider(BaseModelProvider):
                 raise CredentialsValidateFailedError('Task Type must be provided.')
 
             if credentials['task_type'] not in ("text2text-generation", "text-generation", "summarization"):
-                raise CredentialsValidateFailedError('Task Type must be one of text2text-generation, text-generation, summarization.')
+                raise CredentialsValidateFailedError('Task Type must be one of text2text-generation, '
+                                                     'text-generation, summarization.')
 
             try:
                 llm = HuggingFaceEndpointLLM(

api/core/model_providers/providers/localai_provider.py

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

api/core/model_providers/providers/spark_provider.py

@@ -83,14 +83,15 @@ class SparkProvider(BaseModelProvider):
         if 'api_secret' not in credentials:
             raise CredentialsValidateFailedError('Spark api_secret must be provided.')
 
-        try:
-            credential_kwargs = {
-                'app_id': credentials['app_id'],
-                'api_key': credentials['api_key'],
-                'api_secret': credentials['api_secret'],
-            }
+        credential_kwargs = {
+            'app_id': credentials['app_id'],
+            'api_key': credentials['api_key'],
+            'api_secret': credentials['api_secret'],
+        }
 
+        try:
             chat_llm = ChatSpark(
+                model_name='spark-v2',
                 max_tokens=10,
                 temperature=0.01,
                 **credential_kwargs
@@ -104,7 +105,27 @@ class SparkProvider(BaseModelProvider):
 
             chat_llm(messages)
         except SparkError as ex:
-            raise CredentialsValidateFailedError(str(ex))
+            # try spark v1.5 if v2.1 failed
+            try:
+                chat_llm = ChatSpark(
+                    model_name='spark',
+                    max_tokens=10,
+                    temperature=0.01,
+                    **credential_kwargs
+                )
+
+                messages = [
+                    HumanMessage(
+                        content="ping"
+                    )
+                ]
+
+                chat_llm(messages)
+            except SparkError as ex:
+                raise CredentialsValidateFailedError(str(ex))
+            except Exception as ex:
+                logging.exception('Spark config validation failed')
+                raise ex
         except Exception as ex:
             logging.exception('Spark config validation failed')
             raise ex

api/core/model_providers/rules/_providers.json

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

api/core/model_providers/rules/localai.json

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

api/core/orchestrator_rule_parser.py

@@ -36,8 +36,8 @@ class OrchestratorRuleParser:
         self.app_model_config = app_model_config
 
     def to_agent_executor(self, conversation_message_task: ConversationMessageTask, memory: Optional[BaseChatMemory],
-                          rest_tokens: int, chain_callback: MainChainGatherCallbackHandler) \
-            -> Optional[AgentExecutor]:
+                          rest_tokens: int, chain_callback: MainChainGatherCallbackHandler,
+                          return_resource: bool = False, retriever_from: str = 'dev') -> Optional[AgentExecutor]:
         if not self.app_model_config.agent_mode_dict:
             return None
 
@@ -74,7 +74,7 @@ class OrchestratorRuleParser:
 
             # only OpenAI chat model (include Azure) support function call, use ReACT instead
             if agent_model_instance.model_mode != ModelMode.CHAT \
-                         or agent_model_instance.model_provider.provider_name not in ['openai', 'azure_openai']:
+                    or agent_model_instance.model_provider.provider_name not in ['openai', 'azure_openai']:
                 if planning_strategy in [PlanningStrategy.FUNCTION_CALL, PlanningStrategy.MULTI_FUNCTION_CALL]:
                     planning_strategy = PlanningStrategy.REACT
                 elif planning_strategy == PlanningStrategy.ROUTER:
@@ -99,7 +99,9 @@ class OrchestratorRuleParser:
                 tool_configs=tool_configs,
                 conversation_message_task=conversation_message_task,
                 rest_tokens=rest_tokens,
-                callbacks=[agent_callback, DifyStdOutCallbackHandler()]
+                callbacks=[agent_callback, DifyStdOutCallbackHandler()],
+                return_resource=return_resource,
+                retriever_from=retriever_from
             )
 
             if len(tools) == 0:
@@ -145,8 +147,10 @@ class OrchestratorRuleParser:
 
         return None
 
-    def to_tools(self, agent_model_instance: BaseLLM, tool_configs: list, conversation_message_task: ConversationMessageTask,
-                 rest_tokens: int, callbacks: Callbacks = None) -> list[BaseTool]:
+    def to_tools(self, agent_model_instance: BaseLLM, tool_configs: list,
+                 conversation_message_task: ConversationMessageTask,
+                 rest_tokens: int, callbacks: Callbacks = None, return_resource: bool = False,
+                 retriever_from: str = 'dev') -> list[BaseTool]:
         """
         Convert app agent tool configs to tools
 
@@ -155,6 +159,8 @@ class OrchestratorRuleParser:
         :param tool_configs: app agent tool configs
         :param conversation_message_task:
         :param callbacks:
+        :param return_resource:
+        :param retriever_from:
         :return:
         """
         tools = []
@@ -166,7 +172,7 @@ class OrchestratorRuleParser:
 
             tool = None
             if tool_type == "dataset":
-                tool = self.to_dataset_retriever_tool(tool_val, conversation_message_task, rest_tokens)
+                tool = self.to_dataset_retriever_tool(tool_val, conversation_message_task, rest_tokens, return_resource, retriever_from)
             elif tool_type == "web_reader":
                 tool = self.to_web_reader_tool(agent_model_instance)
             elif tool_type == "google_search":
@@ -183,13 +189,15 @@ class OrchestratorRuleParser:
         return tools
 
     def to_dataset_retriever_tool(self, tool_config: dict, conversation_message_task: ConversationMessageTask,
-                                  rest_tokens: int) \
+                                  rest_tokens: int, return_resource: bool = False, retriever_from: str = 'dev') \
             -> Optional[BaseTool]:
         """
         A dataset tool is a tool that can be used to retrieve information from a dataset
         :param rest_tokens:
         :param tool_config:
         :param conversation_message_task:
+        :param return_resource:
+        :param retriever_from:
         :return:
         """
         # get dataset from dataset id
@@ -208,7 +216,10 @@ class OrchestratorRuleParser:
         tool = DatasetRetrieverTool.from_dataset(
             dataset=dataset,
             k=k,
-            callbacks=[DatasetToolCallbackHandler(conversation_message_task)]
+            callbacks=[DatasetToolCallbackHandler(conversation_message_task)],
+            conversation_message_task=conversation_message_task,
+            return_resource=return_resource,
+            retriever_from=retriever_from
         )
 
         return tool
@@ -283,6 +294,7 @@ class OrchestratorRuleParser:
     def _dynamic_calc_retrieve_k(cls, dataset: Dataset, rest_tokens: int) -> int:
         DEFAULT_K = 2
         CONTEXT_TOKENS_PERCENT = 0.3
+        MAX_K = 10
 
         if rest_tokens == -1:
             return DEFAULT_K
@@ -311,5 +323,5 @@ class OrchestratorRuleParser:
         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
+        # Expand the k value when there's still some room left in the 30% rest tokens space, but less than the MAX_K
+        return min(context_limit_tokens // segment_max_tokens, MAX_K)

api/core/prompt/generate_prompts/baichuan_chat.json

@@ -8,6 +8,6 @@
     "pre_prompt",
     "histories_prompt"
   ],
-  "query_prompt": "用户：{{query}}",
+  "query_prompt": "\n\n用户：{{query}}",
   "stops": ["用户："]
 }

api/core/prompt/generate_prompts/common_chat.json

@@ -8,6 +8,6 @@
     "pre_prompt",
     "histories_prompt"
   ],
-  "query_prompt": "Human: {{query}}\n\nAssistant: ",
+  "query_prompt": "\n\nHuman: {{query}}\n\nAssistant: ",
   "stops": ["\nHuman:", "</histories>"]
-}
+}

api/core/prompt/prompts.py

@@ -1,10 +1,65 @@
-CONVERSATION_TITLE_PROMPT = (
-    "Human:{query}\n-----\n"
-    "Help me summarize the intent of what the human said and provide a title, the title should not exceed 20 words.\n"
-    "If what the human said is conducted in English, you should only return an English title.\n" 
-    "If what the human said is conducted in Chinese, you should only return a Chinese title.\n"
-    "title:"
-)
+# Written by YORKI MINAKO🤡
+CONVERSATION_TITLE_PROMPT = """You need to decompose the user's input into "subject" and "intention" in order to accurately figure out what the user's input language actually is. 
+Notice: the language type user use could be diverse, which can be English, Chinese, Español, Arabic, Japanese, French, and etc.
+MAKE SURE your output is the SAME language as the user's input!
+Your output is restricted only to: (Input language) Intention + Subject(short as possible)
+Your output MUST be a valid JSON.
+
+Tip: When the user's question is directed at you (the language model), you can add an emoji to make it more fun.
+
+
+example 1:
+User Input: hi, yesterday i had some burgers.
+{
+  "Language Type": "The user's input is pure English",
+  "Your Reasoning": "The language of my output must be pure English.",
+  "Your Output": "sharing yesterday's food"
+}
+
+example 2:
+User Input: hello
+{
+  "Language Type": "The user's input is written in pure English",
+  "Your Reasoning": "The language of my output must be pure English.",
+  "Your Output": "Greeting myself☺️"
+}
+
+
+example 3:
+User Input: why mmap file: oom
+{
+  "Language Type": "The user's input is written in pure English",
+  "Your Reasoning": "The language of my output must be pure English.",
+  "Your Output": "Asking about the reason for mmap file: oom"
+}
+
+
+example 4:
+User Input: www.convinceme.yesterday-you-ate-seafood.tv讲了什么？
+{
+  "Language Type": "The user's input English-Chinese mixed",
+  "Your Reasoning": "The English-part is an URL, the main intention is still written in Chinese, so the language of my output must be using Chinese.",
+  "Your Output": "询问网站www.convinceme.yesterday-you-ate-seafood.tv"
+}
+
+example 5:
+User Input: why小红的年龄is老than小明？
+{
+  "Language Type": "The user's input is English-Chinese mixed",
+  "Your Reasoning": "The English parts are subjective particles, the main intention is written in Chinese, besides, Chinese occupies a greater \"actual meaning\" than English, so the language of my output must be using Chinese.",
+  "Your Output": "询问小红和小明的年龄"
+}
+
+example 6:
+User Input: yo, 你今天咋样？
+{
+  "Language Type": "The user's input is English-Chinese mixed",
+  "Your Reasoning": "The English-part is a subjective particle, the main intention is written in Chinese, so the language of my output must be using Chinese.",
+  "Your Output": "查询今日我的状态☺️"
+}
+
+User Input: 
+"""
 
 CONVERSATION_SUMMARY_PROMPT = (
     "Please generate a short summary of the following conversation.\n"
@@ -50,7 +105,7 @@ GENERATOR_QA_PROMPT = (
     'Step 3: Decompose or combine multiple pieces of information and concepts.\n'
     'Step 4: Generate 20 questions and answers based on these key information and concepts.'
     'The questions should be clear and detailed, and the answers should be detailed and complete.\n'
-    "Answer must be the language:{language} and in the following format: Q1:\nA1:\nQ2:\nA2:...\n"
+    "Answer according to the the language:{language} and in the following format: Q1:\nA1:\nQ2:\nA2:...\n"
 )
 
 RULE_CONFIG_GENERATE_TEMPLATE = """Given MY INTENDED AUDIENCES and HOPING TO SOLVE using a language model, please select \

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

@@ -0,0 +1,48 @@
+from typing import Dict
+
+from httpx import Limits
+from langchain.chat_models import ChatAnthropic
+from langchain.utils import get_from_dict_or_env, check_package_version
+from pydantic import root_validator
+
+
+class AnthropicLLM(ChatAnthropic):
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values["anthropic_api_key"] = get_from_dict_or_env(
+            values, "anthropic_api_key", "ANTHROPIC_API_KEY"
+        )
+        # Get custom api url from environment.
+        values["anthropic_api_url"] = get_from_dict_or_env(
+            values,
+            "anthropic_api_url",
+            "ANTHROPIC_API_URL",
+            default="https://api.anthropic.com",
+        )
+
+        try:
+            import anthropic
+
+            check_package_version("anthropic", gte_version="0.3")
+            values["client"] = anthropic.Anthropic(
+                base_url=values["anthropic_api_url"],
+                api_key=values["anthropic_api_key"],
+                timeout=values["default_request_timeout"],
+                max_retries=0,
+                connection_pool_limits=Limits(max_connections=200, max_keepalive_connections=100),
+            )
+            values["async_client"] = anthropic.AsyncAnthropic(
+                base_url=values["anthropic_api_url"],
+                api_key=values["anthropic_api_key"],
+                timeout=values["default_request_timeout"],
+            )
+            values["HUMAN_PROMPT"] = anthropic.HUMAN_PROMPT
+            values["AI_PROMPT"] = anthropic.AI_PROMPT
+            values["count_tokens"] = values["client"].count_tokens
+        except ImportError:
+            raise ImportError(
+                "Could not import anthropic python package. "
+                "Please it install it with `pip install anthropic`."
+            )
+        return values

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

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

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

@@ -1,7 +1,11 @@
-from typing import Dict
+from typing import Dict, Any, Optional, List, Iterable, Iterator
 
+from huggingface_hub import InferenceClient
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.embeddings.huggingface_hub import VALID_TASKS
 from langchain.llms import HuggingFaceEndpoint
-from pydantic import Extra, root_validator
+from langchain.llms.utils import enforce_stop_tokens
+from pydantic import root_validator
 
 from langchain.utils import get_from_dict_or_env
 
@@ -27,6 +31,8 @@ class HuggingFaceEndpointLLM(HuggingFaceEndpoint):
                 huggingfacehub_api_token="my-api-key"
             )
     """
+    client: Any
+    streaming: bool = False
 
     @root_validator(allow_reuse=True)
     def validate_environment(cls, values: Dict) -> Dict:
@@ -35,5 +41,88 @@ class HuggingFaceEndpointLLM(HuggingFaceEndpoint):
             values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
         )
 
+        values['client'] = InferenceClient(values['endpoint_url'], token=huggingfacehub_api_token)
+
         values["huggingfacehub_api_token"] = huggingfacehub_api_token
         return values
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to HuggingFace Hub's inference endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+
+        Example:
+            .. code-block:: python
+
+                response = hf("Tell me a joke.")
+        """
+        _model_kwargs = self.model_kwargs or {}
+
+        # payload samples
+        params = {**_model_kwargs, **kwargs}
+
+        # generation parameter
+        gen_kwargs = {
+            **params,
+            'stop_sequences': stop
+        }
+
+        response = self.client.text_generation(prompt, stream=self.streaming, details=True, **gen_kwargs)
+
+        if self.streaming and isinstance(response, Iterable):
+            combined_text_output = ""
+            for token in self._stream_response(response, run_manager):
+                combined_text_output += token
+            completion = combined_text_output
+        else:
+            completion = response.generated_text
+
+        if self.task == "text-generation":
+            text = completion
+            # Remove prompt if included in generated text.
+            if text.startswith(prompt):
+                text = text[len(prompt) :]
+        elif self.task == "text2text-generation":
+            text = completion
+        else:
+            raise ValueError(
+                f"Got invalid task {self.task}, "
+                f"currently only {VALID_TASKS} are supported"
+            )
+
+        if stop is not None:
+            # This is a bit hacky, but I can't figure out a better way to enforce
+            # stop tokens when making calls to huggingface_hub.
+            text = enforce_stop_tokens(text, stop)
+
+        return text
+
+    def _stream_response(
+            self,
+            response: Iterable,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+    ) -> Iterator[str]:
+        for r in response:
+            # skip special tokens
+            if r.token.special:
+                continue
+
+            token = r.token.text
+            if run_manager:
+                run_manager.on_llm_new_token(
+                    token=token, verbose=self.verbose, log_probs=None
+                )
+
+            # yield the generated token
+            yield token

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

@@ -0,0 +1,62 @@
+from typing import Dict, Optional, List, Any
+
+from huggingface_hub import HfApi, InferenceApi
+from langchain import HuggingFaceHub
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.huggingface_hub import VALID_TASKS
+from pydantic import root_validator
+
+from langchain.utils import get_from_dict_or_env
+
+
+class HuggingFaceHubLLM(HuggingFaceHub):
+    """HuggingFaceHub  models.
+
+    To use, you should have the ``huggingface_hub`` python package installed, and the
+    environment variable ``HUGGINGFACEHUB_API_TOKEN`` set with your API token, or pass
+    it as a named parameter to the constructor.
+
+    Only supports `text-generation`, `text2text-generation` and `summarization` for now.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.llms import HuggingFaceHub
+            hf = HuggingFaceHub(repo_id="gpt2", huggingfacehub_api_token="my-api-key")
+    """
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        huggingfacehub_api_token = get_from_dict_or_env(
+            values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
+        )
+        client = InferenceApi(
+            repo_id=values["repo_id"],
+            token=huggingfacehub_api_token,
+            task=values.get("task"),
+        )
+        client.options = {"wait_for_model": False, "use_gpu": False}
+        values["client"] = client
+        return values
+
+    def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> str:
+        hfapi = HfApi(token=self.huggingfacehub_api_token)
+        model_info = hfapi.model_info(repo_id=self.repo_id)
+        if not model_info:
+            raise ValueError(f"Model {self.repo_id} not found.")
+
+        if 'inference' in model_info.cardData and not model_info.cardData['inference']:
+            raise ValueError(f"Inference API has been turned off for this model {self.repo_id}.")
+
+        if model_info.pipeline_tag not in VALID_TASKS:
+            raise ValueError(f"Model {self.repo_id} is not a valid task, "
+                             f"must be one of {VALID_TASKS}.")
+
+        return super()._call(prompt, stop, run_manager, **kwargs)

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

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

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

@@ -3,17 +3,20 @@ from typing import Optional, List, Any, Union, Generator
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 from langchain.llms import Xinference
 from langchain.llms.utils import enforce_stop_tokens
-from xinference.client import RESTfulChatglmCppChatModelHandle, \
-    RESTfulChatModelHandle, RESTfulGenerateModelHandle
+from xinference.client import (
+    RESTfulChatglmCppChatModelHandle,
+    RESTfulChatModelHandle,
+    RESTfulGenerateModelHandle,
+)
 
 
 class XinferenceLLM(Xinference):
     def _call(
-            self,
-            prompt: str,
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            **kwargs: Any,
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
     ) -> str:
         """Call the xinference model and return the output.
 
@@ -29,7 +32,9 @@ class XinferenceLLM(Xinference):
         model = self.client.get_model(self.model_uid)
 
         if isinstance(model, RESTfulChatModelHandle):
-            generate_config: "LlamaCppGenerateConfig" = kwargs.get("generate_config", {})
+            generate_config: "LlamaCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
 
             if stop:
                 generate_config["stop"] = stop
@@ -37,10 +42,10 @@ class XinferenceLLM(Xinference):
             if generate_config and generate_config.get("stream"):
                 combined_text_output = ""
                 for token in self._stream_generate(
-                        model=model,
-                        prompt=prompt,
-                        run_manager=run_manager,
-                        generate_config=generate_config,
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
                 ):
                     combined_text_output += token
                 return combined_text_output
@@ -48,7 +53,9 @@ class XinferenceLLM(Xinference):
                 completion = model.chat(prompt=prompt, generate_config=generate_config)
                 return completion["choices"][0]["message"]["content"]
         elif isinstance(model, RESTfulGenerateModelHandle):
-            generate_config: "LlamaCppGenerateConfig" = kwargs.get("generate_config", {})
+            generate_config: "LlamaCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
 
             if stop:
                 generate_config["stop"] = stop
@@ -56,27 +63,31 @@ class XinferenceLLM(Xinference):
             if generate_config and generate_config.get("stream"):
                 combined_text_output = ""
                 for token in self._stream_generate(
-                        model=model,
-                        prompt=prompt,
-                        run_manager=run_manager,
-                        generate_config=generate_config,
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
                 ):
                     combined_text_output += token
                 return combined_text_output
 
             else:
-                completion = model.generate(prompt=prompt, generate_config=generate_config)
+                completion = model.generate(
+                    prompt=prompt, generate_config=generate_config
+                )
                 return completion["choices"][0]["text"]
         elif isinstance(model, RESTfulChatglmCppChatModelHandle):
-            generate_config: "ChatglmCppGenerateConfig" = kwargs.get("generate_config", {})
+            generate_config: "ChatglmCppGenerateConfig" = kwargs.get(
+                "generate_config", {}
+            )
 
             if generate_config and generate_config.get("stream"):
                 combined_text_output = ""
                 for token in self._stream_generate(
-                        model=model,
-                        prompt=prompt,
-                        run_manager=run_manager,
-                        generate_config=generate_config,
+                    model=model,
+                    prompt=prompt,
+                    run_manager=run_manager,
+                    generate_config=generate_config,
                 ):
                     combined_text_output += token
                 completion = combined_text_output
@@ -90,12 +101,21 @@ class XinferenceLLM(Xinference):
             return completion
 
     def _stream_generate(
-            self,
-            model: Union["RESTfulGenerateModelHandle", "RESTfulChatModelHandle", "RESTfulChatglmCppChatModelHandle"],
-            prompt: str,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            generate_config: Optional[
-                Union["LlamaCppGenerateConfig", "PytorchGenerateConfig", "ChatglmCppGenerateConfig"]] = None,
+        self,
+        model: Union[
+            "RESTfulGenerateModelHandle",
+            "RESTfulChatModelHandle",
+            "RESTfulChatglmCppChatModelHandle",
+        ],
+        prompt: str,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        generate_config: Optional[
+            Union[
+                "LlamaCppGenerateConfig",
+                "PytorchGenerateConfig",
+                "ChatglmCppGenerateConfig",
+            ]
+        ] = None,
     ) -> Generator[str, None, None]:
         """
         Args:
@@ -108,7 +128,9 @@ class XinferenceLLM(Xinference):
         Yields:
             A string token.
         """
-        if isinstance(model, (RESTfulChatModelHandle, RESTfulChatglmCppChatModelHandle)):
+        if isinstance(
+            model, (RESTfulChatModelHandle, RESTfulChatglmCppChatModelHandle)
+        ):
             streaming_response = model.chat(
                 prompt=prompt, generate_config=generate_config
             )
@@ -123,14 +145,10 @@ class XinferenceLLM(Xinference):
                 if choices:
                     choice = choices[0]
                     if isinstance(choice, dict):
-                        if 'finish_reason' in choice and choice['finish_reason'] \
-                                and choice['finish_reason'] in ['stop', 'length']:
-                            break
-
-                        if 'text' in choice:
+                        if "text" in choice:
                             token = choice.get("text", "")
-                        elif 'delta' in choice and 'content' in choice['delta']:
-                            token = choice.get('delta').get('content')
+                        elif "delta" in choice and "content" in choice["delta"]:
+                            token = choice.get("delta").get("content")
                         else:
                             continue
                         log_probs = choice.get("logprobs")

api/core/tool/dataset_retriever_tool.py

@@ -1,4 +1,4 @@
-import re
+import json
 from typing import Type
 
 from flask import current_app
@@ -6,17 +6,17 @@ from langchain.tools import BaseTool
 from pydantic import Field, BaseModel
 
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
+from core.conversation_message_task import ConversationMessageTask
 from core.embedding.cached_embedding import CacheEmbedding
 from core.index.keyword_table_index.keyword_table_index import KeywordTableIndex, KeywordTableConfig
 from core.index.vector_index.vector_index import VectorIndex
 from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_providers.model_factory import ModelFactory
 from extensions.ext_database import db
-from models.dataset import Dataset, DocumentSegment
+from models.dataset import Dataset, DocumentSegment, Document
 
 
 class DatasetRetrieverToolInput(BaseModel):
-    dataset_id: str = Field(..., description="ID of dataset to be queried. MUST be UUID format.")
     query: str = Field(..., description="Query for the dataset to be used to retrieve the dataset.")
 
 
@@ -29,6 +29,10 @@ class DatasetRetrieverTool(BaseTool):
     tenant_id: str
     dataset_id: str
     k: int = 3
+    conversation_message_task: ConversationMessageTask
+    return_resource: str
+    retriever_from: str
+
 
     @classmethod
     def from_dataset(cls, dataset: Dataset, **kwargs):
@@ -37,27 +41,22 @@ class DatasetRetrieverTool(BaseTool):
             description = 'useful for when you want to answer queries about the ' + dataset.name
 
         description = description.replace('\n', '').replace('\r', '')
-        description += '\nID of dataset MUST be ' + dataset.id
         return cls(
+            name=f'dataset-{dataset.id}',
             tenant_id=dataset.tenant_id,
             dataset_id=dataset.id,
             description=description,
             **kwargs
         )
 
-    def _run(self, dataset_id: str, query: str) -> str:
-        pattern = r'\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b'
-        match = re.search(pattern, dataset_id, re.IGNORECASE)
-        if match:
-            dataset_id = match.group()
-
+    def _run(self, query: str) -> str:
         dataset = db.session.query(Dataset).filter(
             Dataset.tenant_id == self.tenant_id,
-            Dataset.id == dataset_id
+            Dataset.id == self.dataset_id
         ).first()
 
         if not dataset:
-            return f'[{self.name} failed to find dataset with id {dataset_id}.]'
+            return f'[{self.name} failed to find dataset with id {self.dataset_id}.]'
 
         if dataset.indexing_technique == "economy":
             # use keyword table query
@@ -93,7 +92,7 @@ class DatasetRetrieverTool(BaseTool):
             if self.k > 0:
                 documents = vector_index.search(
                     query,
-                    search_type='similarity',
+                    search_type='similarity_score_threshold',
                     search_kwargs={
                         'k': self.k
                     }
@@ -101,11 +100,16 @@ class DatasetRetrieverTool(BaseTool):
             else:
                 documents = []
 
-            hit_callback = DatasetIndexToolCallbackHandler(dataset.id)
+            hit_callback = DatasetIndexToolCallbackHandler(dataset.id, self.conversation_message_task)
             hit_callback.on_tool_end(documents)
+            document_score_list = {}
+            if dataset.indexing_technique != "economy":
+                for item in documents:
+                    document_score_list[item.metadata['doc_id']] = item.metadata['score']
             document_context_list = []
             index_node_ids = [document.metadata['doc_id'] for document in documents]
-            segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+            segments = DocumentSegment.query.filter(DocumentSegment.dataset_id == self.dataset_id,
+                                                    DocumentSegment.completed_at.isnot(None),
                                                     DocumentSegment.status == 'completed',
                                                     DocumentSegment.enabled == True,
                                                     DocumentSegment.index_node_id.in_(index_node_ids)
@@ -118,9 +122,43 @@ class DatasetRetrieverTool(BaseTool):
                                                                                            float('inf')))
                 for segment in sorted_segments:
                     if segment.answer:
-                        document_context_list.append(f'question:{segment.content} \nanswer:{segment.answer}')
+                        document_context_list.append(f'question:{segment.content} answer:{segment.answer}')
                     else:
                         document_context_list.append(segment.content)
+                if self.return_resource:
+                    context_list = []
+                    resource_number = 1
+                    for segment in sorted_segments:
+                        context = {}
+                        document = Document.query.filter(Document.id == segment.document_id,
+                                                         Document.enabled == True,
+                                                         Document.archived == False,
+                                                         ).first()
+                        if dataset and document:
+                            source = {
+                                'position': resource_number,
+                                'dataset_id': dataset.id,
+                                'dataset_name': dataset.name,
+                                'document_id': document.id,
+                                'document_name': document.name,
+                                'data_source_type': document.data_source_type,
+                                'segment_id': segment.id,
+                                'retriever_from': self.retriever_from
+                            }
+                            if dataset.indexing_technique != "economy":
+                                source['score'] = document_score_list.get(segment.index_node_id)
+                            if self.retriever_from == 'dev':
+                                source['hit_count'] = segment.hit_count
+                                source['word_count'] = segment.word_count
+                                source['segment_position'] = segment.position
+                                source['index_node_hash'] = segment.index_node_hash
+                            if segment.answer:
+                                source['content'] = f'question:{segment.content} \nanswer:{segment.answer}'
+                            else:
+                                source['content'] = segment.content
+                            context_list.append(source)
+                        resource_number += 1
+                    hit_callback.return_retriever_resource_info(context_list)
 
             return str("\n".join(document_context_list))
 

api/core/tool/web_reader_tool.py

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

api/core/vector_store/milvus_vector_store.py

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

api/core/vector_store/qdrant_vector_store.py

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

api/events/event_handlers/create_document_index.py

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

api/events/event_handlers/generate_conversation_name_when_first_message_created.py

@@ -26,7 +26,7 @@ def handle(sender, **kwargs):
 
                 conversation.name = name
             except:
-                conversation.name = 'New Chat'
+                conversation.name = 'New conversation'
 
             db.session.add(conversation)
             db.session.commit()

api/migrations/versions/4bcffcd64aa4_update_dataset_model_field_null_.py

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