bump version to 0.3.26 (#1307)

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

.devcontainer/Dockerfile

@@ -1,11 +1,8 @@
-FROM mcr.microsoft.com/devcontainers/anaconda:0-3
+FROM mcr.microsoft.com/devcontainers/python:3.10
+
+COPY . .
 
-# Copy environment.yml (if found) to a temp location so we update the environment. Also
-# copy "noop.txt" so the COPY instruction does not fail if no environment.yml exists.
-COPY environment.yml* .devcontainer/noop.txt /tmp/conda-tmp/
-RUN if [ -f "/tmp/conda-tmp/environment.yml" ]; then umask 0002 && /opt/conda/bin/conda env update -n base -f /tmp/conda-tmp/environment.yml; fi \
-    && rm -rf /tmp/conda-tmp
 
 # [Optional] Uncomment this section to install additional OS packages.
 # RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
-#     && apt-get -y install --no-install-recommends <your-package-list-here>
+#     && apt-get -y install --no-install-recommends <your-package-list-here>

.devcontainer/devcontainer.json

@@ -1,13 +1,12 @@
 // For format details, see https://aka.ms/devcontainer.json. For config options, see the
 // README at: https://github.com/devcontainers/templates/tree/main/src/anaconda
 {
-	"name": "Anaconda (Python 3)",
+	"name": "Python 3.10",
 	"build": { 
 		"context": "..",
 		"dockerfile": "Dockerfile"
 	},
 	"features": {
-		"ghcr.io/dhoeric/features/act:1": {},
 		"ghcr.io/devcontainers/features/node:1": {
 			"nodeGypDependencies": true,
 			"version": "lts"

.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/help_wanted.yml

@@ -0,0 +1,11 @@
+name: "🤝 Help Wanted"
+description: "Request help from the community"
+labels:
+- help-wanted
+body:
+- type: textarea
+  attributes:
+    label: Provide a description of the help you need
+    placeholder: Briefly describe what you need help with.
+  validations:
+    required: true

.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/api-unit-tests.yml

@@ -0,0 +1,38 @@
+name: Run Pytest
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - deploy/dev
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.10'
+
+    - name: Cache pip dependencies
+      uses: actions/cache@v2
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('api/requirements.txt') }}
+        restore-keys: ${{ runner.os }}-pip-
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        pip install -r api/requirements.txt
+
+    - name: Run pytest
+      run: pytest api/tests/unit_tests

.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

@@ -144,9 +144,11 @@ docker/volumes/app/storage/*
 docker/volumes/db/data/*
 docker/volumes/redis/data/*
 docker/volumes/weaviate/*
+docker/volumes/qdrant/*
 
 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) に参加してください。私たちがお手伝いします！

README.md

@@ -16,18 +16,31 @@ Out-of-the-box web sites supporting form mode and chat conversation mode
 A single API encompassing plugin capabilities, context enhancement, and more, saving you backend coding effort
 Visual data analysis, log review, and annotation for applications
 
+
+https://github.com/langgenius/dify/assets/100913391/f6e658d5-31b3-4c16-a0af-9e191da4d0f6
+
+
 ## Highlighted Features
 **1. LLMs support:** Choose capabilities based on different models when building your Dify AI apps. Dify is compatible with Langchain, meaning it will support various LLMs. Currently supported:
 
->* OpenAI: GPT-4, GPT-3.5-turbo, GPT-3.5-turbo-16k, text-davinci-003
+- [x] **OpenAI**: GPT4, GPT3.5-turbo, GPT3.5-turbo-16k, text-davinci-003 
->* Azure OpenAI Service
+- [x] **Azure OpenAI Service**
->* Anthropic: Claude2, Claude-instant
+- [x] **Anthropic**: Claude2, Claude-instant
->* Hugging Face Hub (coming soon)
+- [x] **Replicate**
+- [x] **Hugging Face Hub**
+- [x] **ChatGLM**
+- [x] **Llama2**
+- [x] **MiniMax**
+- [x] **Spark**
+- [x] **Wenxin**
+- [x] **Tongyi**
+
 
 We provide the following free resources for registered Dify cloud users (sign up at [dify.ai](https://dify.ai)):
-* 1000 free Claude model queries to build Claude-powered apps
+* 600,000 free Claude model tokens to build Claude-powered apps
 * 200 free OpenAI queries to build OpenAI-based apps
 
+  
 **2. Visual orchestration:** Build an AI app in minutes by writing and debugging prompts visually.
 
 **3. Text embedding:** Fully automated text preprocessing embeds your data as context without complex concepts. Supports PDF, TXT, and syncing data from Notion, webpages, APIs.
@@ -55,7 +68,7 @@ Visit [Dify.ai](https://dify.ai)
 
 Before installing Dify, make sure your machine meets the following minimum system requirements:
 
-- CPU >= 1 Core
+- CPU >= 2 Core
 - RAM >= 4GB
 
 ### Quick Start
@@ -86,8 +99,6 @@ Features under development:
 We will support more datasets, including text, webpages, and even Notion content. Users can build AI applications based on their own data sources.
 - **Plugins**, introducing ChatGPT Plugin-standard plugins for applications, or using Dify-produced plugins
 We will release plugins complying with ChatGPT standard, or Dify's own plugins to enable more capabilities in applications. 
-- **Open-source models**, e.g. adopting Llama as a model provider or for further fine-tuning
-We will work with excellent open-source models like Llama, by providing them as model options in our platform, or using them for further fine-tuning.
 
 
 ## Q&A

README_CN.md

@@ -17,19 +17,29 @@
 - 一套 API 即可包含插件、上下文增强等能力，替你省下了后端代码的编写工作
 - 可视化的对应用进行数据分析，查阅日志或进行标注
 
-
+https://github.com/langgenius/dify/assets/100913391/f6e658d5-31b3-4c16-a0af-9e191da4d0f6
 
 ## 核心能力
 1. **模型支持：** 你可以在 Dify 上选择基于不同模型的能力来开发你的 AI 应用。Dify 兼容 Langchain，这意味着我们将逐步支持多种 LLMs ，目前支持的模型供应商：
 
-> * **OpenAI**：GPT4、GPT3.5-turbo、GPT3.5-turbo-16k、text-davinci-003 
+- [x] **OpenAI**：GPT4、GPT3.5-turbo、GPT3.5-turbo-16k、text-davinci-003 
-> * **Azure OpenAI Service**
+- [x] **Azure OpenAI Service**
-> * **Anthropic**：Claude2、Claude-instant
+- [x] **Anthropic**：Claude2、Claude-instant
-> * **Hugging Face Hub**（即将推出）
+- [x] **Replicate**
+- [x] **Hugging Face Hub**
+- [x] **ChatGLM**
+- [x] **Llama2**
+- [x] **MiniMax**
+- [x] **讯飞星火大模型**
+- [x] **文心一言**
+- [x] **通义千问**
+
 
 我们为所有注册云端版的用户免费提供以下资源（登录 [dify.ai](https://cloud.dify.ai) 即可使用）：
-* 1000 次 Claude 模型的消息调用额度，用于创建基于 Claude 模型的 AI 应用
+* 60 万 Tokens Claude 模型的消息调用额度，用于创建基于 Claude 模型的 AI 应用
 * 200 次 OpenAI 模型的消息调用额度，用于创建基于 OpenAI 模型的 AI 应用
+* 300 万 讯飞星火大模型 Token 的调用额度，用于创建基于讯飞星火大模型的 AI 应用
+* 100 万 MiniMax Token 的调用额度，用于创建基于 MiniMax 模型的 AI 应用
 2. **可视化编排 Prompt：** 通过界面化编写 prompt 并调试，只需几分钟即可发布一个 AI 应用。
 3. **文本 Embedding 处理（数据集）**：全自动完成文本预处理，使用你的数据作为上下文，无需理解晦涩的概念和技术处理。支持 PDF、txt 等文件格式，支持从 Notion、网页、API 同步数据。
 4. **基于 API 开发：** 后端即服务。您可以直接访问网页应用，也可以接入 API 集成到您的应用中，无需关注复杂的后端架构和部署过程。
@@ -53,7 +63,7 @@
 
 在安装 Dify 之前，请确保您的机器满足以下最低系统要求：
 
-- CPU >= 1 Core
+- CPU >= 2 Core
 - RAM >= 4GB
 
 ### 快速启动
@@ -82,8 +92,6 @@ docker compose up -d
 
 - **数据集**，支持更多的数据集，通过网页、API 同步内容。用户可以根据自己的数据源构建 AI 应用程序。
 - **插件**，我们将发布符合 ChatGPT 标准的插件，支持更多 Dify 自己的插件，支持用户自定义插件能力，以在应用程序中启用更多功能，例如以支持以目标为导向的分解推理任务。
-- **开源模型支持**，支持 Hugging face Hub 上的开源模型。例如采用 Llama 作为模型提供者，或进行进一步的微调
-我们将与优秀的开源模型合作，通过在我们的平台中提供它们作为模型选项，或使用它们进行进一步的微调。
 
 ## Q&A
 

README_ES.md

@@ -32,7 +32,7 @@ Visita [Dify.ai](https://dify.ai)
 
 Antes de instalar Dify, asegúrate de que tu máquina cumple con los siguientes requisitos mínimos del sistema:
 
-- CPU >= 1 Core
+- CPU >= 2 Core
 - RAM >= 4GB
 
 ### Inicio rápido

api/.env.example

@@ -50,25 +50,7 @@ S3_REGION=your-region
 WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 
-# Cookie configuration
+# Vector database configuration, support: weaviate, qdrant, milvus
-COOKIE_HTTPONLY=true
-COOKIE_SAMESITE=None
-COOKIE_SECURE=true
-
-# Session configuration
-SESSION_PERMANENT=true
-SESSION_USE_SIGNER=true
-
-## support redis, sqlalchemy
-SESSION_TYPE=redis
-
-# session redis configuration
-SESSION_REDIS_HOST=localhost
-SESSION_REDIS_PORT=6379
-SESSION_REDIS_PASSWORD=difyai123456
-SESSION_REDIS_DB=2
-
-# Vector database configuration, support: weaviate, qdrant
 VECTOR_STORE=weaviate
 
 # Weaviate configuration
@@ -77,9 +59,16 @@ WEAVIATE_API_KEY=WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih
 WEAVIATE_GRPC_ENABLED=false
 WEAVIATE_BATCH_SIZE=100
 
-# Qdrant configuration, use `path:` prefix for local mode or `https://your-qdrant-cluster-url.qdrant.io` for remote mode
+# Qdrant configuration, use `http://localhost:6333` for local mode or `https://your-qdrant-cluster-url.qdrant.io` for remote mode
-QDRANT_URL=path:storage/qdrant
+QDRANT_URL=http://localhost:6333
-QDRANT_API_KEY=your-qdrant-api-key
+QDRANT_API_KEY=difyai123456
+
+# Milvus configuration
+MILVUS_HOST=127.0.0.1
+MILVUS_PORT=19530
+MILVUS_USER=root
+MILVUS_PASSWORD=Milvus
+MILVUS_SECURE=false
 
 # Mail configuration, support: resend
 MAIL_TYPE=
@@ -117,10 +106,12 @@ HOSTED_AZURE_OPENAI_QUOTA_LIMIT=200
 HOSTED_ANTHROPIC_ENABLED=false
 HOSTED_ANTHROPIC_API_BASE=
 HOSTED_ANTHROPIC_API_KEY=
-HOSTED_ANTHROPIC_QUOTA_LIMIT=1000000
+HOSTED_ANTHROPIC_QUOTA_LIMIT=600000
 HOSTED_ANTHROPIC_PAID_ENABLED=false
 HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID=
-HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA=1
+HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA=1000000
+HOSTED_ANTHROPIC_PAID_MIN_QUANTITY=20
+HOSTED_ANTHROPIC_PAID_MAX_QUANTITY=100
 
 STRIPE_API_KEY=
 STRIPE_WEBHOOK_SECRET=

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 && \
+RUN apt-get update \
-    apt-get install -y bash curl wget vim gcc g++ python3-dev libc-dev libffi-dev
+    && apt-get install -y --no-install-recommends bash curl wget vim nodejs \
-
+    && apt-get autoremove \
-COPY requirements.txt /app/api/requirements.txt
+    && rm -rf /var/lib/apt/lists/*
-
-RUN pip install -r requirements.txt
 
+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,8 +1,7 @@
 # -*- coding:utf-8 -*-
 import os
-from datetime import datetime
 
-from werkzeug.exceptions import Forbidden
+from werkzeug.exceptions import Unauthorized
 
 if not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true':
     from gevent import monkey
@@ -12,12 +11,11 @@ import logging
 import json
 import threading
 
-from flask import Flask, request, Response, session
+from flask import Flask, request, Response
-import flask_login
 from flask_cors import CORS
 
 from core.model_providers.providers import hosted
-from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
+from extensions import ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
     ext_database, ext_storage, ext_mail, ext_stripe
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
@@ -27,12 +25,10 @@ from models import model, account, dataset, web, task, source, tool
 from events import event_handlers
 # DO NOT REMOVE ABOVE
 
-import core
 from config import Config, CloudEditionConfig
 from commands import register_commands
-from models.account import TenantAccountJoin, AccountStatus
+from services.account_service import AccountService
-from models.model import Account, EndUser, App
+from libs.passport import PassportService
-from services.account_service import TenantService
 
 import warnings
 warnings.simplefilter("ignore", ResourceWarning)
@@ -85,77 +81,33 @@ def initialize_extensions(app):
     ext_redis.init_app(app)
     ext_storage.init_app(app)
     ext_celery.init_app(app)
-    ext_session.init_app(app)
     ext_login.init_app(app)
     ext_mail.init_app(app)
     ext_sentry.init_app(app)
     ext_stripe.init_app(app)
 
 
-def _create_tenant_for_account(account):
-    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
-
-    TenantService.create_tenant_member(tenant, account, role='owner')
-    account.current_tenant = tenant
-
-    return tenant
-
-
 # Flask-Login configuration
-@login_manager.user_loader
+@login_manager.request_loader
-def load_user(user_id):
+def load_user_from_request(request_from_flask_login):
-    """Load user based on the user_id."""
+    """Load user based on the request."""
     if request.blueprint == 'console':
         # Check if the user_id contains a dot, indicating the old format
-        if '.' in user_id:
+        auth_header = request.headers.get('Authorization', '')
-            tenant_id, account_id = user_id.split('.')
+        if ' ' not in auth_header:
-        else:
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
-            account_id = user_id
+        auth_scheme, auth_token = auth_header.split(None, 1)
+        auth_scheme = auth_scheme.lower()
+        if auth_scheme != 'bearer':
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+        
+        decoded = PassportService().verify(auth_token)
+        user_id = decoded.get('user_id')
 
-        account = db.session.query(Account).filter(Account.id == account_id).first()
+        return AccountService.load_user(user_id)
-
-        if account:
-            if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
-                raise Forbidden('Account is banned or closed.')
-
-            workspace_id = session.get('workspace_id')
-            if workspace_id:
-                tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                    TenantAccountJoin.account_id == account.id,
-                    TenantAccountJoin.tenant_id == workspace_id
-                ).first()
-
-                if not tenant_account_join:
-                    tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                        TenantAccountJoin.account_id == account.id).first()
-
-                    if tenant_account_join:
-                        account.current_tenant_id = tenant_account_join.tenant_id
-                    else:
-                        _create_tenant_for_account(account)
-                    session['workspace_id'] = account.current_tenant_id
-                else:
-                    account.current_tenant_id = workspace_id
-            else:
-                tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                    TenantAccountJoin.account_id == account.id).first()
-                if tenant_account_join:
-                    account.current_tenant_id = tenant_account_join.tenant_id
-                else:
-                    _create_tenant_for_account(account)
-                session['workspace_id'] = account.current_tenant_id
-
-            account.last_active_at = datetime.utcnow()
-            db.session.commit()
-
-            # Log in the user with the updated user_id
-            flask_login.login_user(account, remember=True)
-
-        return account
     else:
         return None
 
-
 @login_manager.unauthorized_handler
 def unauthorized_handler():
     """Handle unauthorized requests."""
@@ -212,6 +164,7 @@ if app.config['TESTING']:
 @app.after_request
 def after_request(response):
     """Add Version headers to the response."""
+    response.set_cookie('remember_token', '', expires=0)
     response.headers.add('X-Version', app.config['CURRENT_VERSION'])
     response.headers.add('X-Env', app.config['DEPLOY_ENV'])
     return response

api/commands.py

@@ -1,26 +1,36 @@
 import datetime
+import json
 import math
 import random
 import string
+import threading
 import time
+import uuid
 
 import click
-from flask import current_app
+from tqdm import tqdm
+from flask import current_app, Flask
+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.dataset import Dataset, DatasetQuery, Document, DatasetCollectionBinding
-from models.model import Account
+from models.model import Account, AppModelConfig, App
 import secrets
 import base64
 
-from models.provider import Provider, ProviderType, ProviderQuotaType
+from models.provider import Provider, ProviderType, ProviderQuotaType, ProviderModel
 
 
 @click.command('reset-password', help='Reset the account password.')
@@ -102,6 +112,7 @@ def reset_encrypt_key_pair():
     tenant.encrypt_public_key = generate_key_pair(tenant.id)
 
     db.session.query(Provider).filter(Provider.provider_type == 'custom').delete()
+    db.session.query(ProviderModel).delete()
     db.session.commit()
 
     click.echo(click.style('Congratulations! '
@@ -230,7 +241,13 @@ def clean_unused_dataset_indexes():
                         kw_index = IndexBuilder.get_index(dataset, 'economy')
                         # delete from vector index
                         if vector_index:
-                            vector_index.delete()
+                            if dataset.collection_binding_id:
+                                vector_index.delete_by_group_id(dataset.id)
+                            else:
+                                if dataset.collection_binding_id:
+                                    vector_index.delete_by_group_id(dataset.id)
+                                else:
+                                    vector_index.delete()
                         kw_index.delete()
                         # update document
                         update_params = {
@@ -258,6 +275,8 @@ def sync_anthropic_hosted_providers():
     click.echo(click.style('Start sync anthropic hosted providers.', fg='green'))
     count = 0
 
+    new_quota_limit = hosted_model_providers.anthropic.quota_limit
+
     page = 1
     while True:
         try:
@@ -265,6 +284,7 @@ def sync_anthropic_hosted_providers():
                 Provider.provider_name == 'anthropic',
                 Provider.provider_type == ProviderType.SYSTEM.value,
                 Provider.quota_type == ProviderQuotaType.TRIAL.value,
+                Provider.quota_limit != new_quota_limit
             ).order_by(Provider.created_at.desc()).paginate(page=page, per_page=100)
         except NotFound:
             break
@@ -272,9 +292,9 @@ def sync_anthropic_hosted_providers():
         page += 1
         for provider in providers:
             try:
-                click.echo('Syncing tenant anthropic hosted provider: {}'.format(provider.tenant_id))
+                click.echo('Syncing tenant anthropic hosted provider: {}, origin: limit {}, used {}'
+                           .format(provider.tenant_id, provider.quota_limit, provider.quota_used))
                 original_quota_limit = provider.quota_limit
-                new_quota_limit = hosted_model_providers.anthropic.quota_limit
                 division = math.ceil(new_quota_limit / 1000)
 
                 provider.quota_limit = new_quota_limit if original_quota_limit == 1000 \
@@ -292,6 +312,412 @@ 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('normalization-collections', help='restore all collections in one')
+def normalization_collections():
+    click.echo(click.style('Start normalization collections.', fg='green'))
+    normalization_count = []
+    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=100)
+        except NotFound:
+            break
+        datasets_result = datasets.items
+        page += 1
+        for i in range(0, len(datasets_result), 5):
+            threads = []
+            sub_datasets = datasets_result[i:i + 5]
+            for dataset in sub_datasets:
+                document_format_thread = threading.Thread(target=deal_dataset_vector, kwargs={
+                    'flask_app': current_app._get_current_object(),
+                    'dataset': dataset,
+                    'normalization_count': normalization_count
+                })
+                threads.append(document_format_thread)
+                document_format_thread.start()
+            for thread in threads:
+                thread.join()
+
+    click.echo(click.style('Congratulations! restore {} dataset indexes.'.format(len(normalization_count)), fg='green'))
+
+
+def deal_dataset_vector(flask_app: Flask, dataset: Dataset, normalization_count: list):
+    with flask_app.app_context():
+        try:
+            click.echo('restore dataset 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)
+            dataset_collection_binding = db.session.query(DatasetCollectionBinding). \
+                filter(DatasetCollectionBinding.provider_name == embedding_model.model_provider.provider_name,
+                       DatasetCollectionBinding.model_name == embedding_model.name). \
+                order_by(DatasetCollectionBinding.created_at). \
+                first()
+
+            if not dataset_collection_binding:
+                dataset_collection_binding = DatasetCollectionBinding(
+                    provider_name=embedding_model.model_provider.provider_name,
+                    model_name=embedding_model.name,
+                    collection_name="Vector_index_" + str(uuid.uuid4()).replace("-", "_") + '_Node'
+                )
+                db.session.add(dataset_collection_binding)
+                db.session.commit()
+
+            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.delete_by_group_id(dataset.id)
+                index.restore_dataset_in_one(dataset, dataset_collection_binding)
+            else:
+                click.echo('passed.')
+            normalization_count.append(1)
+        except Exception as e:
+            click.echo(
+                click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)),
+                            fg='red'))
+
+
+@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))
+
+@click.command('migrate_default_input_to_dataset_query_variable')
+@click.option("--batch-size", default=500, help="Number of records to migrate in each batch.")
+def migrate_default_input_to_dataset_query_variable(batch_size):
+
+    click.secho("Starting...", fg='green')
+
+    total_records = db.session.query(AppModelConfig) \
+        .join(App, App.app_model_config_id == AppModelConfig.id) \
+        .filter(App.mode == 'completion') \
+        .filter(AppModelConfig.dataset_query_variable == None) \
+        .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') \
+                .filter(AppModelConfig.dataset_query_variable == None) \
+                .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:
+                    config = AppModelConfig.to_dict(data)
+
+                    tools = config["agent_mode"]["tools"]
+                    dataset_exists = "dataset" in str(tools)
+                    if not dataset_exists:
+                        continue
+
+                    user_input_form = config.get("user_input_form", [])
+                    for form in user_input_form:
+                        paragraph = form.get('paragraph')
+                        if paragraph \
+                            and paragraph.get('variable') == 'query':
+                                data.dataset_query_variable = 'query'
+                                break
+                        
+                        if paragraph \
+                            and paragraph.get('variable') == 'default_input':
+                                data.dataset_query_variable = 'default_input'
+                                break
+
+                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)
@@ -300,3 +726,8 @@ 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)
+    app.cli.add_command(normalization_collections)
+    app.cli.add_command(migrate_default_input_to_dataset_query_variable)

api/config.py

@@ -10,9 +10,6 @@ from extensions.ext_redis import redis_client
 dotenv.load_dotenv()
 
 DEFAULTS = {
-    'COOKIE_HTTPONLY': 'True',
-    'COOKIE_SECURE': 'True',
-    'COOKIE_SAMESITE': 'None',
     'DB_USERNAME': 'postgres',
     'DB_PASSWORD': '',
     'DB_HOST': 'localhost',
@@ -22,10 +19,6 @@ DEFAULTS = {
     'REDIS_PORT': '6379',
     'REDIS_DB': '0',
     'REDIS_USE_SSL': 'False',
-    'SESSION_REDIS_HOST': 'localhost',
-    'SESSION_REDIS_PORT': '6379',
-    'SESSION_REDIS_DB': '2',
-    'SESSION_REDIS_USE_SSL': 'False',
     'OAUTH_REDIRECT_PATH': '/console/api/oauth/authorize',
     'OAUTH_REDIRECT_INDEX_PATH': '/',
     'CONSOLE_WEB_URL': 'https://cloud.dify.ai',
@@ -36,9 +29,6 @@ DEFAULTS = {
     'STORAGE_TYPE': 'local',
     'STORAGE_LOCAL_PATH': 'storage',
     'CHECK_UPDATE_URL': 'https://updates.dify.ai',
-    'SESSION_TYPE': 'sqlalchemy',
-    'SESSION_PERMANENT': 'True',
-    'SESSION_USE_SIGNER': 'True',
     'DEPLOY_ENV': 'PRODUCTION',
     'SQLALCHEMY_POOL_SIZE': 30,
     'SQLALCHEMY_POOL_RECYCLE': 3600,
@@ -48,21 +38,25 @@ DEFAULTS = {
     'WEAVIATE_GRPC_ENABLED': 'True',
     'WEAVIATE_BATCH_SIZE': 100,
     'CELERY_BACKEND': 'database',
-    'PDF_PREVIEW': 'True',
     'LOG_LEVEL': 'INFO',
-    'DISABLE_PROVIDER_CONFIG_VALIDATION': 'False',
     'HOSTED_OPENAI_QUOTA_LIMIT': 200,
     'HOSTED_OPENAI_ENABLED': 'False',
     'HOSTED_OPENAI_PAID_ENABLED': 'False',
     'HOSTED_OPENAI_PAID_INCREASE_QUOTA': 1,
     'HOSTED_AZURE_OPENAI_ENABLED': 'False',
     'HOSTED_AZURE_OPENAI_QUOTA_LIMIT': 200,
-    'HOSTED_ANTHROPIC_QUOTA_LIMIT': 1000000,
+    'HOSTED_ANTHROPIC_QUOTA_LIMIT': 600000,
     'HOSTED_ANTHROPIC_ENABLED': 'False',
     'HOSTED_ANTHROPIC_PAID_ENABLED': 'False',
-    'HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA': 1,
+    'HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA': 1000000,
+    'HOSTED_ANTHROPIC_PAID_MIN_QUANTITY': 20,
+    'HOSTED_ANTHROPIC_PAID_MAX_QUANTITY': 100,
+    'HOSTED_MODERATION_ENABLED': 'False',
+    'HOSTED_MODERATION_PROVIDERS': '',
     'TENANT_DOCUMENT_COUNT': 100,
-    'CLEAN_DAY_SETTING': 30
+    'CLEAN_DAY_SETTING': 30,
+    'UPLOAD_FILE_SIZE_LIMIT': 15,
+    'UPLOAD_FILE_BATCH_LIMIT': 5,
 }
 
 
@@ -98,13 +92,12 @@ 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.13"
+        self.CURRENT_VERSION = "0.3.26"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = "SELF_HOSTED"
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')
         self.TESTING = False
         self.LOG_LEVEL = get_env('LOG_LEVEL')
-        self.PDF_PREVIEW = get_bool_env('PDF_PREVIEW')
 
         # Your App secret key will be used for securely signing the session cookie
         # Make sure you are changing this key for your deployment with a strong key.
@@ -112,20 +105,6 @@ class Config:
         # Alternatively you can set it with `SECRET_KEY` environment variable.
         self.SECRET_KEY = get_env('SECRET_KEY')
 
-        # cookie settings
-        self.REMEMBER_COOKIE_HTTPONLY = get_bool_env('COOKIE_HTTPONLY')
-        self.SESSION_COOKIE_HTTPONLY = get_bool_env('COOKIE_HTTPONLY')
-        self.REMEMBER_COOKIE_SAMESITE = get_env('COOKIE_SAMESITE')
-        self.SESSION_COOKIE_SAMESITE = get_env('COOKIE_SAMESITE')
-        self.REMEMBER_COOKIE_SECURE = get_bool_env('COOKIE_SECURE')
-        self.SESSION_COOKIE_SECURE = get_bool_env('COOKIE_SECURE')
-        self.PERMANENT_SESSION_LIFETIME = timedelta(days=7)
-
-        # session settings, only support sqlalchemy, redis
-        self.SESSION_TYPE = get_env('SESSION_TYPE')
-        self.SESSION_PERMANENT = get_bool_env('SESSION_PERMANENT')
-        self.SESSION_USE_SIGNER = get_bool_env('SESSION_USE_SIGNER')
-
         # redis settings
         self.REDIS_HOST = get_env('REDIS_HOST')
         self.REDIS_PORT = get_env('REDIS_PORT')
@@ -134,14 +113,6 @@ class Config:
         self.REDIS_DB = get_env('REDIS_DB')
         self.REDIS_USE_SSL = get_bool_env('REDIS_USE_SSL')
 
-        # session redis settings
-        self.SESSION_REDIS_HOST = get_env('SESSION_REDIS_HOST')
-        self.SESSION_REDIS_PORT = get_env('SESSION_REDIS_PORT')
-        self.SESSION_REDIS_USERNAME = get_env('SESSION_REDIS_USERNAME')
-        self.SESSION_REDIS_PASSWORD = get_env('SESSION_REDIS_PASSWORD')
-        self.SESSION_REDIS_DB = get_env('SESSION_REDIS_DB')
-        self.SESSION_REDIS_USE_SSL = get_bool_env('SESSION_REDIS_USE_SSL')
-
         # storage settings
         self.STORAGE_TYPE = get_env('STORAGE_TYPE')
         self.STORAGE_LOCAL_PATH = get_env('STORAGE_LOCAL_PATH')
@@ -164,6 +135,14 @@ class Config:
         self.QDRANT_URL = get_env('QDRANT_URL')
         self.QDRANT_API_KEY = get_env('QDRANT_API_KEY')
 
+        # milvus setting
+        self.MILVUS_HOST = get_env('MILVUS_HOST')
+        self.MILVUS_PORT = get_env('MILVUS_PORT')
+        self.MILVUS_USER = get_env('MILVUS_USER')
+        self.MILVUS_PASSWORD = get_env('MILVUS_PASSWORD')
+        self.MILVUS_SECURE = get_env('MILVUS_SECURE')
+
+
         # cors settings
         self.CONSOLE_CORS_ALLOW_ORIGINS = get_cors_allow_origins(
             'CONSOLE_CORS_ALLOW_ORIGINS', self.CONSOLE_WEB_URL)
@@ -209,7 +188,7 @@ class Config:
         self.HOSTED_OPENAI_API_KEY = get_env('HOSTED_OPENAI_API_KEY')
         self.HOSTED_OPENAI_API_BASE = get_env('HOSTED_OPENAI_API_BASE')
         self.HOSTED_OPENAI_API_ORGANIZATION = get_env('HOSTED_OPENAI_API_ORGANIZATION')
-        self.HOSTED_OPENAI_QUOTA_LIMIT = get_env('HOSTED_OPENAI_QUOTA_LIMIT')
+        self.HOSTED_OPENAI_QUOTA_LIMIT = int(get_env('HOSTED_OPENAI_QUOTA_LIMIT'))
         self.HOSTED_OPENAI_PAID_ENABLED = get_bool_env('HOSTED_OPENAI_PAID_ENABLED')
         self.HOSTED_OPENAI_PAID_STRIPE_PRICE_ID = get_env('HOSTED_OPENAI_PAID_STRIPE_PRICE_ID')
         self.HOSTED_OPENAI_PAID_INCREASE_QUOTA = int(get_env('HOSTED_OPENAI_PAID_INCREASE_QUOTA'))
@@ -217,23 +196,24 @@ class Config:
         self.HOSTED_AZURE_OPENAI_ENABLED = get_bool_env('HOSTED_AZURE_OPENAI_ENABLED')
         self.HOSTED_AZURE_OPENAI_API_KEY = get_env('HOSTED_AZURE_OPENAI_API_KEY')
         self.HOSTED_AZURE_OPENAI_API_BASE = get_env('HOSTED_AZURE_OPENAI_API_BASE')
-        self.HOSTED_AZURE_OPENAI_QUOTA_LIMIT = get_env('HOSTED_AZURE_OPENAI_QUOTA_LIMIT')
+        self.HOSTED_AZURE_OPENAI_QUOTA_LIMIT = int(get_env('HOSTED_AZURE_OPENAI_QUOTA_LIMIT'))
 
         self.HOSTED_ANTHROPIC_ENABLED = get_bool_env('HOSTED_ANTHROPIC_ENABLED')
         self.HOSTED_ANTHROPIC_API_BASE = get_env('HOSTED_ANTHROPIC_API_BASE')
         self.HOSTED_ANTHROPIC_API_KEY = get_env('HOSTED_ANTHROPIC_API_KEY')
-        self.HOSTED_ANTHROPIC_QUOTA_LIMIT = get_env('HOSTED_ANTHROPIC_QUOTA_LIMIT')
+        self.HOSTED_ANTHROPIC_QUOTA_LIMIT = int(get_env('HOSTED_ANTHROPIC_QUOTA_LIMIT'))
         self.HOSTED_ANTHROPIC_PAID_ENABLED = get_bool_env('HOSTED_ANTHROPIC_PAID_ENABLED')
         self.HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID = get_env('HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID')
-        self.HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA = get_env('HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA')
+        self.HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA = int(get_env('HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA'))
+        self.HOSTED_ANTHROPIC_PAID_MIN_QUANTITY = int(get_env('HOSTED_ANTHROPIC_PAID_MIN_QUANTITY'))
+        self.HOSTED_ANTHROPIC_PAID_MAX_QUANTITY = int(get_env('HOSTED_ANTHROPIC_PAID_MAX_QUANTITY'))
+
+        self.HOSTED_MODERATION_ENABLED = get_bool_env('HOSTED_MODERATION_ENABLED')
+        self.HOSTED_MODERATION_PROVIDERS = get_env('HOSTED_MODERATION_PROVIDERS')
 
         self.STRIPE_API_KEY = get_env('STRIPE_API_KEY')
         self.STRIPE_WEBHOOK_SECRET = get_env('STRIPE_WEBHOOK_SECRET')
 
-        # By default it is False
-        # You could disable it for compatibility with certain OpenAPI providers
-        self.DISABLE_PROVIDER_CONFIG_VALIDATION = get_bool_env('DISABLE_PROVIDER_CONFIG_VALIDATION')
-
         # notion import setting
         self.NOTION_CLIENT_ID = get_env('NOTION_CLIENT_ID')
         self.NOTION_CLIENT_SECRET = get_env('NOTION_CLIENT_SECRET')
@@ -244,6 +224,10 @@ class Config:
         self.TENANT_DOCUMENT_COUNT = get_env('TENANT_DOCUMENT_COUNT')
         self.CLEAN_DAY_SETTING = get_env('CLEAN_DAY_SETTING')
 
+        # uploading settings
+        self.UPLOAD_FILE_SIZE_LIMIT = int(get_env('UPLOAD_FILE_SIZE_LIMIT'))
+        self.UPLOAD_FILE_BATCH_LIMIT = int(get_env('UPLOAD_FILE_BATCH_LIMIT'))
+
 
 class CloudEditionConfig(Config):
 

api/constants/model_template.py

@@ -16,7 +16,7 @@ model_templates = {
         },
         'model_config': {
             'provider': 'openai',
-            'model_id': 'text-davinci-003',
+            'model_id': 'gpt-3.5-turbo-instruct',
             'configs': {
                 'prompt_template': '',
                 'prompt_variables': [],
@@ -30,7 +30,7 @@ model_templates = {
             },
             'model': json.dumps({
                 "provider": "openai",
-                "name": "text-davinci-003",
+                "name": "gpt-3.5-turbo-instruct",
                 "completion_params": {
                     "max_tokens": 512,
                     "temperature": 1,
@@ -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}}'
         }
     },
 
@@ -93,7 +104,7 @@ demo_model_templates = {
             'mode': 'completion',
             'model_config': AppModelConfig(
                 provider='openai',
-                model_id='text-davinci-003',
+                model_id='gpt-3.5-turbo-instruct',
                 configs={
                     'prompt_template': "Please translate the following text into {{target_language}}:\n",
                     'prompt_variables': [
@@ -129,7 +140,7 @@ demo_model_templates = {
                 pre_prompt="Please translate the following text into {{target_language}}:\n",
                 model=json.dumps({
                     "provider": "openai",
-                    "name": "text-davinci-003",
+                    "name": "gpt-3.5-turbo-instruct",
                     "completion_params": {
                         "max_tokens": 1000,
                         "temperature": 0,
@@ -211,7 +222,7 @@ demo_model_templates = {
             'mode': 'completion',
             'model_config': AppModelConfig(
                 provider='openai',
-                model_id='text-davinci-003',
+                model_id='gpt-3.5-turbo-instruct',
                 configs={
                     'prompt_template': "请将以下文本翻译为{{target_language}}:\n",
                     'prompt_variables': [
@@ -247,7 +258,7 @@ demo_model_templates = {
                 pre_prompt="请将以下文本翻译为{{target_language}}:\n",
                 model=json.dumps({
                     "provider": "openai",
-                    "name": "text-davinci-003",
+                    "name": "gpt-3.5-turbo-instruct",
                     "completion_params": {
                         "max_tokens": 1000,
                         "temperature": 0,

api/controllers/console/apikey.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from libs.login import login_required
 import flask_restful
 from flask_restful import Resource, fields, marshal_with
 from werkzeug.exceptions import Forbidden
@@ -80,6 +81,7 @@ class BaseApiKeyListResource(Resource):
         key = ApiToken.generate_api_key(self.token_prefix, 24)
         api_token = ApiToken()
         setattr(api_token, self.resource_id_field, resource_id)
+        api_token.tenant_id = current_user.current_tenant_id
         api_token.token = key
         api_token.type = self.resource_type
         db.session.add(api_token)

api/controllers/console/app/app.py

@@ -1,9 +1,11 @@
 # -*- coding:utf-8 -*-
 import json
+import logging
 from datetime import datetime
 
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal_with, abort, inputs
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with, abort, inputs
 from werkzeug.exceptions import Forbidden
 
 from constants.model_template import model_templates, demo_model_templates
@@ -11,42 +13,16 @@ from controllers.console import api
 from controllers.console.app.error import AppNotFoundError, ProviderNotInitializeError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
 from core.model_providers.model_factory import ModelFactory
-from core.model_providers.models.entity.model_params import ModelType
+from core.model_providers.model_provider_factory import ModelProviderFactory
 from events.app_event import app_was_created, app_was_deleted
-from libs.helper import TimestampField
+from fields.app_fields import app_pagination_fields, app_detail_fields, template_list_fields, \
+    app_detail_fields_with_site
 from extensions.ext_database import db
 from models.model import App, AppModelConfig, Site
 from services.app_model_config_service import AppModelConfigService
 
-model_config_fields = {
-    'opening_statement': fields.String,
-    'suggested_questions': fields.Raw(attribute='suggested_questions_list'),
-    'suggested_questions_after_answer': fields.Raw(attribute='suggested_questions_after_answer_dict'),
-    'speech_to_text': fields.Raw(attribute='speech_to_text_dict'),
-    'more_like_this': fields.Raw(attribute='more_like_this_dict'),
-    'sensitive_word_avoidance': fields.Raw(attribute='sensitive_word_avoidance_dict'),
-    'model': fields.Raw(attribute='model_dict'),
-    'user_input_form': fields.Raw(attribute='user_input_form_list'),
-    'pre_prompt': fields.String,
-    'agent_mode': fields.Raw(attribute='agent_mode_dict'),
-}
-
-app_detail_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'mode': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String,
-    'enable_site': fields.Boolean,
-    'enable_api': fields.Boolean,
-    'api_rpm': fields.Integer,
-    'api_rph': fields.Integer,
-    'is_demo': fields.Boolean,
-    'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
-    'created_at': TimestampField
-}
-
 
 def _get_app(app_id, tenant_id):
     app = db.session.query(App).filter(App.id == app_id, App.tenant_id == tenant_id).first()
@@ -56,35 +32,6 @@ def _get_app(app_id, tenant_id):
 
 
 class AppListApi(Resource):
-    prompt_config_fields = {
-        'prompt_template': fields.String,
-    }
-
-    model_config_partial_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    app_partial_fields = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'enable_site': fields.Boolean,
-        'enable_api': fields.Boolean,
-        'is_demo': fields.Boolean,
-        'model_config': fields.Nested(model_config_partial_fields, attribute='app_model_config'),
-        'created_at': TimestampField
-    }
-
-    app_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(app_partial_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
@@ -124,12 +71,40 @@ class AppListApi(Resource):
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
 
+        try:
+            default_model = ModelFactory.get_text_generation_model(
+                tenant_id=current_user.current_tenant_id
+            )
+        except (ProviderTokenNotInitError, LLMBadRequestError):
+            default_model = None
+        except Exception as e:
+            logging.exception(e)
+            default_model = None
+
         if args['model_config'] is not None:
             # validate config
+            model_config_dict = args['model_config']
+
+            # get model provider
+            model_provider = ModelProviderFactory.get_preferred_model_provider(
+                current_user.current_tenant_id,
+                model_config_dict["model"]["provider"]
+            )
+
+            if not model_provider:
+                if not default_model:
+                    raise ProviderNotInitializeError(
+                        f"No Default System Reasoning Model available. Please configure "
+                        f"in the Settings -> Model Provider.")
+                else:
+                    model_config_dict["model"]["provider"] = default_model.model_provider.provider_name
+                    model_config_dict["model"]["name"] = default_model.name
+
             model_configuration = AppModelConfigService.validate_configuration(
                 tenant_id=current_user.current_tenant_id,
                 account=current_user,
-                config=args['model_config']
+                config=model_config_dict,
+                mode=args['mode']
             )
 
             app = App(
@@ -141,21 +116,8 @@ class AppListApi(Resource):
                 status='normal'
             )
 
-            app_model_config = AppModelConfig(
+            app_model_config = AppModelConfig()
-                provider="",
+            app_model_config = app_model_config.from_model_config_dict(model_configuration)
-                model_id="",
-                configs={},
-                opening_statement=model_configuration['opening_statement'],
-                suggested_questions=json.dumps(model_configuration['suggested_questions']),
-                suggested_questions_after_answer=json.dumps(model_configuration['suggested_questions_after_answer']),
-                speech_to_text=json.dumps(model_configuration['speech_to_text']),
-                more_like_this=json.dumps(model_configuration['more_like_this']),
-                sensitive_word_avoidance=json.dumps(model_configuration['sensitive_word_avoidance']),
-                model=json.dumps(model_configuration['model']),
-                user_input_form=json.dumps(model_configuration['user_input_form']),
-                pre_prompt=model_configuration['pre_prompt'],
-                agent_mode=json.dumps(model_configuration['agent_mode']),
-            )
         else:
             if 'mode' not in args or args['mode'] is None:
                 abort(400, message="mode is required")
@@ -165,20 +127,22 @@ class AppListApi(Resource):
             app = App(**model_config_template['app'])
             app_model_config = AppModelConfig(**model_config_template['model_config'])
 
-            default_model = ModelFactory.get_default_model(
+            # get model provider
-                tenant_id=current_user.current_tenant_id,
+            model_provider = ModelProviderFactory.get_preferred_model_provider(
-                model_type=ModelType.TEXT_GENERATION
+                current_user.current_tenant_id,
+                app_model_config.model_dict["provider"]
             )
 
-            if default_model:
+            if not model_provider:
-                model_dict = app_model_config.model_dict
+                if not default_model:
-                model_dict['provider'] = default_model.provider_name
+                    raise ProviderNotInitializeError(
-                model_dict['name'] = default_model.model_name
+                        f"No Default System Reasoning Model available. Please configure "
-                app_model_config.model = json.dumps(model_dict)
+                        f"in the Settings -> Model Provider.")
-            else:
+                else:
-                raise ProviderNotInitializeError(
+                    model_dict = app_model_config.model_dict
-                    f"No Text Generation Model available. Please configure a valid provider "
+                    model_dict['provider'] = default_model.model_provider.provider_name
-                    f"in the Settings -> Model Provider.")
+                    model_dict['name'] = default_model.name
+                    app_model_config.model = json.dumps(model_dict)
 
         app.name = args['name']
         app.mode = args['mode']
@@ -214,18 +178,6 @@ class AppListApi(Resource):
 
 
 class AppTemplateApi(Resource):
-    template_fields = {
-        'name': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'description': fields.String,
-        'mode': fields.String,
-        'model_config': fields.Nested(model_config_fields),
-    }
-
-    template_list_fields = {
-        'data': fields.List(fields.Nested(template_fields)),
-    }
 
     @setup_required
     @login_required
@@ -244,38 +196,6 @@ class AppTemplateApi(Resource):
 
 
 class AppApi(Resource):
-    site_fields = {
-        'access_token': fields.String(attribute='code'),
-        'code': fields.String,
-        'title': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'description': fields.String,
-        'default_language': fields.String,
-        'customize_domain': fields.String,
-        'copyright': fields.String,
-        'privacy_policy': fields.String,
-        'customize_token_strategy': fields.String,
-        'prompt_public': fields.Boolean,
-        'app_base_url': fields.String,
-    }
-
-    app_detail_fields_with_site = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'enable_site': fields.Boolean,
-        'enable_api': fields.Boolean,
-        'api_rpm': fields.Integer,
-        'api_rph': fields.Integer,
-        'is_demo': fields.Boolean,
-        'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
-        'site': fields.Nested(site_fields),
-        'api_base_url': fields.String,
-        'created_at': TimestampField
-    }
 
     @setup_required
     @login_required
@@ -297,7 +217,7 @@ class AppApi(Resource):
 
         if current_user.current_tenant.current_role not in ['admin', 'owner']:
             raise Forbidden()
-        
+
         app = _get_app(app_id, current_user.current_tenant_id)
 
         db.session.delete(app)
@@ -397,29 +317,6 @@ class AppApiStatus(Resource):
         return app
 
 
-class AppRateLimit(Resource):
-    @setup_required
-    @login_required
-    @account_initialization_required
-    @marshal_with(app_detail_fields)
-    def post(self, app_id):
-        parser = reqparse.RequestParser()
-        parser.add_argument('api_rpm', type=inputs.natural, required=False, location='json')
-        parser.add_argument('api_rph', type=inputs.natural, required=False, location='json')
-        args = parser.parse_args()
-
-        app_id = str(app_id)
-        app = _get_app(app_id, current_user.current_tenant_id)
-
-        if args.get('api_rpm'):
-            app.api_rpm = args.get('api_rpm')
-        if args.get('api_rph'):
-            app.api_rph = args.get('api_rph')
-        app.updated_at = datetime.utcnow()
-        db.session.commit()
-        return app
-
-
 class AppCopy(Resource):
     @staticmethod
     def create_app_copy(app):
@@ -439,22 +336,9 @@ class AppCopy(Resource):
 
     @staticmethod
     def create_app_model_config_copy(app_config, copy_app_id):
-        copy_app_model_config = AppModelConfig(
+        copy_app_model_config = app_config.copy()
-            app_id=copy_app_id,
+        copy_app_model_config.app_id = copy_app_id
-            provider=app_config.provider,
+
-            model_id=app_config.model_id,
-            configs=app_config.configs,
-            opening_statement=app_config.opening_statement,
-            suggested_questions=app_config.suggested_questions,
-            suggested_questions_after_answer=app_config.suggested_questions_after_answer,
-            speech_to_text=app_config.speech_to_text,
-            more_like_this=app_config.more_like_this,
-            sensitive_word_avoidance=app_config.sensitive_word_avoidance,
-            model=app_config.model,
-            user_input_form=app_config.user_input_form,
-            pre_prompt=app_config.pre_prompt,
-            agent_mode=app_config.agent_mode
-        )
         return copy_app_model_config
 
     @setup_required
@@ -482,16 +366,6 @@ class AppCopy(Resource):
         return copy_app, 201
 
 
-class AppExport(Resource):
-
-    @setup_required
-    @login_required
-    @account_initialization_required
-    def post(self, app_id):
-        # todo
-        pass
-
-
 api.add_resource(AppListApi, '/apps')
 api.add_resource(AppTemplateApi, '/app-templates')
 api.add_resource(AppApi, '/apps/<uuid:app_id>')
@@ -500,4 +374,3 @@ api.add_resource(AppNameApi, '/apps/<uuid:app_id>/name')
 api.add_resource(AppIconApi, '/apps/<uuid:app_id>/icon')
 api.add_resource(AppSiteStatus, '/apps/<uuid:app_id>/site-enable')
 api.add_resource(AppApiStatus, '/apps/<uuid:app_id>/api-enable')
-api.add_resource(AppRateLimit, '/apps/<uuid:app_id>/rate-limit')

api/controllers/console/app/audio.py

@@ -2,8 +2,8 @@
 import logging
 
 from flask import request
-from flask_login import login_required
+from libs.login import login_required
-from werkzeug.exceptions import InternalServerError, NotFound
+from werkzeug.exceptions import InternalServerError
 
 import services
 from controllers.console import api

api/controllers/console/app/completion.py

@@ -5,7 +5,7 @@ from typing import Generator, Union
 
 import flask_login
 from flask import Response, stream_with_context
-from flask_login import login_required
+from libs.login import login_required
 from werkzeug.exceptions import InternalServerError, NotFound
 
 import services
@@ -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/app/conversation.py

@@ -1,8 +1,9 @@
 from datetime import datetime
 
 import pytz
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal_with
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with
 from flask_restful.inputs import int_range
 from sqlalchemy import or_, func
 from sqlalchemy.orm import joinedload
@@ -12,107 +13,14 @@ from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from libs.helper import TimestampField, datetime_string, uuid_value
+from fields.conversation_fields import conversation_pagination_fields, conversation_detail_fields, \
+    conversation_message_detail_fields, conversation_with_summary_pagination_fields
+from libs.helper import datetime_string
 from extensions.ext_database import db
 from models.model import Message, MessageAnnotation, Conversation
 
-account_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'email': fields.String
-}
-
-feedback_fields = {
-    'rating': fields.String,
-    'content': fields.String,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account': fields.Nested(account_fields, allow_null=True),
-}
-
-annotation_fields = {
-    'content': fields.String,
-    'account': fields.Nested(account_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-message_detail_fields = {
-    'id': fields.String,
-    'conversation_id': fields.String,
-    'inputs': fields.Raw,
-    'query': fields.String,
-    'message': fields.Raw,
-    'message_tokens': fields.Integer,
-    'answer': fields.String,
-    'answer_tokens': fields.Integer,
-    'provider_response_latency': fields.Float,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account_id': fields.String,
-    'feedbacks': fields.List(fields.Nested(feedback_fields)),
-    'annotation': fields.Nested(annotation_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-feedback_stat_fields = {
-    'like': fields.Integer,
-    'dislike': fields.Integer
-}
-
-model_config_fields = {
-    'opening_statement': fields.String,
-    'suggested_questions': fields.Raw,
-    'model': fields.Raw,
-    'user_input_form': fields.Raw,
-    'pre_prompt': fields.String,
-    'agent_mode': fields.Raw,
-}
-
 
 class CompletionConversationApi(Resource):
-    class MessageTextField(fields.Raw):
-        def format(self, value):
-            return value[0]['text'] if value else ''
-
-    simple_configs_fields = {
-        'prompt_template': fields.String,
-    }
-
-    simple_model_config_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    simple_message_detail_fields = {
-        'inputs': fields.Raw,
-        'query': fields.String,
-        'message': MessageTextField,
-        'answer': fields.String,
-    }
-
-    conversation_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_end_user_session_id': fields.String(),
-        'from_account_id': fields.String,
-        'read_at': TimestampField,
-        'created_at': TimestampField,
-        'annotation': fields.Nested(annotation_fields, allow_null=True),
-        'model_config': fields.Nested(simple_model_config_fields),
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields),
-        'message': fields.Nested(simple_message_detail_fields, attribute='first_message')
-    }
-
-    conversation_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(conversation_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
@@ -190,21 +98,11 @@ class CompletionConversationApi(Resource):
 
 
 class CompletionConversationDetailApi(Resource):
-    conversation_detail_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_account_id': fields.String,
-        'created_at': TimestampField,
-        'model_config': fields.Nested(model_config_fields),
-        'message': fields.Nested(message_detail_fields, attribute='first_message'),
-    }
 
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(conversation_detail_fields)
+    @marshal_with(conversation_message_detail_fields)
     def get(self, app_id, conversation_id):
         app_id = str(app_id)
         conversation_id = str(conversation_id)
@@ -233,44 +131,11 @@ class CompletionConversationDetailApi(Resource):
 
 
 class ChatConversationApi(Resource):
-    simple_configs_fields = {
-        'prompt_template': fields.String,
-    }
-
-    simple_model_config_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    conversation_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_end_user_session_id': fields.String,
-        'from_account_id': fields.String,
-        'summary': fields.String(attribute='summary_or_query'),
-        'read_at': TimestampField,
-        'created_at': TimestampField,
-        'annotated': fields.Boolean,
-        'model_config': fields.Nested(simple_model_config_fields),
-        'message_count': fields.Integer,
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields)
-    }
-
-    conversation_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(conversation_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(conversation_pagination_fields)
+    @marshal_with(conversation_with_summary_pagination_fields)
     def get(self, app_id):
         app_id = str(app_id)
 
@@ -355,19 +220,6 @@ class ChatConversationApi(Resource):
 
 
 class ChatConversationDetailApi(Resource):
-    conversation_detail_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_account_id': fields.String,
-        'created_at': TimestampField,
-        'annotated': fields.Boolean,
-        'model_config': fields.Nested(model_config_fields),
-        'message_count': fields.Integer,
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields)
-    }
 
     @setup_required
     @login_required

api/controllers/console/app/generator.py

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

api/controllers/console/app/message.py

@@ -3,7 +3,7 @@ import logging
 from typing import Union, Generator
 
 from flask import Response, stream_with_context
-from flask_login import current_user, login_required
+from flask_login import current_user
 from flask_restful import Resource, reqparse, marshal_with, fields
 from flask_restful.inputs import int_range
 from werkzeug.exceptions import InternalServerError, NotFound
@@ -16,7 +16,9 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
-from libs.helper import uuid_value, TimestampField
+from libs.login import login_required
+from fields.conversation_fields import message_detail_fields
+from libs.helper import uuid_value
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from extensions.ext_database import db
 from models.model import MessageAnnotation, Conversation, Message, MessageFeedback
@@ -26,44 +28,6 @@ from services.errors.conversation import ConversationNotExistsError
 from services.errors.message import MessageNotExistsError
 from services.message_service import MessageService
 
-account_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'email': fields.String
-}
-
-feedback_fields = {
-    'rating': fields.String,
-    'content': fields.String,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account': fields.Nested(account_fields, allow_null=True),
-}
-
-annotation_fields = {
-    'content': fields.String,
-    'account': fields.Nested(account_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-message_detail_fields = {
-    'id': fields.String,
-    'conversation_id': fields.String,
-    'inputs': fields.Raw,
-    'query': fields.String,
-    'message': fields.Raw,
-    'message_tokens': fields.Integer,
-    'answer': fields.String,
-    'answer_tokens': fields.Integer,
-    'provider_response_latency': fields.Float,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account_id': fields.String,
-    'feedbacks': fields.List(fields.Nested(feedback_fields)),
-    'annotation': fields.Nested(annotation_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
 
 class ChatMessageListApi(Resource):
     message_infinite_scroll_pagination_fields = {

api/controllers/console/app/model_config.py

@@ -1,14 +1,14 @@
 # -*- coding:utf-8 -*-
-import json
 
 from flask import request
 from flask_restful import Resource
-from flask_login import login_required, current_user
+from flask_login import current_user
 
 from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from libs.login import login_required
 from events.app_event import app_model_config_was_updated
 from extensions.ext_database import db
 from models.model import AppModelConfig
@@ -30,25 +30,14 @@ class ModelConfigResource(Resource):
         model_configuration = AppModelConfigService.validate_configuration(
             tenant_id=current_user.current_tenant_id,
             account=current_user,
-            config=request.json
+            config=request.json,
+            mode=app_model.mode
         )
 
         new_app_model_config = AppModelConfig(
             app_id=app_model.id,
-            provider="",
-            model_id="",
-            configs={},
-            opening_statement=model_configuration['opening_statement'],
-            suggested_questions=json.dumps(model_configuration['suggested_questions']),
-            suggested_questions_after_answer=json.dumps(model_configuration['suggested_questions_after_answer']),
-            speech_to_text=json.dumps(model_configuration['speech_to_text']),
-            more_like_this=json.dumps(model_configuration['more_like_this']),
-            sensitive_word_avoidance=json.dumps(model_configuration['sensitive_word_avoidance']),
-            model=json.dumps(model_configuration['model']),
-            user_input_form=json.dumps(model_configuration['user_input_form']),
-            pre_prompt=model_configuration['pre_prompt'],
-            agent_mode=json.dumps(model_configuration['agent_mode']),
         )
+        new_app_model_config = new_app_model_config.from_model_config_dict(model_configuration)
 
         db.session.add(new_app_model_config)
         db.session.flush()

api/controllers/console/app/site.py

@@ -1,32 +1,18 @@
 # -*- coding:utf-8 -*-
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal_with
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 
 from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from fields.app_fields import app_site_fields
 from libs.helper import supported_language
 from extensions.ext_database import db
 from models.model import Site
 
-app_site_fields = {
-    'app_id': fields.String,
-    'access_token': fields.String(attribute='code'),
-    'code': fields.String,
-    'title': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String,
-    'description': fields.String,
-    'default_language': fields.String,
-    'customize_domain': fields.String,
-    'copyright': fields.String,
-    'privacy_policy': fields.String,
-    'customize_token_strategy': fields.String,
-    'prompt_public': fields.Boolean
-}
-
 
 def parse_app_site_args():
     parser = reqparse.RequestParser()
@@ -80,6 +66,13 @@ class AppSite(Resource):
             if value is not None:
                 setattr(site, attr_name, value)
 
+                if attr_name == 'title':
+                    app_model.name = value
+                elif attr_name == 'icon':
+                    app_model.icon = value
+                elif attr_name == 'icon_background':
+                    app_model.icon_background = value
+
         db.session.commit()
 
         return site

api/controllers/console/app/statistic.py

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

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('workspace_id', type=str, required=False, nullable=True, location='args')
-        parser.add_argument('email', type=email, required=True, nullable=False, 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(
+        invitation = RegisterService.get_invitation_if_token_valid(workspaceId, reg_email, token)
-            Tenant.id == args['workspace_id'],
-            Tenant.status == 'normal'
-        ).first()
 
-        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('workspace_id', type=str, required=False, nullable=True, location='json')
-        parser.add_argument('email', type=email, required=True, nullable=False, 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'])
+        invitation = RegisterService.get_invitation_if_token_valid(args['workspace_id'], args['email'], args['token'])
-        if account is None:
+        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/auth/data_source_oauth.py

@@ -1,13 +1,13 @@
 import logging
-from datetime import datetime
-from typing import Optional
 
-import flask_login
 import requests
-from flask import request, redirect, current_app, session
+from flask import request, redirect, current_app
-from flask_login import current_user, login_required
+from flask_login import current_user
+
 from flask_restful import Resource
 from werkzeug.exceptions import Forbidden
+
+from libs.login import login_required
 from libs.oauth_data_source import NotionOAuth
 from controllers.console import api
 from ..setup import setup_required
@@ -42,15 +42,34 @@ class OAuthDataSource(Resource):
         if current_app.config.get('NOTION_INTEGRATION_TYPE') == 'internal':
             internal_secret = current_app.config.get('NOTION_INTERNAL_SECRET')
             oauth_provider.save_internal_access_token(internal_secret)
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
+            return { 'data': '' }
         else:
             auth_url = oauth_provider.get_authorization_url()
-            return redirect(auth_url)
+            return { 'data': auth_url }, 200
 
 
 
 
 class OAuthDataSourceCallback(Resource):
+    def get(self, provider: str):
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        if 'code' in request.args:
+            code = request.args.get('code')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&code={code}')
+        elif 'error' in request.args:
+            error = request.args.get('error')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error={error}')
+        else:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error=Access denied')
+        
+
+class OAuthDataSourceBinding(Resource):
     def get(self, provider: str):
         OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
         with current_app.app_context():
@@ -66,12 +85,7 @@ class OAuthDataSourceCallback(Resource):
                     f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")
                 return {'error': 'OAuth data source process failed'}, 400
 
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
+            return {'result': 'success'}, 200
-        elif 'error' in request.args:
-            error = request.args.get('error')
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source={error}')
-        else:
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=access_denied')
 
 
 class OAuthDataSourceSync(Resource):
@@ -98,4 +112,5 @@ class OAuthDataSourceSync(Resource):
 
 api.add_resource(OAuthDataSource, '/oauth/data-source/<string:provider>')
 api.add_resource(OAuthDataSourceCallback, '/oauth/data-source/callback/<string:provider>')
+api.add_resource(OAuthDataSourceBinding, '/oauth/data-source/binding/<string:provider>')
 api.add_resource(OAuthDataSourceSync, '/oauth/data-source/<string:provider>/<uuid:binding_id>/sync')

api/controllers/console/auth/login.py

@@ -6,7 +6,6 @@ from flask_restful import Resource, reqparse
 
 import services
 from controllers.console import api
-from controllers.console.error import AccountNotLinkTenantError
 from controllers.console.setup import setup_required
 from libs.helper import email
 from libs.password import valid_password
@@ -37,12 +36,12 @@ class LoginApi(Resource):
         except Exception:
             pass
 
-        flask_login.login_user(account, remember=args['remember_me'])
         AccountService.update_last_login(account, request)
 
         # todo: return the user info
+        token = AccountService.get_account_jwt_token(account)
 
-        return {'result': 'success'}
+        return {'result': 'success', 'data': token}
 
 
 class LogoutApi(Resource):

api/controllers/console/auth/oauth.py

@@ -2,9 +2,8 @@ import logging
 from datetime import datetime
 from typing import Optional
 
-import flask_login
 import requests
-from flask import request, redirect, current_app, session
+from flask import request, redirect, current_app
 from flask_restful import Resource
 
 from libs.oauth import OAuthUserInfo, GitHubOAuth, GoogleOAuth
@@ -75,12 +74,11 @@ class OAuthCallback(Resource):
             account.initialized_at = datetime.utcnow()
             db.session.commit()
 
-        # login user
-        session.clear()
-        flask_login.login_user(account, remember=True)
         AccountService.update_last_login(account, request)
 
-        return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_login=success')
+        token = AccountService.get_account_jwt_token(account)
+
+        return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?console_token={token}')
 
 
 def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) -> Optional[Account]:

api/controllers/console/datasets/data_source.py

@@ -2,9 +2,10 @@ import datetime
 import json
 
 from cachetools import TTLCache
-from flask import request, current_app
+from flask import request
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, marshal_with, fields, reqparse, marshal
+from libs.login import login_required
+from flask_restful import Resource, marshal_with, reqparse
 from werkzeug.exceptions import NotFound
 
 from controllers.console import api
@@ -13,7 +14,7 @@ from controllers.console.wraps import account_initialization_required
 from core.data_loader.loader.notion import NotionLoader
 from core.indexing_runner import IndexingRunner
 from extensions.ext_database import db
-from libs.helper import TimestampField
+from fields.data_source_fields import integrate_notion_info_list_fields, integrate_list_fields
 from models.dataset import Document
 from models.source import DataSourceBinding
 from services.dataset_service import DatasetService, DocumentService
@@ -21,43 +22,8 @@ from tasks.document_indexing_sync_task import document_indexing_sync_task
 
 cache = TTLCache(maxsize=None, ttl=30)
 
-FILE_SIZE_LIMIT = 15 * 1024 * 1024  # 15MB
-ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm']
-PREVIEW_WORDS_LIMIT = 3000
-
 
 class DataSourceApi(Resource):
-    integrate_icon_fields = {
-        'type': fields.String,
-        'url': fields.String,
-        'emoji': fields.String
-    }
-    integrate_page_fields = {
-        'page_name': fields.String,
-        'page_id': fields.String,
-        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
-        'parent_id': fields.String,
-        'type': fields.String
-    }
-    integrate_workspace_fields = {
-        'workspace_name': fields.String,
-        'workspace_id': fields.String,
-        'workspace_icon': fields.String,
-        'pages': fields.List(fields.Nested(integrate_page_fields)),
-        'total': fields.Integer
-    }
-    integrate_fields = {
-        'id': fields.String,
-        'provider': fields.String,
-        'created_at': TimestampField,
-        'is_bound': fields.Boolean,
-        'disabled': fields.Boolean,
-        'link': fields.String,
-        'source_info': fields.Nested(integrate_workspace_fields)
-    }
-    integrate_list_fields = {
-        'data': fields.List(fields.Nested(integrate_fields)),
-    }
 
     @setup_required
     @login_required
@@ -134,28 +100,6 @@ class DataSourceApi(Resource):
 
 
 class DataSourceNotionListApi(Resource):
-    integrate_icon_fields = {
-        'type': fields.String,
-        'url': fields.String,
-        'emoji': fields.String
-    }
-    integrate_page_fields = {
-        'page_name': fields.String,
-        'page_id': fields.String,
-        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
-        'is_bound': fields.Boolean,
-        'parent_id': fields.String,
-        'type': fields.String
-    }
-    integrate_workspace_fields = {
-        'workspace_name': fields.String,
-        'workspace_id': fields.String,
-        'workspace_icon': fields.String,
-        'pages': fields.List(fields.Nested(integrate_page_fields))
-    }
-    integrate_notion_info_list_fields = {
-        'notion_info': fields.List(fields.Nested(integrate_workspace_fields)),
-    }
 
     @setup_required
     @login_required

api/controllers/console/datasets/datasets.py

@@ -1,7 +1,11 @@
 # -*- coding:utf-8 -*-
-from flask import request
+import flask_restful
-from flask_login import login_required, current_user
+from flask import request, current_app
-from flask_restful import Resource, reqparse, fields, marshal, marshal_with
+from flask_login import current_user
+
+from controllers.console.apikey import api_key_list, api_key_fields
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 import services
 from controllers.console import api
@@ -10,40 +14,16 @@ from controllers.console.datasets.error import DatasetNameDuplicateError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.indexing_runner import IndexingRunner
-from core.model_providers.error import LLMBadRequestError
+from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
-from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.entity.model_params import ModelType
-from libs.helper import TimestampField
+from fields.app_fields import related_app_list
+from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields
+from fields.document_fields import document_status_fields
 from extensions.ext_database import db
 from models.dataset import DocumentSegment, Document
-from models.model import UploadFile
+from models.model import UploadFile, ApiToken
 from services.dataset_service import DatasetService, DocumentService
-
+from services.provider_service import ProviderService
-dataset_detail_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'description': fields.String,
-    'provider': fields.String,
-    'permission': fields.String,
-    'data_source_type': fields.String,
-    'indexing_technique': fields.String,
-    'app_count': fields.Integer,
-    'document_count': fields.Integer,
-    'word_count': fields.Integer,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'updated_by': fields.String,
-    'updated_at': TimestampField,
-}
-
-dataset_query_detail_fields = {
-    "id": fields.String,
-    "content": fields.String,
-    "source": fields.String,
-    "source_app_id": fields.String,
-    "created_by_role": fields.String,
-    "created_by": fields.String,
-    "created_at": TimestampField
-}
 
 
 def _validate_name(name):
@@ -74,8 +54,29 @@ class DatasetListApi(Resource):
             datasets, total = DatasetService.get_datasets(page, limit, provider,
                                                           current_user.current_tenant_id, current_user)
 
+        # check embedding setting
+        provider_service = ProviderService()
+        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
+                                                                 ModelType.EMBEDDINGS.value)
+        # if len(valid_model_list) == 0:
+        #     raise ProviderNotInitializeError(
+        #         f"No Embedding Model available. Please configure a valid provider "
+        #         f"in the Settings -> Model Provider.")
+        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['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'] = True
         response = {
-            'data': marshal(datasets, dataset_detail_fields),
+            'data': data,
             'has_more': len(datasets) == limit,
             'limit': limit,
             'total': total,
@@ -100,15 +101,6 @@ class DatasetListApi(Resource):
         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(
                 tenant_id=current_user.current_tenant_id,
@@ -131,20 +123,40 @@ 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))
-
+        data = marshal(dataset, dataset_detail_fields)
-        return marshal(dataset, dataset_detail_fields), 200
+        # 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,
@@ -232,7 +244,11 @@ 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')
         args = parser.parse_args()
         # validate args
         DocumentService.estimate_args_validate(args)
@@ -250,11 +266,15 @@ class DatasetIndexingEstimateApi(Resource):
 
             try:
                 response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
-                                                                  args['process_rule'], args['doc_form'])
+                                                                  args['process_rule'], args['doc_form'],
+                                                                  args['doc_language'], args['dataset_id'],
+                                                                  args['indexing_technique'])
             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)
         elif args['info_list']['data_source_type'] == 'notion_import':
 
             indexing_runner = IndexingRunner()
@@ -262,29 +282,21 @@ class DatasetIndexingEstimateApi(Resource):
             try:
                 response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
                                                                     args['info_list']['notion_info_list'],
-                                                                    args['process_rule'], args['doc_form'])
+                                                                    args['process_rule'], args['doc_form'],
+                                                                    args['doc_language'], args['dataset_id'],
+                                                                    args['indexing_technique'])
             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)
         else:
             raise ValueError('Data source type not support')
         return response, 200
 
 
 class DatasetRelatedAppListApi(Resource):
-    app_detail_kernel_fields = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-    }
-
-    related_app_list = {
-        'data': fields.List(fields.Nested(app_detail_kernel_fields)),
-        'total': fields.Integer,
-    }
 
     @setup_required
     @login_required
@@ -316,24 +328,6 @@ class DatasetRelatedAppListApi(Resource):
 
 
 class DatasetIndexingStatusApi(Resource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
-
-    document_status_fields_list = {
-        'data': fields.List(fields.Nested(document_status_fields))
-    }
 
     @setup_required
     @login_required
@@ -353,16 +347,101 @@ class DatasetIndexingStatusApi(Resource):
                                                           DocumentSegment.status != 're_segment').count()
             document.completed_segments = completed_segments
             document.total_segments = total_segments
-            documents_status.append(marshal(document, self.document_status_fields))
+            documents_status.append(marshal(document, document_status_fields))
         data = {
             'data': documents_status
         }
         return data
 
 
+class DatasetApiKeyApi(Resource):
+    max_keys = 10
+    token_prefix = 'dataset-'
+    resource_type = 'dataset'
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_key_list)
+    def get(self):
+        keys = db.session.query(ApiToken). \
+            filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
+            all()
+        return {"items": keys}
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_key_fields)
+    def post(self):
+        # 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()
+
+        current_key_count = db.session.query(ApiToken). \
+            filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
+            count()
+
+        if current_key_count >= self.max_keys:
+            flask_restful.abort(
+                400,
+                message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
+                code='max_keys_exceeded'
+            )
+
+        key = ApiToken.generate_api_key(self.token_prefix, 24)
+        api_token = ApiToken()
+        api_token.tenant_id = current_user.current_tenant_id
+        api_token.token = key
+        api_token.type = self.resource_type
+        db.session.add(api_token)
+        db.session.commit()
+        return api_token, 200
+
+
+class DatasetApiDeleteApi(Resource):
+    resource_type = 'dataset'
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, api_key_id):
+        api_key_id = str(api_key_id)
+
+        # 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()
+
+        key = db.session.query(ApiToken). \
+            filter(ApiToken.tenant_id == current_user.current_tenant_id, ApiToken.type == self.resource_type,
+                   ApiToken.id == api_key_id). \
+            first()
+
+        if key is None:
+            flask_restful.abort(404, message='API key not found')
+
+        db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete()
+        db.session.commit()
+
+        return {'result': 'success'}, 204
+
+
+class DatasetApiBaseUrlApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        return {
+            'api_base_url': (current_app.config['SERVICE_API_URL'] if current_app.config['SERVICE_API_URL']
+                             else request.host_url.rstrip('/')) + '/v1'
+        }
+
+
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
 api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
 api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
 api.add_resource(DatasetRelatedAppListApi, '/datasets/<uuid:dataset_id>/related-apps')
 api.add_resource(DatasetIndexingStatusApi, '/datasets/<uuid:dataset_id>/indexing-status')
+api.add_resource(DatasetApiKeyApi, '/datasets/api-keys')
+api.add_resource(DatasetApiDeleteApi, '/datasets/api-keys/<uuid:api_key_id>')
+api.add_resource(DatasetApiBaseUrlApi, '/datasets/api-base-info')

api/controllers/console/datasets/datasets_document.py

@@ -1,10 +1,10 @@
 # -*- coding:utf-8 -*-
-import random
 from datetime import datetime
 from typing import List
 
-from flask import request
+from flask import request, current_app
-from flask_login import login_required, current_user
+from flask_login import current_user
+from libs.login import login_required
 from flask_restful import Resource, fields, marshal, marshal_with, reqparse
 from sqlalchemy import desc, asc
 from werkzeug.exceptions import NotFound, Forbidden
@@ -22,7 +22,8 @@ from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededE
     LLMBadRequestError
 from core.model_providers.model_factory import ModelFactory
 from extensions.ext_redis import redis_client
-from libs.helper import TimestampField
+from fields.document_fields import document_with_segments_fields, document_fields, \
+    dataset_and_document_fields, document_status_fields
 from extensions.ext_database import db
 from models.dataset import DatasetProcessRule, Dataset
 from models.dataset import Document, DocumentSegment
@@ -31,64 +32,6 @@ from services.dataset_service import DocumentService, DatasetService
 from tasks.add_document_to_index_task import add_document_to_index_task
 from tasks.remove_document_from_index_task import remove_document_from_index_task
 
-dataset_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'description': fields.String,
-    'permission': fields.String,
-    'data_source_type': fields.String,
-    'indexing_technique': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-}
-
-document_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'data_source_type': fields.String,
-    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
-    'dataset_process_rule_id': fields.String,
-    'name': fields.String,
-    'created_from': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'tokens': fields.Integer,
-    'indexing_status': fields.String,
-    'error': fields.String,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'archived': fields.Boolean,
-    'display_status': fields.String,
-    'word_count': fields.Integer,
-    'hit_count': fields.Integer,
-    'doc_form': fields.String,
-}
-
-document_with_segments_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'data_source_type': fields.String,
-    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
-    'dataset_process_rule_id': fields.String,
-    'name': fields.String,
-    'created_from': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'tokens': fields.Integer,
-    'indexing_status': fields.String,
-    'error': fields.String,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'archived': fields.Boolean,
-    'display_status': fields.String,
-    'word_count': fields.Integer,
-    'hit_count': fields.Integer,
-    'completed_segments': fields.Integer,
-    'total_segments': fields.Integer
-}
-
 
 class DocumentResource(Resource):
     def get_document(self, dataset_id: str, document_id: str) -> Document:
@@ -137,6 +80,10 @@ class GetProcessRuleApi(Resource):
         req_data = request.args
 
         document_id = req_data.get('document_id')
+        
+        # get default rules
+        mode = DocumentService.DEFAULT_RULES['mode']
+        rules = DocumentService.DEFAULT_RULES['rules']
         if document_id:
             # get the latest process rule
             document = Document.query.get_or_404(document_id)
@@ -157,11 +104,9 @@ class GetProcessRuleApi(Resource):
                 order_by(DatasetProcessRule.created_at.desc()). \
                 limit(1). \
                 one_or_none()
-            mode = dataset_process_rule.mode
+            if dataset_process_rule:
-            rules = dataset_process_rule.rules_dict
+                mode = dataset_process_rule.mode
-        else:
+                rules = dataset_process_rule.rules_dict
-            mode = DocumentService.DEFAULT_RULES['mode']
-            rules = DocumentService.DEFAULT_RULES['rules']
 
         return {
             'mode': mode,
@@ -274,6 +219,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')
         args = parser.parse_args()
 
         if not dataset.indexing_technique and not args['indexing_technique']:
@@ -282,15 +229,6 @@ class DatasetDocumentListApi(Resource):
         # validate args
         DocumentService.document_create_args_validate(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.")
-
         try:
             documents, batch = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
         except ProviderTokenNotInitError as ex:
@@ -307,11 +245,6 @@ class DatasetDocumentListApi(Resource):
 
 
 class DatasetInitApi(Resource):
-    dataset_and_document_fields = {
-        'dataset': fields.Nested(dataset_fields),
-        'documents': fields.List(fields.Nested(document_fields)),
-        'batch': fields.String
-    }
 
     @setup_required
     @login_required
@@ -328,16 +261,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')
         args = parser.parse_args()
-
+        if args['indexing_technique'] == 'high_quality':
-        try:
+            try:
-            ModelFactory.get_embedding_model(
+                ModelFactory.get_embedding_model(
-                tenant_id=current_user.current_tenant_id
+                    tenant_id=current_user.current_tenant_id
-            )
+                )
-        except LLMBadRequestError:
+            except LLMBadRequestError:
-            raise ProviderNotInitializeError(
+                raise ProviderNotInitializeError(
-                f"No Embedding Model available. Please configure a valid provider "
+                    f"No Embedding Model available. Please configure a valid provider "
-                f"in the Settings -> Model Provider.")
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
 
         # validate args
         DocumentService.document_create_args_validate(args)
@@ -406,11 +343,14 @@ class DocumentIndexingEstimateApi(DocumentResource):
 
                 try:
                     response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, [file],
-                                                                      data_process_rule_dict)
+                                                                      data_process_rule_dict, None,
+                                                                      'English', dataset_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)
 
         return response
 
@@ -473,46 +413,34 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
             indexing_runner = IndexingRunner()
             try:
                 response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
-                                                                  data_process_rule_dict)
+                                                                  data_process_rule_dict, None,
+                                                                  'English', dataset_id)
             except LLMBadRequestError:
                 raise ProviderNotInitializeError(
                     f"No Embedding Model available. Please configure a valid provider "
                     f"in the Settings -> Model Provider.")
-        elif dataset.data_source_type:
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+        elif dataset.data_source_type == 'notion_import':
 
             indexing_runner = IndexingRunner()
             try:
                 response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
                                                                     info_list,
-                                                                    data_process_rule_dict)
+                                                                    data_process_rule_dict,
+                                                                    None, 'English', dataset_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)
         else:
             raise ValueError('Data source type not support')
         return response
 
 
 class DocumentBatchIndexingStatusApi(DocumentResource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
-
-    document_status_fields_list = {
-        'data': fields.List(fields.Nested(document_status_fields))
-    }
 
     @setup_required
     @login_required
@@ -532,7 +460,7 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
             document.total_segments = total_segments
             if document.is_paused:
                 document.indexing_status = 'paused'
-            documents_status.append(marshal(document, self.document_status_fields))
+            documents_status.append(marshal(document, document_status_fields))
         data = {
             'data': documents_status
         }
@@ -540,20 +468,6 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
 
 
 class DocumentIndexingStatusApi(DocumentResource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
 
     @setup_required
     @login_required
@@ -575,8 +489,9 @@ class DocumentIndexingStatusApi(DocumentResource):
 
         document.completed_segments = completed_segments
         document.total_segments = total_segments
-
+        if document.is_paused:
-        return marshal(document, self.document_status_fields)
+            document.indexing_status = 'paused'
+        return marshal(document, document_status_fields)
 
 
 class DocumentDetailApi(DocumentResource):
@@ -709,6 +624,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:
@@ -749,11 +670,13 @@ class DocumentMetadataApi(DocumentResource):
         metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
 
         document.doc_metadata = {}
-
+        if doc_type == 'others':
-        for key, value_type in metadata_schema.items():
+            document.doc_metadata = doc_metadata
-            value = doc_metadata.get(key)
+        else:
-            if value is not None and isinstance(value, value_type):
+            for key, value_type in metadata_schema.items():
-                document.doc_metadata[key] = value
+                value = doc_metadata.get(key)
+                if value is not None and isinstance(value, value_type):
+                    document.doc_metadata[key] = value
 
         document.doc_type = doc_type
         document.updated_at = datetime.utcnow()
@@ -769,6 +692,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
@@ -832,12 +761,40 @@ class DocumentStatusApi(DocumentResource):
 
                 remove_document_from_index_task.delay(document_id)
 
+            return {'result': 'success'}, 200
+        elif action == "un_archive":
+            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
+            document.updated_at = datetime.utcnow()
+            db.session.commit()
+
+            # Set cache to prevent indexing the same document multiple times
+            redis_client.setex(indexing_cache_key, 600, 1)
+
+            add_document_to_index_task.delay(document_id)
+
             return {'result': 'success'}, 200
         else:
             raise InvalidActionError()
 
 
 class DocumentPauseApi(DocumentResource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
     def patch(self, dataset_id, document_id):
         """pause document."""
         dataset_id = str(dataset_id)
@@ -867,6 +824,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)
@@ -892,6 +852,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')
@@ -917,3 +892,4 @@ api.add_resource(DocumentStatusApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/status/<string:action>')
 api.add_resource(DocumentPauseApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause')
 api.add_resource(DocumentRecoverApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume')
+api.add_resource(DocumentLimitApi, '/datasets/limit')

api/controllers/console/datasets/datasets_segments.py

@@ -1,53 +1,30 @@
 # -*- coding:utf-8 -*-
+import uuid
 from datetime import datetime
-
+from flask import request
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal
+from flask_restful import Resource, reqparse, marshal
 from werkzeug.exceptions import NotFound, Forbidden
 
 import services
 from controllers.console import api
-from controllers.console.datasets.error import InvalidActionError
+from controllers.console.app.error import ProviderNotInitializeError
+from controllers.console.datasets.error import InvalidActionError, NoFileUploadedError, TooManyFilesError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
+from core.model_providers.model_factory import ModelFactory
+from libs.login import login_required
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
+from fields.segment_fields import segment_fields
 from models.dataset import DocumentSegment
 
-from libs.helper import TimestampField
 from services.dataset_service import DatasetService, DocumentService, SegmentService
 from tasks.enable_segment_to_index_task import enable_segment_to_index_task
-from tasks.remove_segment_from_index_task import remove_segment_from_index_task
+from tasks.disable_segment_from_index_task import disable_segment_from_index_task
-
+from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task
-segment_fields = {
+import pandas as pd
-    'id': fields.String,
-    'position': fields.Integer,
-    'document_id': fields.String,
-    'content': fields.String,
-    'answer': fields.String,
-    'word_count': fields.Integer,
-    'tokens': fields.Integer,
-    'keywords': fields.List(fields.String),
-    'index_node_id': fields.String,
-    'index_node_hash': fields.String,
-    'hit_count': fields.Integer,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'status': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'indexing_at': TimestampField,
-    'completed_at': TimestampField,
-    'error': fields.String,
-    'stopped_at': TimestampField
-}
-
-segment_list_response = {
-    'data': fields.List(fields.Nested(segment_fields)),
-    'has_more': fields.Boolean,
-    'limit': fields.Integer
-}
 
 
 class DatasetDocumentSegmentListApi(Resource):
@@ -142,7 +119,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()
@@ -151,6 +129,20 @@ class DatasetDocumentSegmentApi(Resource):
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
+        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),
@@ -197,7 +189,7 @@ class DatasetDocumentSegmentApi(Resource):
             # Set cache to prevent indexing the same segment multiple times
             redis_client.setex(indexing_cache_key, 600, 1)
 
-            remove_segment_from_index_task.delay(segment.id)
+            disable_segment_from_index_task.delay(segment.id)
 
             return {'result': 'success'}, 200
         else:
@@ -222,6 +214,20 @@ class DatasetDocumentSegmentAddApi(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()
+        # check embedding model setting
+        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:
@@ -233,7 +239,7 @@ class DatasetDocumentSegmentAddApi(Resource):
         parser.add_argument('keywords', type=list, required=False, nullable=True, location='json')
         args = parser.parse_args()
         SegmentService.segment_create_args_validate(args, document)
-        segment = SegmentService.create_segment(args, document)
+        segment = SegmentService.create_segment(args, document, dataset)
         return {
             'data': marshal(segment, segment_fields),
             'doc_form': document.doc_form
@@ -250,12 +256,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 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),
@@ -277,12 +299,115 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         parser.add_argument('keywords', type=list, required=False, nullable=True, location='json')
         args = parser.parse_args()
         SegmentService.segment_create_args_validate(args, document)
-        segment = SegmentService.update_segment(args, segment, document)
+        segment = SegmentService.update_segment(args, segment, document, dataset)
         return {
             'data': marshal(segment, segment_fields),
             'doc_form': document.doc_form
         }, 200
 
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, dataset_id, document_id, segment_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        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 segment
+        segment_id = str(segment_id)
+        segment = DocumentSegment.query.filter(
+            DocumentSegment.id == str(segment_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        ).first()
+        if not segment:
+            raise NotFound('Segment not found.')
+        # 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:
+            DatasetService.check_dataset_permission(dataset, current_user)
+        except services.errors.account.NoPermissionError as e:
+            raise Forbidden(str(e))
+        SegmentService.delete_segment(segment, document, dataset)
+        return {'result': 'success'}, 200
+
+
+class DatasetDocumentSegmentBatchImportApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, dataset_id, document_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id)
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # check document
+        document_id = str(document_id)
+        document = DocumentService.get_document(dataset_id, document_id)
+        if not document:
+            raise NotFound('Document not found.')
+        # get file from request
+        file = request.files['file']
+        # check file
+        if 'file' not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+        # check file type
+        if not file.filename.endswith('.csv'):
+            raise ValueError("Invalid file type. Only CSV files are allowed")
+
+        try:
+            # Skip the first row
+            df = pd.read_csv(file)
+            result = []
+            for index, row in df.iterrows():
+                if document.doc_form == 'qa_model':
+                    data = {'content': row[0], 'answer': row[1]}
+                else:
+                    data = {'content': row[0]}
+                result.append(data)
+            if len(result) == 0:
+                raise ValueError("The CSV file is empty.")
+            # async job
+            job_id = str(uuid.uuid4())
+            indexing_cache_key = 'segment_batch_import_{}'.format(str(job_id))
+            # send batch add segments task
+            redis_client.setnx(indexing_cache_key, 'waiting')
+            batch_create_segment_to_index_task.delay(str(job_id), result, dataset_id, document_id,
+                                                     current_user.current_tenant_id, current_user.id)
+        except Exception as e:
+            return {'error': str(e)}, 500
+        return {
+            'job_id': job_id,
+            'job_status': 'waiting'
+        }, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, job_id):
+        job_id = str(job_id)
+        indexing_cache_key = 'segment_batch_import_{}'.format(job_id)
+        cache_result = redis_client.get(indexing_cache_key)
+        if cache_result is None:
+            raise ValueError("The job is not exist.")
+
+        return {
+            'job_id': job_id,
+            'job_status': cache_result.decode()
+        }, 200
+
 
 api.add_resource(DatasetDocumentSegmentListApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments')
@@ -292,3 +417,6 @@ api.add_resource(DatasetDocumentSegmentAddApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segment')
 api.add_resource(DatasetDocumentSegmentUpdateApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>')
+api.add_resource(DatasetDocumentSegmentBatchImportApi,
+                 '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/batch_import',
+                 '/datasets/batch_import_status/<uuid:job_id>')

api/controllers/console/datasets/file.py

@@ -1,45 +1,39 @@
-import datetime
-import hashlib
-import tempfile
-import chardet
-import time
-import uuid
-from pathlib import Path
-
 from cachetools import TTLCache
 from flask import request, current_app
-from flask_login import login_required, current_user
+
-from flask_restful import Resource, marshal_with, fields
+import services
-from werkzeug.exceptions import NotFound
+from libs.login import login_required
+from flask_restful import Resource, marshal_with
 
 from controllers.console import api
 from controllers.console.datasets.error import NoFileUploadedError, TooManyFilesError, FileTooLargeError, \
     UnsupportedFileTypeError
+
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.data_loader.file_extractor import FileExtractor
+from fields.file_fields import upload_config_fields, file_fields
-from extensions.ext_storage import storage
+
-from libs.helper import TimestampField
+from services.file_service import FileService
-from extensions.ext_database import db
-from models.model import UploadFile
 
 cache = TTLCache(maxsize=None, ttl=30)
 
-FILE_SIZE_LIMIT = 15 * 1024 * 1024  # 15MB
+ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'docx', 'csv']
-ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx']
 PREVIEW_WORDS_LIMIT = 3000
 
 
 class FileApi(Resource):
-    file_fields = {
+
-        'id': fields.String,
+    @setup_required
-        'name': fields.String,
+    @login_required
-        'size': fields.Integer,
+    @account_initialization_required
-        'extension': fields.String,
+    @marshal_with(upload_config_fields)
-        'mime_type': fields.String,
+    def get(self):
-        'created_by': fields.String,
+        file_size_limit = current_app.config.get("UPLOAD_FILE_SIZE_LIMIT")
-        'created_at': TimestampField,
+        batch_count_limit = current_app.config.get("UPLOAD_FILE_BATCH_LIMIT")
-    }
+        return {
+            'file_size_limit': file_size_limit,
+            'batch_count_limit': batch_count_limit
+        }, 200
 
     @setup_required
     @login_required
@@ -56,44 +50,13 @@ class FileApi(Resource):
 
         if len(request.files) > 1:
             raise TooManyFilesError()
-
+        try:
-        file_content = file.read()
+            upload_file = FileService.upload_file(file)
-        file_size = len(file_content)
+        except services.errors.file.FileTooLargeError as file_too_large_error:
-
+            raise FileTooLargeError(file_too_large_error.description)
-        if file_size > FILE_SIZE_LIMIT:
+        except services.errors.file.UnsupportedFileTypeError:
-            message = "({file_size} > {FILE_SIZE_LIMIT})"
-            raise FileTooLargeError(message)
-
-        extension = file.filename.split('.')[-1]
-        if extension not in ALLOWED_EXTENSIONS:
             raise UnsupportedFileTypeError()
 
-        # user uuid as file name
-        file_uuid = str(uuid.uuid4())
-        file_key = 'upload_files/' + current_user.current_tenant_id + '/' + file_uuid + '.' + extension
-
-        # save file to storage
-        storage.save(file_key, file_content)
-
-        # save file to db
-        config = current_app.config
-        upload_file = UploadFile(
-            tenant_id=current_user.current_tenant_id,
-            storage_type=config['STORAGE_TYPE'],
-            key=file_key,
-            name=file.filename,
-            size=file_size,
-            extension=extension,
-            mime_type=file.mimetype,
-            created_by=current_user.id,
-            created_at=datetime.datetime.utcnow(),
-            used=False,
-            hash=hashlib.sha3_256(file_content).hexdigest()
-        )
-
-        db.session.add(upload_file)
-        db.session.commit()
-
         return upload_file, 201
 
 
@@ -103,26 +66,7 @@ class FilePreviewApi(Resource):
     @account_initialization_required
     def get(self, file_id):
         file_id = str(file_id)
-
+        text = FileService.get_file_preview(file_id)
-        key = file_id + request.path
-        cached_response = cache.get(key)
-        if cached_response and time.time() - cached_response['timestamp'] < cache.ttl:
-            return cached_response['response']
-
-        upload_file = db.session.query(UploadFile) \
-            .filter(UploadFile.id == file_id) \
-            .first()
-
-        if not upload_file:
-            raise NotFound("File not found")
-
-        # extract text from file
-        extension = upload_file.extension
-        if extension not in ALLOWED_EXTENSIONS:
-            raise UnsupportedFileTypeError()
-
-        text = FileExtractor.load(upload_file, return_text=True)
-        text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
         return {'content': text}
 
 

api/controllers/console/datasets/hit_testing.py

@@ -1,7 +1,8 @@
 import logging
 
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, marshal, fields
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal
 from werkzeug.exceptions import InternalServerError, NotFound, Forbidden
 
 import services
@@ -11,49 +12,12 @@ from controllers.console.app.error import ProviderNotInitializeError, ProviderQu
 from controllers.console.datasets.error import HighQualityDatasetOnlyError, DatasetNotInitializedError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError, \
-from libs.helper import TimestampField
+    LLMBadRequestError
+from fields.hit_testing_fields import hit_testing_record_fields
 from services.dataset_service import DatasetService
 from services.hit_testing_service import HitTestingService
 
-document_fields = {
-    'id': fields.String,
-    'data_source_type': fields.String,
-    'name': fields.String,
-    'doc_type': fields.String,
-}
-
-segment_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'document_id': fields.String,
-    'content': fields.String,
-    'answer': fields.String,
-    'word_count': fields.Integer,
-    'tokens': fields.Integer,
-    'keywords': fields.List(fields.String),
-    'index_node_id': fields.String,
-    'index_node_hash': fields.String,
-    'hit_count': fields.Integer,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'status': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'indexing_at': TimestampField,
-    'completed_at': TimestampField,
-    'error': fields.String,
-    'stopped_at': TimestampField,
-    'document': fields.Nested(document_fields),
-}
-
-hit_testing_record_fields = {
-    'segment': fields.Nested(segment_fields),
-    'score': fields.Float,
-    'tsne_position': fields.Raw
-}
-
 
 class HitTestingApi(Resource):
 
@@ -102,6 +66,10 @@ class HitTestingApi(Resource):
             raise ProviderQuotaExceededError()
         except ModelCurrentlyNotSupportError:
             raise ProviderModelCurrentlyNotSupportError()
+        except LLMBadRequestError:
+            raise ProviderNotInitializeError(
+                f"No Embedding Model available. Please configure a valid provider "
+                f"in the Settings -> Model Provider.")
         except ValueError as e:
             raise ValueError(str(e))
         except Exception as e:

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/conversation.py

@@ -7,26 +7,12 @@ from werkzeug.exceptions import NotFound
 from controllers.console import api
 from controllers.console.explore.error import NotChatAppError
 from controllers.console.explore.wraps import InstalledAppResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationListApi(InstalledAppResource):
 
@@ -76,7 +62,7 @@ class ConversationApi(InstalledAppResource):
 
 class ConversationRenameApi(InstalledAppResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, installed_app, c_id):
         app_model = installed_app.app
         if app_model.mode != 'chat':

api/controllers/console/explore/installed_app.py

@@ -1,8 +1,9 @@
 # -*- coding:utf-8 -*-
 from datetime import datetime
 
-from flask_login import login_required, current_user
+from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal_with, inputs
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with, inputs
 from sqlalchemy import and_
 from werkzeug.exceptions import NotFound, Forbidden, BadRequest
 
@@ -10,32 +11,10 @@ from controllers.console import api
 from controllers.console.explore.wraps import InstalledAppResource
 from controllers.console.wraps import account_initialization_required
 from extensions.ext_database import db
-from libs.helper import TimestampField
+from fields.installed_app_fields import installed_app_list_fields
 from models.model import App, InstalledApp, RecommendedApp
 from services.account_service import TenantService
 
-app_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'mode': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String
-}
-
-installed_app_fields = {
-    'id': fields.String,
-    'app': fields.Nested(app_fields),
-    'app_owner_tenant_id': fields.String,
-    'is_pinned': fields.Boolean,
-    'last_used_at': TimestampField,
-    'editable': fields.Boolean,
-    'uninstallable': fields.Boolean,
-}
-
-installed_app_list_fields = {
-    'installed_apps': fields.List(fields.Nested(installed_app_fields))
-}
-
 
 class InstalledAppsListApi(Resource):
     @login_required

api/controllers/console/explore/message.py

@@ -17,6 +17,7 @@ from controllers.console.explore.error import NotCompletionAppError, AppSuggeste
 from controllers.console.explore.wraps import InstalledAppResource
 from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
+from fields.message_fields import message_infinite_scroll_pagination_fields
 from libs.helper import uuid_value, TimestampField
 from services.completion_service import CompletionService
 from services.errors.app import MoreLikeThisDisabledError
@@ -26,25 +27,6 @@ from services.message_service import MessageService
 
 
 class MessageListApi(InstalledAppResource):
-    feedback_fields = {
-        'rating': fields.String
-    }
-
-    message_fields = {
-        'id': fields.String,
-        'conversation_id': fields.String,
-        'inputs': fields.Raw,
-        'query': fields.String,
-        'answer': fields.String,
-        'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
-        'created_at': TimestampField
-    }
-
-    message_infinite_scroll_pagination_fields = {
-        'limit': fields.Integer,
-        'has_more': fields.Boolean,
-        'data': fields.List(fields.Nested(message_fields))
-    }
 
     @marshal_with(message_infinite_scroll_pagination_fields)
     def get(self, installed_app):

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/explore/recommended_app.py

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

api/controllers/console/explore/wraps.py

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

api/controllers/console/setup.py

@@ -1,7 +1,6 @@
 # -*- coding:utf-8 -*-
 from functools import wraps
 
-import flask_login
 from flask import request, current_app
 from flask_restful import Resource, reqparse
 
@@ -58,9 +57,6 @@ class SetupApi(Resource):
         )
 
         setup()
-
-        # Login
-        flask_login.login_user(account)
         AccountService.update_last_login(account, request)
 
         return {'result': 'success'}, 201

api/controllers/console/universal_chat/chat.py

@@ -29,6 +29,7 @@ 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

api/controllers/console/universal_chat/conversation.py

@@ -6,31 +6,17 @@ from werkzeug.exceptions import NotFound
 
 from controllers.console import api
 from controllers.console.universal_chat.wraps import UniversalChatResource
+from fields.conversation_fields import conversation_with_model_config_infinite_scroll_pagination_fields, \
+    conversation_with_model_config_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField,
-    'model_config': fields.Raw,
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class UniversalChatConversationListApi(UniversalChatResource):
 
-    @marshal_with(conversation_infinite_scroll_pagination_fields)
+    @marshal_with(conversation_with_model_config_infinite_scroll_pagination_fields)
     def get(self, universal_app):
         app_model = universal_app
 
@@ -73,7 +59,7 @@ class UniversalChatConversationApi(UniversalChatResource):
 
 class UniversalChatConversationRenameApi(UniversalChatResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(conversation_with_model_config_fields)
     def post(self, universal_app, c_id):
         app_model = universal_app
         conversation_id = str(c_id)

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

@@ -1,7 +1,8 @@
 import json
 from functools import wraps
 
-from flask_login import login_required, current_user
+from flask_login import current_user
+from libs.login import login_required
 from flask_restful import Resource
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
@@ -46,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/webhook/stripe.py

@@ -38,12 +38,20 @@ class StripeWebhookApi(Resource):
             logging.debug(event['data']['object']['payment_status'])
             logging.debug(event['data']['object']['metadata'])
 
+            session = stripe.checkout.Session.retrieve(
+                event['data']['object']['id'],
+                expand=['line_items'],
+            )
+
+            logging.debug(session.line_items['data'][0]['quantity'])
+
             # Fulfill the purchase...
             provider_checkout_service = ProviderCheckoutService()
 
             try:
-                provider_checkout_service.fulfill_provider_order(event)
+                provider_checkout_service.fulfill_provider_order(event, session.line_items)
             except Exception as e:
+
                 logging.debug(str(e))
                 return 'success', 200
 

api/controllers/console/workspace/account.py

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

api/controllers/console/workspace/members.py

@@ -1,6 +1,7 @@
 # -*- coding:utf-8 -*-
 from flask import current_app
-from flask_login import login_required, current_user
+from flask_login import current_user
+from libs.login import login_required
 from flask_restful import Resource, reqparse, marshal_with, abort, fields, marshal
 
 import services
@@ -48,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
-
+        invitation_results = []
-        try:
+        console_web_url = current_app.config.get("CONSOLE_WEB_URL")
-            token = RegisterService.invite_new_member(inviter.current_tenant, invitee_email, role=invitee_role,
+        for invitee_email in invitee_emails:
-                                                      inviter=inviter)
+            try:
-            account = db.session.query(Account, TenantAccountJoin.role).join(
+                token = RegisterService.invite_new_member(inviter.current_tenant, invitee_email, role=invitee_role,
-                TenantAccountJoin, Account.id == TenantAccountJoin.account_id
+                                                        inviter=inviter)
-            ).filter(Account.email == args['email']).first()
+                account = db.session.query(Account, TenantAccountJoin.role).join(
-            account, role = account
+                    TenantAccountJoin, Account.id == TenantAccountJoin.account_id
-            account = marshal(account, account_fields)
+                ).filter(Account.email == invitee_email).first()
-            account['role'] = role
+                account, role = account
-        except services.errors.account.CannotOperateSelfError as e:
+                invitation_results.append({
-            return {'code': 'cannot-operate-self', 'message': str(e)}, 400
+                    'status': 'success',
-        except services.errors.account.NoPermissionError as e:
+                    'email': invitee_email,
-            return {'code': 'forbidden', 'message': str(e)}, 403
+                    'url': f'{console_web_url}/activate?email={invitee_email}&token={token}'
-        except services.errors.account.AccountAlreadyInTenantError as e:
+                })
-            return {'code': 'email-taken', 'message': str(e)}, 409
+                account = marshal(account, account_fields)
-        except Exception as e:
+                account['role'] = role
-            return {'code': 'unexpected-error', 'message': str(e)}, 500
+            except Exception as e:
-
+                invitation_results.append({
-        # todo:413
+                    'status': 'failed',
+                    'email': invitee_email,
+                    'message': str(e)
+                })
 
         return {
             'result': 'success',
-            'account': account,
+            'invitation_results': invitation_results,
-            'invite_url': '{}/activate?workspace_id={}&email={}&token={}'.format(
-                current_app.config.get("CONSOLE_WEB_URL"),
-                str(current_user.current_tenant_id),
-                invitee_email,
-                token
-            )
         }, 201
 
 

api/controllers/console/workspace/model_providers.py

@@ -1,4 +1,5 @@
-from flask_login import login_required, current_user
+from flask_login import current_user
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 
@@ -245,7 +246,8 @@ class ModelProviderModelParameterRuleApi(Resource):
                 'enabled': v.enabled,
                 'min': v.min,
                 'max': v.max,
-                'default': v.default
+                'default': v.default,
+                'precision': v.precision
             }
             for k, v in vars(parameter_rules).items()
         }
@@ -284,6 +286,25 @@ class ModelProviderFreeQuotaSubmitApi(Resource):
         return result
 
 
+class ModelProviderFreeQuotaQualificationVerifyApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_name: str):
+        parser = reqparse.RequestParser()
+        parser.add_argument('token', type=str, required=False, nullable=True, location='args')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+        result = provider_service.free_quota_qualification_verify(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name,
+            token=args['token']
+        )
+
+        return result
+
+
 api.add_resource(ModelProviderListApi, '/workspaces/current/model-providers')
 api.add_resource(ModelProviderValidateApi, '/workspaces/current/model-providers/<string:provider_name>/validate')
 api.add_resource(ModelProviderUpdateApi, '/workspaces/current/model-providers/<string:provider_name>')
@@ -299,3 +320,5 @@ api.add_resource(ModelProviderPaymentCheckoutUrlApi,
                  '/workspaces/current/model-providers/<string:provider_name>/checkout-url')
 api.add_resource(ModelProviderFreeQuotaSubmitApi,
                  '/workspaces/current/model-providers/<string:provider_name>/free-quota-submit')
+api.add_resource(ModelProviderFreeQuotaQualificationVerifyApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/free-quota-qualification-verify')

api/controllers/console/workspace/models.py

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

api/controllers/console/workspace/providers.py

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

api/controllers/console/workspace/tool_providers.py

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

api/controllers/console/workspace/workspace.py

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

api/controllers/service_api/__init__.py

@@ -9,4 +9,4 @@ api = ExternalApi(bp)
 
 from .app import completion, app, conversation, message, audio
 
-from .dataset import document
+from .dataset import document, segment, dataset

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/conversation.py

@@ -8,25 +8,11 @@ from controllers.service_api import api
 from controllers.service_api.app import create_or_update_end_user_for_user_id
 from controllers.service_api.app.error import NotChatAppError
 from controllers.service_api.wraps import AppApiResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 import services
 from services.conversation_service import ConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationApi(AppApiResource):
 
@@ -50,7 +36,7 @@ class ConversationApi(AppApiResource):
             raise NotFound("Last Conversation Not Exists.")
 
 class ConversationDetailApi(AppApiResource):
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def delete(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()
@@ -70,7 +56,7 @@ class ConversationDetailApi(AppApiResource):
 
 class ConversationRenameApi(AppApiResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()

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/service_api/dataset/dataset.py

@@ -0,0 +1,81 @@
+from flask import request
+from flask_restful import reqparse, marshal
+import services.dataset_service
+from controllers.service_api import api
+from controllers.service_api.dataset.error import DatasetNameDuplicateError
+from controllers.service_api.wraps import DatasetApiResource
+from libs.login import current_user
+from core.model_providers.models.entity.model_params import ModelType
+from fields.dataset_fields import dataset_detail_fields
+from services.dataset_service import DatasetService
+from services.provider_service import ProviderService
+
+
+def _validate_name(name):
+    if not name or len(name) < 1 or len(name) > 40:
+        raise ValueError('Name must be between 1 to 40 characters.')
+    return name
+
+
+class DatasetApi(DatasetApiResource):
+    """Resource for get datasets."""
+
+    def get(self, tenant_id):
+        page = request.args.get('page', default=1, type=int)
+        limit = request.args.get('limit', default=20, type=int)
+        provider = request.args.get('provider', default="vendor")
+        datasets, total = DatasetService.get_datasets(page, limit, provider,
+                                                      tenant_id, current_user)
+        # check embedding setting
+        provider_service = ProviderService()
+        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']}")
+        data = marshal(datasets, dataset_detail_fields)
+        for item in data:
+            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'] = True
+        response = {
+            'data': data,
+            'has_more': len(datasets) == limit,
+            'limit': limit,
+            'total': total,
+            'page': page
+        }
+        return response, 200
+
+    """Resource for datasets."""
+
+    def post(self, tenant_id):
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', nullable=False, required=True,
+                            help='type is required. Name must be between 1 to 40 characters.',
+                            type=_validate_name)
+        parser.add_argument('indexing_technique', type=str, location='json',
+                            choices=('high_quality', 'economy'),
+                            help='Invalid indexing technique.')
+        args = parser.parse_args()
+
+        try:
+            dataset = DatasetService.create_empty_dataset(
+                tenant_id=tenant_id,
+                name=args['name'],
+                indexing_technique=args['indexing_technique'],
+                account=current_user
+            )
+        except services.errors.dataset.DatasetNameDuplicateError:
+            raise DatasetNameDuplicateError()
+
+        return marshal(dataset, dataset_detail_fields), 200
+
+
+api.add_resource(DatasetApi, '/datasets')
+

api/controllers/service_api/dataset/document.py

@@ -1,114 +1,287 @@
-import datetime
+import json
-import uuid
 
-from flask import current_app
+from flask import request
-from flask_restful import reqparse
+from flask_restful import reqparse, marshal
+from sqlalchemy import desc
 from werkzeug.exceptions import NotFound
 
 import services.dataset_service
 from controllers.service_api import api
 from controllers.service_api.app.error import ProviderNotInitializeError
 from controllers.service_api.dataset.error import ArchivedDocumentImmutableError, DocumentIndexingError, \
-    DatasetNotInitedError
+    NoFileUploadedError, TooManyFilesError
 from controllers.service_api.wraps import DatasetApiResource
+from libs.login import current_user
 from core.model_providers.error import ProviderTokenNotInitError
 from extensions.ext_database import db
-from extensions.ext_storage import storage
+from fields.document_fields import document_fields, document_status_fields
-from models.model import UploadFile
+from models.dataset import Dataset, Document, DocumentSegment
 from services.dataset_service import DocumentService
+from services.file_service import FileService
 
 
-class DocumentListApi(DatasetApiResource):
+class DocumentAddByTextApi(DatasetApiResource):
     """Resource for documents."""
 
-    def post(self, dataset):
+    def post(self, tenant_id, dataset_id):
-        """Create document."""
+        """Create document by text."""
         parser = reqparse.RequestParser()
         parser.add_argument('name', type=str, required=True, nullable=False, location='json')
         parser.add_argument('text', type=str, required=True, nullable=False, location='json')
-        parser.add_argument('doc_type', type=str, location='json')
+        parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
-        parser.add_argument('doc_metadata', type=dict, 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('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False,
+                            location='json')
         args = parser.parse_args()
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
 
-        if not dataset.indexing_technique:
+        if not dataset:
-            raise DatasetNotInitedError("Dataset indexing technique must be set.")
+            raise ValueError('Dataset is not exist.')
 
-        doc_type = args.get('doc_type')
+        if not dataset.indexing_technique and not args['indexing_technique']:
-        doc_metadata = args.get('doc_metadata')
+            raise ValueError('indexing_technique is required.')
 
-        if doc_type and doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
+        upload_file = FileService.upload_text(args.get('text'), args.get('name'))
-            raise ValueError('Invalid doc_type.')
+        data_source = {
-
+            'type': 'upload_file',
-        # user uuid as file name
+            'info_list': {
-        file_uuid = str(uuid.uuid4())
+                'data_source_type': 'upload_file',
-        file_key = 'upload_files/' + dataset.tenant_id + '/' + file_uuid + '.txt'
+                'file_info_list': {
-
+                    'file_ids': [upload_file.id]
-        # save file to storage
+                }
-        storage.save(file_key, args.get('text'))
-
-        # save file to db
-        config = current_app.config
-        upload_file = UploadFile(
-            tenant_id=dataset.tenant_id,
-            storage_type=config['STORAGE_TYPE'],
-            key=file_key,
-            name=args.get('name') + '.txt',
-            size=len(args.get('text')),
-            extension='txt',
-            mime_type='text/plain',
-            created_by=dataset.created_by,
-            created_at=datetime.datetime.utcnow(),
-            used=True,
-            used_by=dataset.created_by,
-            used_at=datetime.datetime.utcnow()
-        )
-
-        db.session.add(upload_file)
-        db.session.commit()
-
-        document_data = {
-            'data_source': {
-                'type': 'upload_file',
-                'info': [
-                    {
-                        'upload_file_id': upload_file.id
-                    }
-                ]
             }
         }
+        args['data_source'] = data_source
+        # validate args
+        DocumentService.document_create_args_validate(args)
 
         try:
             documents, batch = DocumentService.save_document_with_dataset_id(
                 dataset=dataset,
-                document_data=document_data,
+                document_data=args,
-                account=dataset.created_by_account,
+                account=current_user,
-                dataset_process_rule=dataset.latest_process_rule,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
                 created_from='api'
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
-        if doc_type and doc_metadata:
-            metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
 
-            document.doc_metadata = {}
+        documents_and_batch_fields = {
-
+            'document': marshal(document, document_fields),
-            for key, value_type in metadata_schema.items():
+            'batch': batch
-                value = doc_metadata.get(key)
+        }
-                if value is not None and isinstance(value, value_type):
+        return documents_and_batch_fields, 200
-                    document.doc_metadata[key] = value
-
-            document.doc_type = doc_type
-            document.updated_at = datetime.datetime.utcnow()
-            db.session.commit()
-
-        return {'id': document.id}
 
 
-class DocumentApi(DatasetApiResource):
+class DocumentUpdateByTextApi(DatasetApiResource):
-    def delete(self, dataset, document_id):
+    """Resource for update documents."""
+
+    def post(self, tenant_id, dataset_id, document_id):
+        """Update document by text."""
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=False, nullable=True, location='json')
+        parser.add_argument('text', type=str, required=False, nullable=True, location='json')
+        parser.add_argument('process_rule', type=dict, required=False, 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')
+        args = parser.parse_args()
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+
+        if args['text']:
+            upload_file = FileService.upload_text(args.get('text'), args.get('name'))
+            data_source = {
+                'type': 'upload_file',
+                'info_list': {
+                    'data_source_type': 'upload_file',
+                    'file_info_list': {
+                        'file_ids': [upload_file.id]
+                    }
+                }
+            }
+            args['data_source'] = data_source
+        # validate args
+        args['original_document_id'] = str(document_id)
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=current_user,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentAddByFileApi(DatasetApiResource):
+    """Resource for documents."""
+    def post(self, tenant_id, dataset_id):
+        """Create document by upload file."""
+        args = {}
+        if 'data' in request.form:
+            args = json.loads(request.form['data'])
+        if 'doc_form' not in args:
+            args['doc_form'] = 'text_model'
+        if 'doc_language' not in args:
+            args['doc_language'] = 'English'
+        # get dataset info
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+        if not dataset.indexing_technique and not args['indexing_technique']:
+            raise ValueError('indexing_technique is required.')
+
+        # save file info
+        file = request.files['file']
+        # check file
+        if 'file' not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        upload_file = FileService.upload_file(file)
+        data_source = {
+            'type': 'upload_file',
+            'info_list': {
+                'file_info_list': {
+                    'file_ids': [upload_file.id]
+                }
+            }
+        }
+        args['data_source'] = data_source
+        # validate args
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=dataset.created_by_account,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentUpdateByFileApi(DatasetApiResource):
+    """Resource for update documents."""
+
+    def post(self, tenant_id, dataset_id, document_id):
+        """Update document by upload file."""
+        args = {}
+        if 'data' in request.form:
+            args = json.loads(request.form['data'])
+        if 'doc_form' not in args:
+            args['doc_form'] = 'text_model'
+        if 'doc_language' not in args:
+            args['doc_language'] = 'English'
+
+        # get dataset info
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+        if 'file' in request.files:
+            # save file info
+            file = request.files['file']
+
+
+            if len(request.files) > 1:
+                raise TooManyFilesError()
+
+            upload_file = FileService.upload_file(file)
+            data_source = {
+                'type': 'upload_file',
+                'info_list': {
+                    'file_info_list': {
+                        'file_ids': [upload_file.id]
+                    }
+                }
+            }
+            args['data_source'] = data_source
+        # validate args
+        args['original_document_id'] = str(document_id)
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=dataset.created_by_account,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentDeleteApi(DatasetApiResource):
+    def delete(self, tenant_id, dataset_id, document_id):
         """Delete document."""
         document_id = str(document_id)
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+
+        # get dataset info
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
 
         document = DocumentService.get_document(dataset.id, document_id)
 
@@ -126,8 +299,85 @@ class DocumentApi(DatasetApiResource):
         except services.errors.document.DocumentIndexingError:
             raise DocumentIndexingError('Cannot delete document during indexing.')
 
-        return {'result': 'success'}, 204
+        return {'result': 'success'}, 200
 
 
-api.add_resource(DocumentListApi, '/documents')
+class DocumentListApi(DatasetApiResource):
-api.add_resource(DocumentApi, '/documents/<uuid:document_id>')
+    def get(self, tenant_id, dataset_id):
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        page = request.args.get('page', default=1, type=int)
+        limit = request.args.get('limit', default=20, type=int)
+        search = request.args.get('keyword', default=None, type=str)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+
+        query = Document.query.filter_by(
+            dataset_id=str(dataset_id), tenant_id=tenant_id)
+
+        if search:
+            search = f'%{search}%'
+            query = query.filter(Document.name.like(search))
+
+        query = query.order_by(desc(Document.created_at))
+
+        paginated_documents = query.paginate(
+            page=page, per_page=limit, max_per_page=100, error_out=False)
+        documents = paginated_documents.items
+
+        response = {
+            'data': marshal(documents, document_fields),
+            'has_more': len(documents) == limit,
+            'limit': limit,
+            'total': paginated_documents.total,
+            'page': page
+        }
+
+        return response
+
+
+class DocumentIndexingStatusApi(DatasetApiResource):
+    def get(self, tenant_id, dataset_id, batch):
+        dataset_id = str(dataset_id)
+        batch = str(batch)
+        tenant_id = str(tenant_id)
+        # get dataset
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # get documents
+        documents = DocumentService.get_batch_documents(dataset_id, batch)
+        if not documents:
+            raise NotFound('Documents not found.')
+        documents_status = []
+        for document in documents:
+            completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+                                                              DocumentSegment.document_id == str(document.id),
+                                                              DocumentSegment.status != 're_segment').count()
+            total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
+                                                          DocumentSegment.status != 're_segment').count()
+            document.completed_segments = completed_segments
+            document.total_segments = total_segments
+            if document.is_paused:
+                document.indexing_status = 'paused'
+            documents_status.append(marshal(document, document_status_fields))
+        data = {
+            'data': documents_status
+        }
+        return data
+
+
+api.add_resource(DocumentAddByTextApi, '/datasets/<uuid:dataset_id>/document/create_by_text')
+api.add_resource(DocumentAddByFileApi, '/datasets/<uuid:dataset_id>/document/create_by_file')
+api.add_resource(DocumentUpdateByTextApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text')
+api.add_resource(DocumentUpdateByFileApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_file')
+api.add_resource(DocumentDeleteApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
+api.add_resource(DocumentListApi, '/datasets/<uuid:dataset_id>/documents')
+api.add_resource(DocumentIndexingStatusApi, '/datasets/<uuid:dataset_id>/documents/<string:batch>/indexing-status')

api/controllers/service_api/dataset/error.py

@@ -1,20 +1,73 @@
-# -*- coding:utf-8 -*-
 from libs.exception import BaseHTTPException
 
 
+class NoFileUploadedError(BaseHTTPException):
+    error_code = 'no_file_uploaded'
+    description = "Please upload your file."
+    code = 400
+
+
+class TooManyFilesError(BaseHTTPException):
+    error_code = 'too_many_files'
+    description = "Only one file is allowed."
+    code = 400
+
+
+class FileTooLargeError(BaseHTTPException):
+    error_code = 'file_too_large'
+    description = "File size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedFileTypeError(BaseHTTPException):
+    error_code = 'unsupported_file_type'
+    description = "File type not allowed."
+    code = 415
+
+
+class HighQualityDatasetOnlyError(BaseHTTPException):
+    error_code = 'high_quality_dataset_only'
+    description = "Current operation only supports 'high-quality' datasets."
+    code = 400
+
+
+class DatasetNotInitializedError(BaseHTTPException):
+    error_code = 'dataset_not_initialized'
+    description = "The dataset is still being initialized or indexing. Please wait a moment."
+    code = 400
+
+
 class ArchivedDocumentImmutableError(BaseHTTPException):
     error_code = 'archived_document_immutable'
-    description = "Cannot operate when document was archived."
+    description = "The archived document is not editable."
     code = 403
 
 
+class DatasetNameDuplicateError(BaseHTTPException):
+    error_code = 'dataset_name_duplicate'
+    description = "The dataset name already exists. Please modify your dataset name."
+    code = 409
+
+
+class InvalidActionError(BaseHTTPException):
+    error_code = 'invalid_action'
+    description = "Invalid action."
+    code = 400
+
+
+class DocumentAlreadyFinishedError(BaseHTTPException):
+    error_code = 'document_already_finished'
+    description = "The document has been processed. Please refresh the page or go to the document details."
+    code = 400
+
+
 class DocumentIndexingError(BaseHTTPException):
     error_code = 'document_indexing'
-    description = "Cannot operate document during indexing."
+    description = "The document is being processed and cannot be edited."
-    code = 403
+    code = 400
 
 
-class DatasetNotInitedError(BaseHTTPException):
+class InvalidMetadataError(BaseHTTPException):
-    error_code = 'dataset_not_inited'
+    error_code = 'invalid_metadata'
-    description = "The dataset is still being initialized or indexing. Please wait a moment."
+    description = "The metadata content is incorrect. Please check and verify."
-    code = 403
+    code = 400

api/controllers/service_api/dataset/segment.py

@@ -0,0 +1,201 @@
+from flask_login import current_user
+from flask_restful import reqparse, marshal
+from werkzeug.exceptions import NotFound
+from controllers.service_api import api
+from controllers.service_api.app.error import ProviderNotInitializeError
+from controllers.service_api.wraps import DatasetApiResource
+from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
+from core.model_providers.model_factory import ModelFactory
+from extensions.ext_database import db
+from fields.segment_fields import segment_fields
+from models.dataset import Dataset, DocumentSegment
+from services.dataset_service import DatasetService, DocumentService, SegmentService
+
+
+class SegmentApi(DatasetApiResource):
+    """Resource for segments."""
+    def post(self, tenant_id, dataset_id, document_id):
+        """Create single segment."""
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # 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
+        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)
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('segments', type=list, required=False, nullable=True, location='json')
+        args = parser.parse_args()
+        for args_item in args['segments']:
+            SegmentService.segment_create_args_validate(args_item, document)
+        segments = SegmentService.multi_create_segment(args['segments'], document, dataset)
+        return {
+            'data': marshal(segments, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+    def get(self, tenant_id, dataset_id, document_id):
+        """Create single segment."""
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # 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
+        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)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('status', type=str,
+                            action='append', default=[], location='args')
+        parser.add_argument('keyword', type=str, default=None, location='args')
+        args = parser.parse_args()
+
+        status_list = args['status']
+        keyword = args['keyword']
+
+        query = DocumentSegment.query.filter(
+            DocumentSegment.document_id == str(document_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        )
+
+        if status_list:
+            query = query.filter(DocumentSegment.status.in_(status_list))
+
+        if keyword:
+            query = query.where(DocumentSegment.content.ilike(f'%{keyword}%'))
+
+        total = query.count()
+        segments = query.order_by(DocumentSegment.position).all()
+        return {
+            'data': marshal(segments, segment_fields),
+            'doc_form': document.doc_form,
+            'total': total
+        }, 200
+
+
+class DatasetSegmentApi(DatasetApiResource):
+    def delete(self, tenant_id, dataset_id, document_id, segment_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        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 segment
+        segment = DocumentSegment.query.filter(
+            DocumentSegment.id == str(segment_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        ).first()
+        if not segment:
+            raise NotFound('Segment not found.')
+        SegmentService.delete_segment(segment, document, dataset)
+        return {'result': 'success'}, 200
+
+    def post(self, tenant_id, dataset_id, document_id, segment_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        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.')
+        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),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        ).first()
+        if not segment:
+            raise NotFound('Segment not found.')
+
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('segments', type=dict, required=False, nullable=True, location='json')
+        args = parser.parse_args()
+
+        SegmentService.segment_create_args_validate(args['segments'], document)
+        segment = SegmentService.update_segment(args['segments'], segment, document, dataset)
+        return {
+            'data': marshal(segment, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+
+api.add_resource(SegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments')
+api.add_resource(DatasetSegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>')

api/controllers/service_api/wraps.py

@@ -2,12 +2,14 @@
 from datetime import datetime
 from functools import wraps
 
-from flask import request
+from flask import request, current_app
+from flask_login import user_logged_in
 from flask_restful import Resource
 from werkzeug.exceptions import NotFound, Unauthorized
 
+from libs.login import _get_user
 from extensions.ext_database import db
-from models.dataset import Dataset
+from models.account import Tenant, TenantAccountJoin, Account
 from models.model import ApiToken, App
 
 
@@ -17,7 +19,7 @@ def validate_app_token(view=None):
         def decorated(*args, **kwargs):
             api_token = validate_and_get_api_token('app')
 
-            app_model = db.session.query(App).get(api_token.app_id)
+            app_model = db.session.query(App).filter(App.id == api_token.app_id).first()
             if not app_model:
                 raise NotFound()
 
@@ -43,12 +45,24 @@ def validate_dataset_token(view=None):
         @wraps(view)
         def decorated(*args, **kwargs):
             api_token = validate_and_get_api_token('dataset')
-
+            tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \
-            dataset = db.session.query(Dataset).get(api_token.dataset_id)
+                .filter(Tenant.id == api_token.tenant_id) \
-            if not dataset:
+                .filter(TenantAccountJoin.tenant_id == Tenant.id) \
-                raise NotFound()
+                .filter(TenantAccountJoin.role == 'owner') \
-
+                .one_or_none()
-            return view(dataset, *args, **kwargs)
+            if tenant_account_join:
+                tenant, ta = tenant_account_join
+                account = Account.query.filter_by(id=ta.account_id).first()
+                # Login admin
+                if account:
+                    account.current_tenant = tenant
+                    current_app.login_manager._update_request_context_with_user(account)
+                    user_logged_in.send(current_app._get_current_object(), user=_get_user())
+                else:
+                    raise Unauthorized("Tenant owner account is not exist.")
+            else:
+                raise Unauthorized("Tenant is not exist.")
+            return view(api_token.tenant_id, *args, **kwargs)
         return decorated
 
     if view:
@@ -64,14 +78,14 @@ def validate_and_get_api_token(scope=None):
     Validate and get API token.
     """
     auth_header = request.headers.get('Authorization')
-    if auth_header is None:
+    if auth_header is None or ' ' not in auth_header:
-        raise Unauthorized()
+        raise Unauthorized("Authorization header must be provided and start with 'Bearer'")
 
     auth_scheme, auth_token = auth_header.split(None, 1)
     auth_scheme = auth_scheme.lower()
 
     if auth_scheme != 'bearer':
-        raise Unauthorized()
+        raise Unauthorized("Authorization scheme must be 'Bearer'")
 
     api_token = db.session.query(ApiToken).filter(
         ApiToken.token == auth_token,
@@ -79,7 +93,7 @@ def validate_and_get_api_token(scope=None):
     ).first()
 
     if not api_token:
-        raise Unauthorized()
+        raise Unauthorized("Access token is invalid")
 
     api_token.last_used_at = datetime.utcnow()
     db.session.commit()

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/conversation.py

@@ -6,26 +6,12 @@ from werkzeug.exceptions import NotFound
 from controllers.web import api
 from controllers.web.error import NotChatAppError
 from controllers.web.wraps import WebApiResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationListApi(WebApiResource):
 
@@ -73,7 +59,7 @@ class ConversationApi(WebApiResource):
 
 class ConversationRenameApi(WebApiResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()

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,14 +53,28 @@ 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:
             _, observation = intermediate_steps[-1]
             return AgentFinish(return_values={"output": observation}, log=observation)
 
-        return super().plan(intermediate_steps, callbacks, **kwargs)
+        try:
+            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
 
     async def aplan(
             self,

api/core/agent/agent/openai_function_call.py

@@ -45,14 +45,18 @@ 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()
 
-        predicted_message = self.llm.predict_messages(
+        try:
-            messages, functions=self.functions, callbacks=None
+            predicted_message = self.llm.predict_messages(
-        )
+                messages, functions=self.functions, callbacks=None
+            )
+        except Exception as e:
+            new_exception = self.model_instance.handle_exceptions(e)
+            raise new_exception
 
         function_call = predicted_message.additional_kwargs.get("function_call", {})
 
@@ -93,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/openai_function_call_summarize_mixin.py

@@ -14,7 +14,7 @@ from core.model_providers.models.llm.base import BaseLLM
 class OpenAIFunctionCallSummarizeMixin(BaseModel, CalcTokenMixin):
     moving_summary_buffer: str = ""
     moving_summary_index: int = 0
-    summary_llm: BaseLanguageModel
+    summary_llm: BaseLanguageModel = None
     model_instance: BaseLLM
 
     class Config:
@@ -66,12 +66,12 @@ class OpenAIFunctionCallSummarizeMixin(BaseModel, CalcTokenMixin):
 
         return new_messages
 
-    def get_num_tokens_from_messages(self, llm: BaseLanguageModel, messages: List[BaseMessage], **kwargs) -> int:
+    def get_num_tokens_from_messages(self, model_instance: BaseLLM, messages: List[BaseMessage], **kwargs) -> int:
         """Calculate num tokens for gpt-3.5-turbo and gpt-4 with tiktoken package.
 
         Official documentation: https://github.com/openai/openai-cookbook/blob/
         main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb"""
-        llm = cast(ChatOpenAI, llm)
+        llm = cast(ChatOpenAI, model_instance.client)
         model, encoding = llm._get_encoding_model()
         if model.startswith("gpt-3.5-turbo"):
             # every message follows <im_start>{role/name}\n{content}<im_end>\n

api/core/agent/agent/openai_multi_function_call.py

@@ -50,9 +50,13 @@ class AutoSummarizingOpenMultiAIFunctionCallAgent(OpenAIMultiFunctionsAgent, Ope
         prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
         messages = prompt.to_messages()
 
-        predicted_message = self.llm.predict_messages(
+        try:
-            messages, functions=self.functions, callbacks=None
+            predicted_message = self.llm.predict_messages(
-        )
+                messages, functions=self.functions, callbacks=None
+            )
+        except Exception as e:
+            new_exception = self.model_instance.handle_exceptions(e)
+            raise new_exception
 
         function_call = predicted_message.additional_kwargs.get("function_call", {})
 

api/core/agent/agent/output_parser/structured_chat.py

@@ -10,7 +10,7 @@ from langchain.schema import AgentAction, AgentFinish, OutputParserException
 class StructuredChatOutputParser(LCStructuredChatOutputParser):
     def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
         try:
-            action_match = re.search(r"```(.*?)\n(.*?)```?", text, re.DOTALL)
+            action_match = re.search(r"```(\w*)\n?({.*?)```", text, re.DOTALL)
             if action_match is not None:
                 response = json.loads(action_match.group(2).strip(), strict=False)
                 if isinstance(response, list):
@@ -26,4 +26,4 @@ class StructuredChatOutputParser(LCStructuredChatOutputParser):
             else:
                 return AgentFinish({"output": text}, text)
         except Exception as e:
-            raise OutputParserException(f"Could not parse LLM output: {text}") from e
+            raise OutputParserException(f"Could not parse LLM output: {text}")

api/core/agent/agent/structed_multi_dataset_router_agent.py

@@ -90,14 +90,25 @@ 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)
-        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
 
         try:
-            return self.output_parser.parse(full_output)
+            full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
+        except Exception as e:
+            new_exception = self.model_instance.handle_exceptions(e)
+            raise new_exception
+
+        try:
+            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

@@ -52,7 +52,7 @@ Action:
 class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
     moving_summary_buffer: str = ""
     moving_summary_index: int = 0
-    summary_llm: BaseLanguageModel
+    summary_llm: BaseLanguageModel = None
     model_instance: BaseLLM
 
     class Config:
@@ -89,8 +89,8 @@ class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
             Action specifying what tool to use.
         """
         full_inputs = self.get_full_inputs(intermediate_steps, **kwargs)
-
         prompts, _ = self.llm_chain.prep_prompts(input_list=[self.llm_chain.prep_inputs(full_inputs)])
+
         messages = []
         if prompts:
             messages = prompts[0].to_messages()
@@ -99,16 +99,26 @@ class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
         if rest_tokens < 0:
             full_inputs = self.summarize_messages(intermediate_steps, **kwargs)
 
-        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
+        try:
+            full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
+        except Exception as e:
+            new_exception = self.model_instance.handle_exceptions(e)
+            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."}, "")
 
     def summarize_messages(self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs):
-        if len(intermediate_steps) >= 2:
+        if len(intermediate_steps) >= 2 and self.summary_llm:
             should_summary_intermediate_steps = intermediate_steps[self.moving_summary_index:-1]
             should_summary_messages = [AIMessage(content=observation)
                                        for _, observation in should_summary_intermediate_steps]

api/core/agent/agent_executor.py

@@ -16,6 +16,8 @@ from core.agent.agent.structed_multi_dataset_router_agent import StructuredMulti
 from core.agent.agent.structured_chat import AutoSummarizingStructuredChatAgent
 from langchain.agents import AgentExecutor as LCAgentExecutor
 
+from core.helper import moderation
+from core.model_providers.error import LLMError
 from core.model_providers.models.llm.base import BaseLLM
 from core.tool.dataset_retriever_tool import DatasetRetrieverTool
 
@@ -32,7 +34,7 @@ class AgentConfiguration(BaseModel):
     strategy: PlanningStrategy
     model_instance: BaseLLM
     tools: list[BaseTool]
-    summary_model_instance: BaseLLM
+    summary_model_instance: BaseLLM = None
     memory: Optional[BaseChatMemory] = None
     callbacks: Callbacks = None
     max_iterations: int = 6
@@ -65,7 +67,8 @@ class AgentExecutor:
                 llm=self.configuration.model_instance.client,
                 tools=self.configuration.tools,
                 output_parser=StructuredChatOutputParser(),
-                summary_llm=self.configuration.summary_model_instance.client,
+                summary_llm=self.configuration.summary_model_instance.client
+                if self.configuration.summary_model_instance else None,
                 verbose=True
             )
         elif self.configuration.strategy == PlanningStrategy.FUNCTION_CALL:
@@ -74,7 +77,8 @@ class AgentExecutor:
                 llm=self.configuration.model_instance.client,
                 tools=self.configuration.tools,
                 extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,  # used for read chat histories memory
-                summary_llm=self.configuration.summary_model_instance.client,
+                summary_llm=self.configuration.summary_model_instance.client
+                if self.configuration.summary_model_instance else None,
                 verbose=True
             )
         elif self.configuration.strategy == PlanningStrategy.MULTI_FUNCTION_CALL:
@@ -83,7 +87,8 @@ class AgentExecutor:
                 llm=self.configuration.model_instance.client,
                 tools=self.configuration.tools,
                 extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,  # used for read chat histories memory
-                summary_llm=self.configuration.summary_model_instance.client,
+                summary_llm=self.configuration.summary_model_instance.client
+                if self.configuration.summary_model_instance else None,
                 verbose=True
             )
         elif self.configuration.strategy == PlanningStrategy.ROUTER:
@@ -113,6 +118,18 @@ class AgentExecutor:
         return self.agent.should_use_agent(query)
 
     def run(self, query: str) -> AgentExecuteResult:
+        moderation_result = moderation.check_moderation(
+            self.configuration.model_instance.model_provider,
+            query
+        )
+
+        if not moderation_result:
+            return AgentExecuteResult(
+                output="I apologize for any confusion, but I'm an AI assistant to be helpful, harmless, and honest.",
+                strategy=self.configuration.strategy,
+                configuration=self.configuration
+            )
+
         agent_executor = LCAgentExecutor.from_agent_and_tools(
             agent=self.agent,
             tools=self.configuration.tools,
@@ -125,7 +142,9 @@ class AgentExecutor:
 
         try:
             output = agent_executor.run(query)
-        except Exception:
+        except LLMError as ex:
+            raise ex
+        except Exception as ex:
             logging.exception("agent_executor run failed")
             output = None
 

api/core/callback_handler/agent_loop_gather_callback_handler.py

@@ -6,10 +6,11 @@ from typing import Any, Dict, List, Union, Optional
 
 from langchain.agents import openai_functions_agent, openai_functions_multi_agent
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult, ChatGeneration
+from langchain.schema import AgentAction, AgentFinish, LLMResult, ChatGeneration, BaseMessage
 
 from core.callback_handler.entity.agent_loop import AgentLoop
 from core.conversation_message_task import ConversationMessageTask
+from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.llm.base import BaseLLM
 
 
@@ -17,9 +18,9 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
     """Callback Handler that prints to std out."""
     raise_error: bool = True
 
-    def __init__(self, model_instant: BaseLLM, conversation_message_task: ConversationMessageTask) -> None:
+    def __init__(self, model_instance: BaseLLM, conversation_message_task: ConversationMessageTask) -> None:
         """Initialize callback handler."""
-        self.model_instant = model_instant
+        self.model_instance = model_instance
         self.conversation_message_task = conversation_message_task
         self._agent_loops = []
         self._current_loop = None
@@ -45,6 +46,21 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
         """Whether to ignore chain callbacks."""
         return True
 
+    def on_chat_model_start(
+            self,
+            serialized: Dict[str, Any],
+            messages: List[List[BaseMessage]],
+            **kwargs: Any
+    ) -> Any:
+        if not self._current_loop:
+            # Agent start with a LLM query
+            self._current_loop = AgentLoop(
+                position=len(self._agent_loops) + 1,
+                prompt="\n".join([message.content for message in messages[0]]),
+                status='llm_started',
+                started_at=time.perf_counter()
+            )
+
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
@@ -68,6 +84,10 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             self._current_loop.status = 'llm_end'
             if response.llm_output:
                 self._current_loop.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
+            else:
+                self._current_loop.prompt_tokens = self.model_instance.get_num_tokens(
+                    [PromptMessage(content=self._current_loop.prompt)]
+                )
             completion_generation = response.generations[0][0]
             if isinstance(completion_generation, ChatGeneration):
                 completion_message = completion_generation.message
@@ -81,11 +101,15 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
 
             if response.llm_output:
                 self._current_loop.completion_tokens = response.llm_output['token_usage']['completion_tokens']
+            else:
+                self._current_loop.completion_tokens = self.model_instance.get_num_tokens(
+                    [PromptMessage(content=self._current_loop.completion)]
+                )
 
     def on_llm_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
-        logging.exception(error)
+        logging.debug("Agent on_llm_error: %s", error)
         self._agent_loops = []
         self._current_loop = None
         self._message_agent_thought = None
@@ -153,7 +177,7 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             self._current_loop.latency = self._current_loop.completed_at - self._current_loop.started_at
 
             self.conversation_message_task.on_agent_end(
-                self._message_agent_thought, self.model_instant, self._current_loop
+                self._message_agent_thought, self.model_instance, self._current_loop
             )
 
             self._agent_loops.append(self._current_loop)
@@ -164,7 +188,7 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
         """Do nothing."""
-        logging.exception(error)
+        logging.debug("Agent on_tool_error: %s", error)
         self._agent_loops = []
         self._current_loop = None
         self._message_agent_thought = None
@@ -184,7 +208,7 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             )
 
             self.conversation_message_task.on_agent_end(
-                self._message_agent_thought, self.model_instant, self._current_loop
+                self._message_agent_thought, self.model_instance, self._current_loop
             )
 
             self._agent_loops.append(self._current_loop)

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')
+        tool_name: str = serialized.get('name')
-        input_dict = json.loads(input_str.replace("'", "\""))
+        dataset_id = tool_name.removeprefix('dataset-')
-        dataset_id = input_dict.get('dataset_id')
+
-        query = input_dict.get('query')
+        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,14 +64,11 @@ 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:
         """Do nothing."""
-        logging.exception(error)
+        logging.debug("Dataset tool on_llm_error: %s", error)

api/core/callback_handler/entity/llm_message.py

@@ -6,4 +6,3 @@ class LLMMessage(BaseModel):
     prompt_tokens: int = 0
     completion: str = ''
     completion_tokens: int = 0
-    latency: float = 0.0

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/callback_handler/llm_callback_handler.py

@@ -1,5 +1,4 @@
 import logging
-import time
 from typing import Any, Dict, List, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
@@ -32,7 +31,6 @@ class LLMCallbackHandler(BaseCallbackHandler):
             messages: List[List[BaseMessage]],
             **kwargs: Any
     ) -> Any:
-        self.start_at = time.perf_counter()
         real_prompts = []
         for message in messages[0]:
             if message.type == 'human':
@@ -53,8 +51,6 @@ class LLMCallbackHandler(BaseCallbackHandler):
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
     ) -> None:
-        self.start_at = time.perf_counter()
-
         self.llm_message.prompt = [{
             "role": 'user',
             "text": prompts[0]
@@ -63,14 +59,22 @@ class LLMCallbackHandler(BaseCallbackHandler):
         self.llm_message.prompt_tokens = self.model_instance.get_num_tokens([PromptMessage(content=prompts[0])])
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        end_at = time.perf_counter()
-        self.llm_message.latency = end_at - self.start_at
-
         if not self.conversation_message_task.streaming:
             self.conversation_message_task.append_message_text(response.generations[0][0].text)
             self.llm_message.completion = response.generations[0][0].text
 
-        self.llm_message.completion_tokens = self.model_instance.get_num_tokens([PromptMessage(content=self.llm_message.completion)])
+        if response.llm_output and 'token_usage' in response.llm_output:
+            if 'prompt_tokens' in response.llm_output['token_usage']:
+                self.llm_message.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
+
+            if 'completion_tokens' in response.llm_output['token_usage']:
+                self.llm_message.completion_tokens = response.llm_output['token_usage']['completion_tokens']
+            else:
+                self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                    [PromptMessage(content=self.llm_message.completion)])
+        else:
+            self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                [PromptMessage(content=self.llm_message.completion)])
 
         self.conversation_message_task.save_message(self.llm_message)
 
@@ -89,8 +93,6 @@ class LLMCallbackHandler(BaseCallbackHandler):
         """Do nothing."""
         if isinstance(error, ConversationTaskStoppedException):
             if self.conversation_message_task.streaming:
-                end_at = time.perf_counter()
-                self.llm_message.latency = end_at - self.start_at
                 self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
                     [PromptMessage(content=self.llm_message.completion)]
                 )

api/core/callback_handler/main_chain_gather_callback_handler.py

@@ -72,5 +72,5 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
     def on_chain_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
-        logging.exception(error)
+        logging.debug("Dataset tool on_chain_error: %s", error)
         self.clear_chain_results()

api/core/chain/sensitive_word_avoidance_chain.py

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

api/core/completion.py

@@ -1,31 +1,32 @@
+import json
 import logging
-import re
+from typing import Optional, List, Union
-from typing import Optional, List, Union, Tuple
 
-from langchain.schema import BaseMessage
 from requests.exceptions import ChunkedEncodingError
 
 from core.agent.agent_executor import AgentExecuteResult, PlanningStrategy
 from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
 from core.callback_handler.llm_callback_handler import LLMCallbackHandler
+from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceError
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
 from core.model_providers.error import LLMBadRequestError
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
     ReadOnlyConversationTokenDBBufferSharedMemory
 from core.model_providers.model_factory import ModelFactory
-from core.model_providers.models.entity.message import PromptMessage, to_prompt_messages
+from core.model_providers.models.entity.message import PromptMessage
 from core.model_providers.models.llm.base import BaseLLM
 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
         """
@@ -76,29 +77,55 @@ class Completion:
             app_model_config=app_model_config
         )
 
-        # parse sensitive_word_avoidance_chain
-        chain_callback = MainChainGatherCallbackHandler(conversation_message_task)
-        sensitive_word_avoidance_chain = orchestrator_rule_parser.to_sensitive_word_avoidance_chain([chain_callback])
-        if sensitive_word_avoidance_chain:
-            query = sensitive_word_avoidance_chain.run(query)
-
-        # get agent executor
-        agent_executor = orchestrator_rule_parser.to_agent_executor(
-            conversation_message_task=conversation_message_task,
-            memory=memory,
-            rest_tokens=rest_tokens_for_context_and_memory,
-            chain_callback=chain_callback
-        )
-
-        # run agent executor
-        agent_execute_result = None
-        if agent_executor:
-            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:
+            # parse sensitive_word_avoidance_chain
+            chain_callback = MainChainGatherCallbackHandler(conversation_message_task)
+            sensitive_word_avoidance_chain = orchestrator_rule_parser.to_sensitive_word_avoidance_chain(
+                final_model_instance, [chain_callback])
+            if sensitive_word_avoidance_chain:
+                try:
+                    query = sensitive_word_avoidance_chain.run(query)
+                except SensitiveWordAvoidanceError as ex:
+                    cls.run_final_llm(
+                        model_instance=final_model_instance,
+                        mode=app.mode,
+                        app_model_config=app_model_config,
+                        query=query,
+                        inputs=inputs,
+                        agent_execute_result=None,
+                        conversation_message_task=conversation_message_task,
+                        memory=memory,
+                        fake_response=ex.message
+                    )
+                    return
+
+            # get agent executor
+            agent_executor = orchestrator_rule_parser.to_agent_executor(
+                conversation_message_task=conversation_message_task,
+                memory=memory,
+                rest_tokens=rest_tokens_for_context_and_memory,
+                chain_callback=chain_callback,
+                retriever_from=retriever_from
+            )
+
+            query_for_agent = cls.get_query_for_agent(app, app_model_config, query, inputs)
+
+            # run agent executor
+            agent_execute_result = None
+            if query_for_agent and agent_executor:
+                should_use_agent = agent_executor.should_use_agent(query_for_agent)
+                if should_use_agent:
+                    agent_execute_result = agent_executor.run(query_for_agent)
+
+            # When no extra pre prompt is specified,
+            # the output of the agent can be used directly as the main output content without calling LLM again
+            fake_response = None
+            if not app_model_config.pre_prompt and agent_execute_result and agent_execute_result.output \
+                    and agent_execute_result.strategy not in [PlanningStrategy.ROUTER,
+                                                              PlanningStrategy.REACT_ROUTER]:
+                fake_response = agent_execute_result.output
+
+            # run the final llm
             cls.run_final_llm(
                 model_instance=final_model_instance,
                 mode=app.mode,
@@ -107,7 +134,8 @@ class Completion:
                 inputs=inputs,
                 agent_execute_result=agent_execute_result,
                 conversation_message_task=conversation_message_task,
-                memory=memory
+                memory=memory,
+                fake_response=fake_response
             )
         except ConversationTaskStoppedException:
             return
@@ -116,27 +144,28 @@ class Completion:
             logging.warning(f'ChunkedEncodingError: {e}')
             conversation_message_task.end()
             return
+        
+    @classmethod
+    def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str:
+        if app.mode != 'completion':
+            return query
+        
+        return inputs.get(app_model_config.dataset_query_variable, "")
 
     @classmethod
-    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]):
+                      memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory],
-        # When no extra pre prompt is specified,
+                      fake_response: Optional[str]):
-        # the output of the agent can be used directly as the main output content without calling LLM again
-        fake_response = None
-        if not app_model_config.pre_prompt and agent_execute_result and agent_execute_result.output \
-                and agent_execute_result.strategy != PlanningStrategy.ROUTER:
-            fake_response = agent_execute_result.output
-
         # get llm prompt
-        prompt_messages, stop_words = cls.get_main_llm_prompt(
+        prompt_messages, stop_words = model_instance.get_prompt(
             mode=mode,
-            model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
-            query=query,
             inputs=inputs,
-            agent_execute_result=agent_execute_result,
+            query=query,
+            context=agent_execute_result.output if agent_execute_result else None,
             memory=memory
         )
 
@@ -151,116 +180,8 @@ class Completion:
             callbacks=[LLMCallbackHandler(model_instance, conversation_message_task)],
             fake_response=fake_response
         )
-
         return response
 
-    @classmethod
-    def get_main_llm_prompt(cls, mode: str, model: dict,
-                            pre_prompt: str, query: str, inputs: dict,
-                            agent_execute_result: Optional[AgentExecuteResult],
-                            memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]) -> \
-            Tuple[List[PromptMessage], Optional[List[str]]]:
-        if mode == 'completion':
-            prompt_template = JinjaPromptTemplate.from_template(
-                template=("""Use the following context as your learned knowledge, inside <context></context> XML tags.
-
-<context>
-{{context}}
-</context>
-
-When answer to user:
-- If you don't know, just say that you don't know.
-- If you don't know when you are not sure, ask for clarification. 
-Avoid mentioning that you obtained the information from the context.
-And answer according to the language of the user's question.
-""" if agent_execute_result else "")
-                         + (pre_prompt + "\n" if pre_prompt else "")
-                         + "{{query}}\n"
-            )
-
-            if agent_execute_result:
-                inputs['context'] = agent_execute_result.output
-
-            prompt_inputs = {k: inputs[k] for k in prompt_template.input_variables if k in inputs}
-            prompt_content = prompt_template.format(
-                query=query,
-                **prompt_inputs
-            )
-
-            return [PromptMessage(content=prompt_content)], None
-        else:
-            messages: List[BaseMessage] = []
-
-            human_inputs = {
-                "query": query
-            }
-
-            human_message_prompt = ""
-
-            if pre_prompt:
-                pre_prompt_inputs = {k: inputs[k] for k in
-                                     JinjaPromptTemplate.from_template(template=pre_prompt).input_variables
-                                     if k in inputs}
-
-                if pre_prompt_inputs:
-                    human_inputs.update(pre_prompt_inputs)
-
-            if agent_execute_result:
-                human_inputs['context'] = agent_execute_result.output
-                human_message_prompt += """Use the following context as your learned knowledge, inside <context></context> XML tags.
-
-<context>
-{{context}}
-</context>
-
-When answer to user:
-- If you don't know, just say that you don't know.
-- If you don't know when you are not sure, ask for clarification. 
-Avoid mentioning that you obtained the information from the context.
-And answer according to the language of the user's question.
-"""
-
-            if pre_prompt:
-                human_message_prompt += pre_prompt
-
-            query_prompt = "\n\nHuman: {{query}}\n\nAssistant: "
-
-            if memory:
-                # append chat histories
-                tmp_human_message = PromptBuilder.to_human_message(
-                    prompt_content=human_message_prompt + query_prompt,
-                    inputs=human_inputs
-                )
-
-                if memory.model_instance.model_rules.max_tokens.max:
-                    curr_message_tokens = memory.model_instance.get_num_tokens(to_prompt_messages([tmp_human_message]))
-                    max_tokens = model.get("completion_params").get('max_tokens')
-                    rest_tokens = memory.model_instance.model_rules.max_tokens.max - max_tokens - curr_message_tokens
-                    rest_tokens = max(rest_tokens, 0)
-                else:
-                    rest_tokens = 2000
-
-                histories = cls.get_history_messages_from_memory(memory, rest_tokens)
-                human_message_prompt += "\n\n" if human_message_prompt else ""
-                human_message_prompt += "Here is the chat histories between human and assistant, " \
-                                        "inside <histories></histories> XML tags.\n\n<histories>\n"
-                human_message_prompt += histories + "\n</histories>"
-
-            human_message_prompt += query_prompt
-
-            # construct main prompt
-            human_message = PromptBuilder.to_human_message(
-                prompt_content=human_message_prompt,
-                inputs=human_inputs
-            )
-
-            messages.append(human_message)
-
-            for message in messages:
-                message.content = re.sub(r'<\|.*?\|>', '', message.content)
-
-            return to_prompt_messages(messages), ['\nHuman:', '</histories>']
-
     @classmethod
     def get_history_messages_from_memory(cls, memory: ReadOnlyConversationTokenDBBufferSharedMemory,
                                          max_token_limit: int) -> str:
@@ -307,13 +228,12 @@ And answer according to the language of the user's question.
             max_tokens = 0
 
         # get prompt without memory and context
-        prompt_messages, _ = cls.get_main_llm_prompt(
+        prompt_messages, _ = model_instance.get_prompt(
             mode=mode,
-            model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
-            query=query,
             inputs=inputs,
-            agent_execute_result=None,
+            query=query,
+            context=None,
             memory=None
         )
 
@@ -358,13 +278,12 @@ And answer according to the language of the user's question.
         )
 
         # get llm prompt
-        old_prompt_messages, _ = cls.get_main_llm_prompt(
+        old_prompt_messages, _ = final_model_instance.get_prompt(
-            mode="completion",
+            mode='completion',
-            model=app_model_config.model_dict,
             pre_prompt=pre_prompt,
-            query=message.query,
             inputs=message.inputs,
-            agent_execute_result=None,
+            query=message.query,
+            context=None,
             memory=None
         )
 

api/core/conversation_message_task.py

@@ -1,6 +1,6 @@
-import decimal
 import json
-from typing import Optional, Union
+import time
+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,13 +15,16 @@ 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:
     def __init__(self, task_id: str, app: App, app_model_config: AppModelConfig, user: Account,
                  inputs: dict, query: str, streaming: bool, model_instance: BaseLLM,
                  conversation: Optional[Conversation] = None, is_override: bool = False):
+        self.start_at = time.perf_counter()
+
         self.task_id = task_id
 
         self.app = app
@@ -41,6 +44,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')
@@ -58,19 +63,10 @@ class ConversationMessageTask:
         )
 
     def init(self):
+
         override_model_configs = None
         if self.is_override:
-            override_model_configs = {
+            override_model_configs = self.app_model_config.to_dict()
-                "model": self.app_model_config.model_dict,
-                "pre_prompt": self.app_model_config.pre_prompt,
-                "agent_mode": self.app_model_config.agent_mode_dict,
-                "opening_statement": self.app_model_config.opening_statement,
-                "suggested_questions": self.app_model_config.suggested_questions_list,
-                "suggested_questions_after_answer": self.app_model_config.suggested_questions_after_answer_dict,
-                "more_like_this": self.app_model_config.more_like_this_dict,
-                "sensitive_word_avoidance": self.app_model_config.sensitive_word_avoidance_dict,
-                "user_input_form": self.app_model_config.user_input_form_list,
-            }
 
         introduction = ''
         system_instruction = ''
@@ -98,7 +94,7 @@ class ConversationMessageTask:
         if not self.conversation:
             self.is_new_conversation = True
             self.conversation = Conversation(
-                app_id=self.app_model_config.app_id,
+                app_id=self.app.id,
                 app_model_config_id=self.app_model_config.id,
                 model_provider=self.provider_name,
                 model_id=self.model_name,
@@ -116,10 +112,10 @@ class ConversationMessageTask:
             )
 
             db.session.add(self.conversation)
-            db.session.flush()
+            db.session.commit()
 
         self.message = Message(
-            app_id=self.app_model_config.app_id,
+            app_id=self.app.id,
             model_provider=self.provider_name,
             model_id=self.model_name,
             override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
@@ -129,9 +125,11 @@ class ConversationMessageTask:
             message="",
             message_tokens=0,
             message_unit_price=0,
+            message_price_unit=0,
             answer="",
             answer_tokens=0,
             answer_unit_price=0,
+            answer_price_unit=0,
             provider_response_latency=0,
             total_price=0,
             currency=self.model_instance.get_currency(),
@@ -142,26 +140,35 @@ class ConversationMessageTask:
         )
 
         db.session.add(self.message)
-        db.session.flush()
+        db.session.commit()
 
     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
         answer_tokens = llm_message.completion_tokens
-        message_unit_price = self.model_instance.get_token_price(1, MessageType.HUMAN)
-        answer_unit_price = self.model_instance.get_token_price(1, MessageType.ASSISTANT)
 
-        total_price = self.calc_total_price(message_tokens, message_unit_price, answer_tokens, answer_unit_price)
+        message_unit_price = self.model_instance.get_tokens_unit_price(MessageType.HUMAN)
+        message_price_unit = self.model_instance.get_price_unit(MessageType.HUMAN)
+        answer_unit_price = self.model_instance.get_tokens_unit_price(MessageType.ASSISTANT)
+        answer_price_unit = self.model_instance.get_price_unit(MessageType.ASSISTANT)
+
+        message_total_price = self.model_instance.calc_tokens_price(message_tokens, MessageType.HUMAN)
+        answer_total_price = self.model_instance.calc_tokens_price(answer_tokens, MessageType.ASSISTANT)
+        total_price = message_total_price + answer_total_price
 
         self.message.message = llm_message.prompt
         self.message.message_tokens = message_tokens
         self.message.message_unit_price = message_unit_price
-        self.message.answer = PromptBuilder.process_template(llm_message.completion.strip()) if llm_message.completion else ''
+        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_tokens = answer_tokens
         self.message.answer_unit_price = answer_unit_price
-        self.message.provider_response_latency = llm_message.latency
+        self.message.answer_price_unit = answer_price_unit
+        self.message.provider_response_latency = time.perf_counter() - self.start_at
         self.message.total_price = total_price
 
         db.session.commit()
@@ -184,12 +191,13 @@ class ConversationMessageTask:
         )
 
         db.session.add(message_chain)
-        db.session.flush()
+        db.session.commit()
 
         return message_chain
 
     def on_chain_end(self, message_chain: MessageChain, chain_result: ChainResult):
         message_chain.output = json.dumps(chain_result.completion)
+        db.session.commit()
 
         self._pub_handler.pub_chain(message_chain)
 
@@ -202,44 +210,47 @@ class ConversationMessageTask:
             tool=agent_loop.tool_name,
             tool_input=agent_loop.tool_input,
             message=agent_loop.prompt,
+            message_price_unit=0,
             answer=agent_loop.completion,
+            answer_price_unit=0,
             created_by_role=('account' if isinstance(self.user, Account) else 'end_user'),
             created_by=self.user.id
         )
 
         db.session.add(message_agent_thought)
-        db.session.flush()
+        db.session.commit()
 
         self._pub_handler.pub_agent_thought(message_agent_thought)
 
         return message_agent_thought
 
-    def on_agent_end(self, message_agent_thought: MessageAgentThought, agent_model_instant: BaseLLM,
+    def on_agent_end(self, message_agent_thought: MessageAgentThought, agent_model_instance: BaseLLM,
                      agent_loop: AgentLoop):
-        agent_message_unit_price = agent_model_instant.get_token_price(1, MessageType.HUMAN)
+        agent_message_unit_price = agent_model_instance.get_tokens_unit_price(MessageType.HUMAN)
-        agent_answer_unit_price = agent_model_instant.get_token_price(1, MessageType.ASSISTANT)
+        agent_message_price_unit = agent_model_instance.get_price_unit(MessageType.HUMAN)
+        agent_answer_unit_price = agent_model_instance.get_tokens_unit_price(MessageType.ASSISTANT)
+        agent_answer_price_unit = agent_model_instance.get_price_unit(MessageType.ASSISTANT)
 
         loop_message_tokens = agent_loop.prompt_tokens
         loop_answer_tokens = agent_loop.completion_tokens
 
-        loop_total_price = self.calc_total_price(
+        loop_message_total_price = agent_model_instance.calc_tokens_price(loop_message_tokens, MessageType.HUMAN)
-            loop_message_tokens,
+        loop_answer_total_price = agent_model_instance.calc_tokens_price(loop_answer_tokens, MessageType.ASSISTANT)
-            agent_message_unit_price,
+        loop_total_price = loop_message_total_price + loop_answer_total_price
-            loop_answer_tokens,
-            agent_answer_unit_price
-        )
 
         message_agent_thought.observation = agent_loop.tool_output
         message_agent_thought.tool_process_data = ''  # currently not support
         message_agent_thought.message_token = loop_message_tokens
         message_agent_thought.message_unit_price = agent_message_unit_price
+        message_agent_thought.message_price_unit = agent_message_price_unit
         message_agent_thought.answer_token = loop_answer_tokens
         message_agent_thought.answer_unit_price = agent_answer_unit_price
+        message_agent_thought.answer_price_unit = agent_answer_price_unit
         message_agent_thought.latency = agent_loop.latency
         message_agent_thought.tokens = agent_loop.prompt_tokens + agent_loop.completion_tokens
         message_agent_thought.total_price = loop_total_price
-        message_agent_thought.currency = agent_model_instant.get_currency()
+        message_agent_thought.currency = agent_model_instance.get_currency()
-        db.session.flush()
+        db.session.commit()
 
     def on_dataset_query_end(self, dataset_query_obj: DatasetQueryObj):
         dataset_query = DatasetQuery(
@@ -252,17 +263,38 @@ class ConversationMessageTask:
         )
 
         db.session.add(dataset_query)
+        db.session.commit()
 
-    def calc_total_price(self, message_tokens, message_unit_price, answer_tokens, answer_unit_price):
+    def on_dataset_query_finish(self, resource: List):
-        message_tokens_per_1k = (decimal.Decimal(message_tokens) / 1000).quantize(decimal.Decimal('0.001'),
+        if resource and len(resource) > 0:
-                                                                                  rounding=decimal.ROUND_HALF_UP)
+            for item in resource:
-        answer_tokens_per_1k = (decimal.Decimal(answer_tokens) / 1000).quantize(decimal.Decimal('0.001'),
+                dataset_retriever_resource = DatasetRetrieverResource(
-                                                                                rounding=decimal.ROUND_HALF_UP)
+                    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.commit()
+            self.retriever_resource = resource
 
-        total_price = message_tokens_per_1k * message_unit_price + answer_tokens_per_1k * answer_unit_price
+    def message_end(self):
-        return total_price.quantize(decimal.Decimal('0.0000001'), rounding=decimal.ROUND_HALF_UP)
+        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()
 
 
@@ -356,6 +388,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 = {