feat: bump version 0.3.14 (#861)

Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>

.devcontainer/Dockerfile

@@ -0,0 +1,11 @@
+FROM mcr.microsoft.com/devcontainers/anaconda:0-3
+
+# 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>

.devcontainer/README.md

@@ -0,0 +1,37 @@
+# Devlopment with devcontainer
+This project includes a devcontainer configuration that allows you to open the project in a container with a fully configured development environment.
+Both frontend and backend environments are initialized when the container is started.
+## GitHub Codespaces
+[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/langgenius/dify)
+
+you can simply click the button above to open this project in GitHub Codespaces.
+
+For more info, check out the [GitHub documentation](https://docs.github.com/en/free-pro-team@latest/github/developing-online-with-codespaces/creating-a-codespace#creating-a-codespace).
+
+
+## VS Code Dev Containers
+[![Open in Dev Containers](https://img.shields.io/static/v1?label=Dev%20Containers&message=Open&color=blue&logo=visualstudiocode)](https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/langgenius/dify)
+
+if you have VS Code installed, you can click the button above to open this project in VS Code Dev Containers.
+
+You can learn more in the [Dev Containers documentation](https://code.visualstudio.com/docs/devcontainers/containers).
+
+
+## Pros of Devcontainer
+Unified Development Environment: By using devcontainers, you can ensure that all developers are developing in the same environment, reducing the occurrence of "it works on my machine" type of issues.
+
+Quick Start: New developers can set up their development environment in a few simple steps, without spending a lot of time on environment configuration.
+
+Isolation: Devcontainers isolate your project from your host operating system, reducing the chance of OS updates or other application installations impacting the development environment.
+
+## Cons of Devcontainer
+Learning Curve: For developers unfamiliar with Docker and VS Code, using devcontainers may be somewhat complex.
+
+Performance Impact: While usually minimal, programs running inside a devcontainer may be slightly slower than those running directly on the host.
+
+## Troubleshooting
+if you see such error message when you open this project in codespaces:
+![Alt text](troubleshooting.png)
+
+a simple workaround is change `/signin` endpoint into another one, then login with github account and close the tab, then change it back to `/signin` endpoint. Then all things will be fine.
+The reason is `signin` endpoint is not allowed in codespaces, details can be found [here](https://github.com/orgs/community/discussions/5204)

.devcontainer/devcontainer.json

@@ -0,0 +1,53 @@
+// 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)",
+	"build": { 
+		"context": "..",
+		"dockerfile": "Dockerfile"
+	},
+	"features": {
+		"ghcr.io/dhoeric/features/act:1": {},
+		"ghcr.io/devcontainers/features/node:1": {
+			"nodeGypDependencies": true,
+			"version": "lts"
+		},
+		"ghcr.io/devcontainers-contrib/features/npm-package:1": {
+			"package": "typescript",
+			"version": "latest"
+		},
+		"ghcr.io/devcontainers/features/docker-in-docker:2": {
+			"moby": true,
+			"azureDnsAutoDetection": true,
+			"installDockerBuildx": true,
+			"version": "latest",
+			"dockerDashComposeVersion": "v2"
+		}
+	},
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"ms-python.pylint",
+				"GitHub.copilot",
+				"ms-python.python"
+			]
+		}
+	},
+	"postStartCommand": "cd api && pip install -r requirements.txt",
+	"postCreateCommand": "cd web && npm install"
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "python --version",
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "root"
+}

.devcontainer/noop.txt

@@ -0,0 +1,3 @@
+This file copied into the container along with environment.yml* from the parent
+folder. This file is included to prevents the Dockerfile COPY instruction from 
+failing if no environment.yml is found.

.github/workflows/build-api-image.yml

@@ -42,12 +42,14 @@ jobs:
       uses: docker/build-push-action@v4
       with:
         context: "{{defaultContext}}:api"
-        platforms: linux/amd64,linux/arm64
+        platforms: ${{ startsWith(github.ref, 'refs/tags/') && 'linux/amd64,linux/arm64' || 'linux/amd64' }}
         build-args: |
           COMMIT_SHA=${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.revision'] }}
         push: true
         tags: ${{ steps.meta.outputs.tags }}
         labels: ${{ steps.meta.outputs.labels }}
+        cache-from: type=gha
+        cache-to: type=gha,mode=max
 
     - name: Deploy to server
       if: github.ref == 'refs/heads/deploy/dev'

.github/workflows/build-web-image.yml

@@ -42,12 +42,14 @@ jobs:
       uses: docker/build-push-action@v4
       with:
         context: "{{defaultContext}}:web"
-        platforms: linux/amd64,linux/arm64
+        platforms: ${{ startsWith(github.ref, 'refs/tags/') && 'linux/amd64,linux/arm64' || 'linux/amd64' }}
         build-args: |
           COMMIT_SHA=${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.revision'] }}
         push: true
         tags: ${{ steps.meta.outputs.tags }}
         labels: ${{ steps.meta.outputs.labels }}
+        cache-from: type=gha
+        cache-to: type=gha,mode=max
 
     - name: Deploy to server
       if: github.ref == 'refs/heads/deploy/dev'

.github/workflows/check_no_chinese_comments.py

@@ -19,7 +19,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']
+    excluded_files = ["model_template.py", 'stopwords.py', 'commands.py',
+                      'indexing_runner.py', 'web_reader_tool.py', 'spark_provider.py']
 
     for root, _, files in os.walk("."):
         for file in files:

LICENSE

@@ -1,26 +1,26 @@
 # Dify Open Source License
 
-The Dify project uses a combination of the Apache License 2.0, MIT License, and an additional agreement to protect against direct competition with Dify Cloud services.
+The Dify project is licensed under the Apache License 2.0, with the following additional conditions:
 
-As a contributor, you should agree that your contributed code:
-a. Might be subject to a more permissive open source license in the future.
+1. Dify is permitted to be used for commercialization, such as using Dify as a "backend-as-a-service" for your other applications, or delivering it to enterprises as an application development platform. However, when the following conditions are met, you must contact the producer to obtain a commercial license:
+
+a. Multi-tenant SaaS service: Unless explicitly authorized by Dify in writing, you may not use the Dify.AI source code to operate a multi-tenant SaaS service that is similar to the Dify.AI service edition.
+b. LOGO and copyright information: In the process of using Dify, you may not remove or modify the LOGO or copyright information in the Dify console.
+
+Please contact business@dify.ai by email to inquire about licensing matters.
+
+2. As a contributor, you should agree that your contributed code:
+
+a. The producer can adjust the open-source agreement to be more strict or relaxed.
 b. Can be used for commercial purposes, such as Dify's cloud business.
 
-The following components are open source under the MIT license, allowing you to build and develop applications based on them:
-- WebApp elements, e.g., web/app/components/share
-- Derived WebApp Template projects
-
-The remaining parts of the project are open source under the Apache License 2.0.
-
-With the Apache License 2.0, MIT License, and this supplementary agreement, anyone can freely use, modify, and distribute Dify, provided that:
-
-- If you use Dify solely as a backend service for other applications, no authorization is needed for commercial or closed source purposes.
-- If you wish to use Dify for commercial and closed source SaaS services similar to Dify Cloud, please contact us for authorization.
+Apart from this, all other rights and restrictions follow the Apache License 2.0. If you need more detailed information, you can refer to the full version of Apache License 2.0.
 
 The interactive design of this product is protected by appearance patent.
 
 © 2023 LangGenius, Inc.
 
+
 ----------
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -34,13 +34,3 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
-
-----------
-The MIT License
-
-Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
-
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -6,7 +6,7 @@
   <a href="./README_ES.md">Español</a>
 </p>
 
-[Website](https://dify.ai) • [Docs](https://docs.dify.ai) • [Twitter](https://twitter.com/dify_ai) • [Discord](https://discord.gg/FngNHpbcY7)
+#### [Website](https://dify.ai) • [Docs](https://docs.dify.ai) • [Deployment Docs](https://docs.dify.ai/getting-started/install-self-hosted) •  [FAQ](https://docs.dify.ai/getting-started/faq) • [Twitter](https://twitter.com/dify_ai) • [Discord](https://discord.gg/FngNHpbcY7)
 
 **Dify** is an easy-to-use LLMOps platform designed to empower more people to create sustainable, AI-native applications. With visual orchestration for various application types, Dify offers out-of-the-box, ready-to-use applications that can also serve as Backend-as-a-Service APIs. Unify your development process with one API for plugins and datasets integration, and streamline your operations using a single interface for prompt engineering, visual analytics, and continuous improvement.
 
@@ -15,17 +15,43 @@ Applications created with Dify include:
 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
-Dify is compatible with Langchain, meaning we'll gradually support multiple LLMs, currently supported:
 
-* **OpenAI** ：GPT4、GPT3.5-turbo、GPT3.5-turbo-16k、text-davinci-003
+## 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:
 
-* **Azure OpenAI**
+- [x] **OpenAI**: GPT4, GPT3.5-turbo, GPT3.5-turbo-16k, text-davinci-003 
+- [x] **Azure OpenAI Service**
+- [x] **Anthropic**: Claude2, Claude-instant
+- [x] **Replicate**
+- [x] **Hugging Face Hub**
+- [x] **MiniMax**
+- [x] **Spark**
+- [x] **Wenxin**
+- [x] **Tongyi**
+- [x] **ChatGLM**
 
-* **Antropic**：Claude2、Claude-instant
-> We've got 1000 free trial credits available for all cloud service users to try out the Claude model.Visit [Dify.ai](https://dify.ai) and
-try it now.
+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
+* 200 free OpenAI queries to build OpenAI-based apps
+* 3 million Xunfei Spark Tokens are provided for creating AI applications based on Spark.
+* 1 million MiniMax Tokens are provided for creating AI applications based on the MiniMax.
+  
+**2. Visual orchestration:** Build an AI app in minutes by writing and debugging prompts visually.
 
-* **hugging face Hub**：Coming soon.
+**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.
+
+**4. API-based:**  Backend-as-a-service. Access web apps directly or integrate via APIs without complex backend setup.
+
+**5. Plugins:** Dify "Smart Chat" now supports first-party plugins like web browsing, Google search, Wikipedia to enable online lookup, analyzing web content, and explaining the AI's reasoning process conversationally.
+
+**6. Team workspaces:** Team members can join workspaces to collaboratively edit, manage, and use team AI apps.
+
+**7. Data labeling and improvement:**  Visually inspect AI app logs and improve data via labeling. Observe the AI's reasoning process to continuously enhance performance. (Coming soon)
+
+## Use cases
+* [Create an AI ChatBot with Business Data in Minutes.](https://docs.dify.ai/use-cases/create-an-ai-chatbot-with-business-data-in-minutes)
+* [How to Build an Notion AI Assistant Based on Your Own Notes?](https://docs.dify.ai/use-cases/build-an-notion-ai-assistant)
+* [Create a Midjoureny Prompt Bot Without Code in Just a Few Minutes.](https://docs.dify.ai/use-cases/create-a-midjoureny-prompt-bot-with-dify)
 
 ## Use Cloud Services
 
@@ -37,7 +63,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
@@ -139,7 +165,6 @@ To protect your privacy, please avoid posting security issues on GitHub. Instead
 This software uses the following open-source software:
 
 - Chase, H. (2022). LangChain [Computer software]. https://github.com/hwchase17/langchain
-- Liu, J. (2022). LlamaIndex [Computer software]. doi: 10.5281/zenodo.1234.
 
 For more information, please refer to the official website or license text of the respective software.
 

README_CN.md

@@ -7,9 +7,9 @@
 </p>
 
 
-[官方网站](https://dify.ai) • [文档](https://docs.dify.ai/v/zh-hans) • [Twitter](https://twitter.com/dify_ai) •  [Discord](https://discord.gg/FngNHpbcY7)
+####  [官方网站](https://dify.ai) •  [使用文档](https://docs.dify.ai/v/zh-hans)  · [部署文档](https://docs.dify.ai/v/zh-hans/getting-started/install-self-hosted)  ·  [FAQ](https://docs.dify.ai/v/zh-hans/getting-started/faq) • [Twitter](https://twitter.com/dify_ai) •  [Discord](https://discord.gg/FngNHpbcY7)
 
-**Dify** 是一个易用的 LLMOps 平台，旨在让更多人可以创建可持续运营的原生 AI 应用。Dify 提供多种类型应用的可视化编排，应用可开箱即用，也能以“后端即服务”的 API 提供服务。
+**Dify** 是一个易用的 LLMOps 平台，基于不同的大型语言模型能力，让更多人可以简易地创建可持续运营的原生 AI 应用。Dify 提供多种类型应用的可视化编排，应用可开箱即用，也能以“后端即服务”的 API 提供服务。
 
 通过 Dify 创建的应用包含了：
 
@@ -17,20 +17,43 @@
 - 一套 API 即可包含插件、上下文增强等能力，替你省下了后端代码的编写工作
 - 可视化的对应用进行数据分析，查阅日志或进行标注
 
-Dify 兼容 Langchain，这意味着我们将逐步支持多种 LLMs ，目前支持的模型供应商：
 
-* **OpenAI**：GPT4、GPT3.5-turbo、GPT3.5-turbo-16k、text-davinci-003 
 
-* **Azure OpenAI Service**
-* **Anthropic**：Claude2、Claude-instant
+## 核心能力
+1. **模型支持：** 你可以在 Dify 上选择基于不同模型的能力来开发你的 AI 应用。Dify 兼容 Langchain，这意味着我们将逐步支持多种 LLMs ，目前支持的模型供应商：
 
-> 我们为所有注册云端版的用户免费提供了 1000 次 Claude 模型的消息调用额度，登录 [dify.ai](https://cloud.dify.ai) 即可使用。
+- [x] **OpenAI**：GPT4、GPT3.5-turbo、GPT3.5-turbo-16k、text-davinci-003 
+- [x] **Azure OpenAI Service**
+- [x] **Anthropic**：Claude2、Claude-instant
+- [x] **Replicate**
+- [x] **Hugging Face Hub**
+- [x] **MiniMax**
+- [x] **讯飞星火大模型**
+- [x] **文心一言**
+- [x] **通义千问**
+- [x] **ChatGLM**
+
+我们为所有注册云端版的用户免费提供以下资源（登录 [dify.ai](https://cloud.dify.ai) 即可使用）：
+* 1000 次 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 集成到您的应用中，无需关注复杂的后端架构和部署过程。
+5. **插件能力：** Dify 「智聊」平台已支持网页浏览、Google 搜索、Wikipedia 查询等第一方插件，可在对话中实现联网搜索、分析网页内容、展示 AI 的推理过程。
+6. **团队 Workspace：** 团队成员可加入 Workspace 编辑、管理和使用团队内的 AI 应用。
+6. **数据标注与改进：** 可视化查阅 AI 应用日志并对数据进行改进标注，观测 AI 的推理过程，不断提高其性能。（Coming soon）
+ -----------------------------
+ ## Use cases
+ * [几分钟创建一个带有业务数据的官网 AI 智能客服](https://docs.dify.ai/v/zh-hans/use-cases/create-an-ai-chatbot-with-business-data-in-minutes)
+ * [构建一个 Notion AI 助手](https://docs.dify.ai/v/zh-hans/use-cases/build-an-notion-ai-assistant)
+ * [创建 Midjoureny 提示词机器人](https://docs.dify.ai/v/zh-hans/use-cases/create-a-midjoureny-prompt-word-robot-with-zero-code)
 
-* **Hugging Face Hub**（即将推出）
 
 ## 使用云服务
 
-访问 [Dify.ai](https://cloud.dify.ai)
+访问 [Dify.ai](https://cloud.dify.ai) 使用云端版。
 
 ## 安装社区版
 
@@ -38,7 +61,7 @@ Dify 兼容 Langchain，这意味着我们将逐步支持多种 LLMs ，目前
 
 在安装 Dify 之前，请确保您的机器满足以下最低系统要求：
 
-- CPU >= 1 Core
+- CPU >= 2 Core
 - RAM >= 4GB
 
 ### 快速启动
@@ -65,12 +88,10 @@ docker compose up -d
 
 我们正在开发中的功能：
 
-- **数据集**，支持更多的数据集，例如同步 Notion 或网页的内容
-我们将支持更多的数据集，包括文本、网页，甚至 Notion 内容。用户可以根据自己的数据源构建 AI 应用程序。
-- **插件**，推出符合 ChatGPT 标准的插件，或使用 Dify 产生的插件
-我们将发布符合 ChatGPT 标准的插件，或者 Dify 自己的插件，以在应用程序中启用更多功能。
-- **开源模型**，例如采用 Llama 作为模型提供者，或进行进一步的微调
-我们将与优秀的开源模型如 Llama 合作，通过在我们的平台中提供它们作为模型选项，或使用它们进行进一步的微调。
+- **数据集**，支持更多的数据集，通过网页、API 同步内容。用户可以根据自己的数据源构建 AI 应用程序。
+- **插件**，我们将发布符合 ChatGPT 标准的插件，支持更多 Dify 自己的插件，支持用户自定义插件能力，以在应用程序中启用更多功能，例如以支持以目标为导向的分解推理任务。
+- **开源模型支持**，支持 Hugging face Hub 上的开源模型。例如采用 Llama 作为模型提供者，或进行进一步的微调
+我们将与优秀的开源模型合作，通过在我们的平台中提供它们作为模型选项，或使用它们进行进一步的微调。
 
 ## Q&A
 
@@ -84,11 +105,11 @@ A: 一个有价值的应用由 Prompt Engineering、上下文增强和 Fine-tune
 
 **Q: 如果要创建一个自己的应用，我需要准备什么？**
 
-A: 我们假定你已经有了 OpenAI API Key，如果没有请去注册一个。如果你已经有了一些内容可以作为训练上下文，就太好了。
+A: 我们假定你已经有了 OpenAI 或 Claude 等模型的 API Key，如果没有请去注册一个。如果你已经有了一些内容可以作为训练上下文，就太好了。
 
 **Q: 提供哪些界面语言？**
 
-A: 现已支持英文与中文，你可以为我们贡献语言包。
+A:  支持英文、中文，你可以为我们贡献语言包并提供维护支持。
 
 ## Star History
 
@@ -134,7 +155,6 @@ A: 现已支持英文与中文，你可以为我们贡献语言包。
 本软件使用了以下开源软件：
 
 - Chase, H. (2022). LangChain [Computer software]. https://github.com/hwchase17/langchain
-- Liu, J. (2022). LlamaIndex [Computer software]. doi: 10.5281/zenodo.1234.
 
 更多信息，请参考相应软件的官方网站或许可证文本。
 

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
@@ -115,7 +115,6 @@ Para proteger tu privacidad, evita publicar problemas de seguridad en GitHub. En
 Este software utiliza el siguiente software de código abierto:
 
 - Chase, H. (2022). LangChain [Software de computadora]. https://github.com/hwchase17/langchain
-- Liu, J. (2022). LlamaIndex [Software de computadora]. doi: 10.5281/zenodo.1234.
 
 Para obtener más información, consulta el sitio web oficial o el texto de la licencia del software correspondiente.
 

README_JA.md

@@ -114,7 +114,6 @@ A: 現在、英語と中国語に対応しており、言語パックを寄贈
 本ソフトウェアは、以下のオープンソースソフトウェアを使用しています:
 
 - Chase, H. (2022). LangChain [Computer software]. https://github.com/hwchase17/langchain
-- Liu, J. (2022). LlamaIndex [Computer software]. doi: 10.5281/zenodo.1234.
 
 詳しくは、各ソフトウェアの公式サイトまたはライセンス文をご参照ください。
 

api/.dockerignore

@@ -1,2 +1,11 @@
 .env
-storage/privkeys/*
+*.env.*
+
+storage/privkeys/*
+
+# Logs
+logs
+*.log*
+
+# jetbrains
+.idea

api/.env.example

@@ -9,17 +9,13 @@ SECRET_KEY=
 
 # Console API base URL
 CONSOLE_API_URL=http://127.0.0.1:5001
-
-# Console frontend web base URL
 CONSOLE_WEB_URL=http://127.0.0.1:3000
 
 # Service API base URL
 SERVICE_API_URL=http://127.0.0.1:5001
 
-# Web APP API base URL
+# Web APP base URL
 APP_API_URL=http://127.0.0.1:5001
-
-# Web APP frontend web base URL
 APP_WEB_URL=http://127.0.0.1:3000
 
 # celery configuration
@@ -102,3 +98,29 @@ NOTION_INTEGRATION_TYPE=public
 NOTION_CLIENT_SECRET=you-client-secret
 NOTION_CLIENT_ID=you-client-id
 NOTION_INTERNAL_SECRET=you-internal-secret
+
+# Hosted Model Credentials
+HOSTED_OPENAI_ENABLED=false
+HOSTED_OPENAI_API_KEY=
+HOSTED_OPENAI_API_BASE=
+HOSTED_OPENAI_API_ORGANIZATION=
+HOSTED_OPENAI_QUOTA_LIMIT=200
+HOSTED_OPENAI_PAID_ENABLED=false
+HOSTED_OPENAI_PAID_STRIPE_PRICE_ID=
+HOSTED_OPENAI_PAID_INCREASE_QUOTA=1
+
+HOSTED_AZURE_OPENAI_ENABLED=false
+HOSTED_AZURE_OPENAI_API_KEY=
+HOSTED_AZURE_OPENAI_API_BASE=
+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_PAID_ENABLED=false
+HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID=
+HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA=1
+
+STRIPE_API_KEY=
+STRIPE_WEBHOOK_SECRET=

api/Dockerfile

@@ -1,4 +1,4 @@
-FROM langgenius/base:1.0.0-bullseye-slim as langgenius-api
+FROM python:3.10-slim
 
 LABEL maintainer="takatost@gmail.com"
 
@@ -15,6 +15,9 @@ EXPOSE 5001
 
 WORKDIR /app/api
 
+RUN apt-get update && \
+    apt-get install -y bash curl wget vim gcc g++ python3-dev libc-dev libffi-dev
+
 COPY requirements.txt /app/api/requirements.txt
 
 RUN pip install -r requirements.txt

api/README.md

@@ -8,7 +8,7 @@
    
    ```bash
    cd ../docker
-   docker-compose -f docker-compose.middleware.yaml up -d
+   docker-compose -f docker-compose.middleware.yaml -p dify up -d
    cd ../api
    ```
 2. Copy `.env.example` to `.env`
@@ -33,9 +33,30 @@
    ```bash
    flask db upgrade
    ```
+
+   ⚠️ If you encounter problems with jieba, for example
+
+   ```
+   > flask db upgrade
+   Error: While importing 'app', an ImportError was raised:
+   ```
+
+   Please run the following command instead.
+
+   ```
+   pip install -r requirements.txt --upgrade --force-reinstall
+   ```
+   
 6. Start backend:
    ```bash
    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`, 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 -Q dataset,generation,mail`, celery can do dataset importing and other async tasks.
+
+8. Start frontend:
+
+   ```
+   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
+   ```
+   This will start a dify frontend, now you are all set, happy coding!

api/app.py

@@ -16,13 +16,14 @@ from flask import Flask, request, Response, session
 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, \
-    ext_database, ext_storage, ext_mail
+    ext_database, ext_storage, ext_mail, ext_stripe
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
 
 # DO NOT REMOVE BELOW
-from models import model, account, dataset, web, task, source
+from models import model, account, dataset, web, task, source, tool
 from events import event_handlers
 # DO NOT REMOVE ABOVE
 
@@ -31,6 +32,7 @@ from config import Config, CloudEditionConfig
 from commands import register_commands
 from models.account import TenantAccountJoin, AccountStatus
 from models.model import Account, EndUser, App
+from services.account_service import TenantService
 
 import warnings
 warnings.simplefilter("ignore", ResourceWarning)
@@ -70,7 +72,7 @@ def create_app(test_config=None) -> Flask:
     register_blueprints(app)
     register_commands(app)
 
-    core.init_app(app)
+    hosted.init_app(app)
 
     return app
 
@@ -87,6 +89,16 @@ def initialize_extensions(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
@@ -119,7 +131,9 @@ def load_user(user_id):
 
                     if tenant_account_join:
                         account.current_tenant_id = tenant_account_join.tenant_id
-                        session['workspace_id'] = account.current_tenant_id
+                    else:
+                        _create_tenant_for_account(account)
+                    session['workspace_id'] = account.current_tenant_id
                 else:
                     account.current_tenant_id = workspace_id
             else:
@@ -127,7 +141,9 @@ def load_user(user_id):
                     TenantAccountJoin.account_id == account.id).first()
                 if tenant_account_join:
                     account.current_tenant_id = tenant_account_join.tenant_id
-                    session['workspace_id'] = account.current_tenant_id
+                else:
+                    _create_tenant_for_account(account)
+                session['workspace_id'] = account.current_tenant_id
 
             account.last_active_at = datetime.utcnow()
             db.session.commit()
@@ -232,5 +248,18 @@ def threads():
     }
 
 
+@app.route('/db-pool-stat')
+def pool_stat():
+    engine = db.engine
+    return {
+        'pool_size': engine.pool.size(),
+        'checked_in_connections': engine.pool.checkedin(),
+        'checked_out_connections': engine.pool.checkedout(),
+        'overflow_connections': engine.pool.overflow(),
+        'connection_timeout': engine.pool.timeout(),
+        'recycle_time': db.engine.pool._recycle
+    }
+
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5001)

api/commands.py

@@ -1,25 +1,26 @@
 import datetime
-import logging
+import math
 import random
 import string
+import time
 
 import click
 from flask import current_app
 from werkzeug.exceptions import NotFound
 
 from core.index.index import IndexBuilder
+from core.model_providers.providers.hosted import hosted_model_providers
 from libs.password import password_pattern, valid_password, hash_password
 from libs.helper import email as email_validate
 from extensions.ext_database import db
 from libs.rsa import generate_key_pair
 from models.account import InvitationCode, Tenant
-from models.dataset import Dataset
+from models.dataset import Dataset, DatasetQuery, Document
 from models.model import Account
 import secrets
 import base64
 
-from models.provider import Provider, ProviderName
-from services.provider_service import ProviderService
+from models.provider import Provider, ProviderType, ProviderQuotaType
 
 
 @click.command('reset-password', help='Reset the account password.')
@@ -172,7 +173,7 @@ def recreate_all_dataset_indexes():
     page = 1
     while True:
         try:
-            datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality')\
+            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
@@ -188,37 +189,104 @@ def recreate_all_dataset_indexes():
                 else:
                     click.echo('passed.')
             except Exception as e:
-                click.echo(click.style('Recreate dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                click.echo(
+                    click.style('Recreate dataset index error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
                 continue
 
     click.echo(click.style('Congratulations! Recreate {} dataset indexes.'.format(recreate_count), fg='green'))
 
 
+@click.command('clean-unused-dataset-indexes', help='Clean unused dataset indexes.')
+def clean_unused_dataset_indexes():
+    click.echo(click.style('Start clean unused dataset indexes.', fg='green'))
+    clean_days = int(current_app.config.get('CLEAN_DAY_SETTING'))
+    start_at = time.perf_counter()
+    thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days)
+    page = 1
+    while True:
+        try:
+            datasets = db.session.query(Dataset).filter(Dataset.created_at < thirty_days_ago) \
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+        except NotFound:
+            break
+        page += 1
+        for dataset in datasets:
+            dataset_query = db.session.query(DatasetQuery).filter(
+                DatasetQuery.created_at > thirty_days_ago,
+                DatasetQuery.dataset_id == dataset.id
+            ).all()
+            if not dataset_query or len(dataset_query) == 0:
+                documents = db.session.query(Document).filter(
+                    Document.dataset_id == dataset.id,
+                    Document.indexing_status == 'completed',
+                    Document.enabled == True,
+                    Document.archived == False,
+                    Document.updated_at > thirty_days_ago
+                ).all()
+                if not documents or len(documents) == 0:
+                    try:
+                        # remove index
+                        vector_index = IndexBuilder.get_index(dataset, 'high_quality')
+                        kw_index = IndexBuilder.get_index(dataset, 'economy')
+                        # delete from vector index
+                        if vector_index:
+                            vector_index.delete()
+                        kw_index.delete()
+                        # update document
+                        update_params = {
+                            Document.enabled: False
+                        }
+
+                        Document.query.filter_by(dataset_id=dataset.id).update(update_params)
+                        db.session.commit()
+                        click.echo(click.style('Cleaned unused dataset {} from db success!'.format(dataset.id),
+                                               fg='green'))
+                    except Exception as e:
+                        click.echo(
+                            click.style('clean dataset index error: {} {}'.format(e.__class__.__name__, str(e)),
+                                        fg='red'))
+    end_at = time.perf_counter()
+    click.echo(click.style('Cleaned unused dataset from db success latency: {}'.format(end_at - start_at), fg='green'))
+
+
 @click.command('sync-anthropic-hosted-providers', help='Sync anthropic hosted providers.')
 def sync_anthropic_hosted_providers():
+    if not hosted_model_providers.anthropic:
+        click.echo(click.style('Anthropic hosted provider is not configured.', fg='red'))
+        return
+
     click.echo(click.style('Start sync anthropic hosted providers.', fg='green'))
     count = 0
 
     page = 1
     while True:
         try:
-            tenants = db.session.query(Tenant).order_by(Tenant.created_at.desc()).paginate(page=page, per_page=50)
+            providers = db.session.query(Provider).filter(
+                Provider.provider_name == 'anthropic',
+                Provider.provider_type == ProviderType.SYSTEM.value,
+                Provider.quota_type == ProviderQuotaType.TRIAL.value,
+            ).order_by(Provider.created_at.desc()).paginate(page=page, per_page=100)
         except NotFound:
             break
 
         page += 1
-        for tenant in tenants:
+        for provider in providers:
             try:
-                click.echo('Syncing tenant anthropic hosted provider: {}'.format(tenant.id))
-                ProviderService.create_system_provider(
-                    tenant,
-                    ProviderName.ANTHROPIC.value,
-                    current_app.config['ANTHROPIC_HOSTED_QUOTA_LIMIT'],
-                    True
-                )
+                click.echo('Syncing tenant anthropic hosted provider: {}'.format(provider.tenant_id))
+                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 \
+                    else original_quota_limit * division
+                provider.quota_used = division * provider.quota_used
+                db.session.commit()
+
                 count += 1
             except Exception as e:
-                click.echo(click.style('Sync tenant anthropic hosted provider error: {} {}'.format(e.__class__.__name__, str(e)), fg='red'))
+                click.echo(click.style(
+                    'Sync tenant anthropic hosted provider error: {} {}'.format(e.__class__.__name__, str(e)),
+                    fg='red'))
                 continue
 
     click.echo(click.style('Congratulations! Synced {} anthropic hosted providers.'.format(count), fg='green'))
@@ -231,3 +299,4 @@ def register_commands(app):
     app.cli.add_command(reset_encrypt_key_pair)
     app.cli.add_command(recreate_all_dataset_indexes)
     app.cli.add_command(sync_anthropic_hosted_providers)
+    app.cli.add_command(clean_unused_dataset_indexes)

api/config.py

@@ -41,6 +41,7 @@ DEFAULTS = {
     'SESSION_USE_SIGNER': 'True',
     'DEPLOY_ENV': 'PRODUCTION',
     'SQLALCHEMY_POOL_SIZE': 30,
+    'SQLALCHEMY_POOL_RECYCLE': 3600,
     'SQLALCHEMY_ECHO': 'False',
     'SENTRY_TRACES_SAMPLE_RATE': 1.0,
     'SENTRY_PROFILES_SAMPLE_RATE': 1.0,
@@ -50,10 +51,18 @@ DEFAULTS = {
     'PDF_PREVIEW': 'True',
     'LOG_LEVEL': 'INFO',
     'DISABLE_PROVIDER_CONFIG_VALIDATION': 'False',
-    'DEFAULT_LLM_PROVIDER': 'openai',
-    'OPENAI_HOSTED_QUOTA_LIMIT': 200,
-    'ANTHROPIC_HOSTED_QUOTA_LIMIT': 1000,
-    'TENANT_DOCUMENT_COUNT': 100
+    '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_ENABLED': 'False',
+    'HOSTED_ANTHROPIC_PAID_ENABLED': 'False',
+    'HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA': 1,
+    'TENANT_DOCUMENT_COUNT': 100,
+    'CLEAN_DAY_SETTING': 30
 }
 
 
@@ -89,7 +98,7 @@ class Config:
         self.CONSOLE_URL = get_env('CONSOLE_URL')
         self.API_URL = get_env('API_URL')
         self.APP_URL = get_env('APP_URL')
-        self.CURRENT_VERSION = "0.3.10"
+        self.CURRENT_VERSION = "0.3.14"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = "SELF_HOSTED"
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')
@@ -181,7 +190,10 @@ class Config:
         }
 
         self.SQLALCHEMY_DATABASE_URI = f"postgresql://{db_credentials['DB_USERNAME']}:{db_credentials['DB_PASSWORD']}@{db_credentials['DB_HOST']}:{db_credentials['DB_PORT']}/{db_credentials['DB_DATABASE']}"
-        self.SQLALCHEMY_ENGINE_OPTIONS = {'pool_size': int(get_env('SQLALCHEMY_POOL_SIZE'))}
+        self.SQLALCHEMY_ENGINE_OPTIONS = {
+            'pool_size': int(get_env('SQLALCHEMY_POOL_SIZE')),
+            'pool_recycle': int(get_env('SQLALCHEMY_POOL_RECYCLE'))
+        }
 
         self.SQLALCHEMY_ECHO = get_bool_env('SQLALCHEMY_ECHO')
 
@@ -193,20 +205,35 @@ class Config:
         self.BROKER_USE_SSL = self.CELERY_BROKER_URL.startswith('rediss://')
 
         # hosted provider credentials
-        self.OPENAI_API_KEY = get_env('OPENAI_API_KEY')
-        self.ANTHROPIC_API_KEY = get_env('ANTHROPIC_API_KEY')
+        self.HOSTED_OPENAI_ENABLED = get_bool_env('HOSTED_OPENAI_ENABLED')
+        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_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'))
 
-        self.OPENAI_HOSTED_QUOTA_LIMIT = get_env('OPENAI_HOSTED_QUOTA_LIMIT')
-        self.ANTHROPIC_HOSTED_QUOTA_LIMIT = get_env('ANTHROPIC_HOSTED_QUOTA_LIMIT')
+        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_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_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.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')
 
-        # For temp use only
-        # set default LLM provider, default is 'openai', support `azure_openai`
-        self.DEFAULT_LLM_PROVIDER = get_env('DEFAULT_LLM_PROVIDER')
-
         # notion import setting
         self.NOTION_CLIENT_ID = get_env('NOTION_CLIENT_ID')
         self.NOTION_CLIENT_SECRET = get_env('NOTION_CLIENT_SECRET')
@@ -215,6 +242,7 @@ class Config:
         self.NOTION_INTEGRATION_TOKEN = get_env('NOTION_INTEGRATION_TOKEN')
 
         self.TENANT_DOCUMENT_COUNT = get_env('TENANT_DOCUMENT_COUNT')
+        self.CLEAN_DAY_SETTING = get_env('CLEAN_DAY_SETTING')
 
 
 class CloudEditionConfig(Config):

api/controllers/console/__init__.py

@@ -18,7 +18,13 @@ from .auth import login, oauth, data_source_oauth, activate
 from .datasets import datasets, datasets_document, datasets_segments, file, hit_testing, data_source
 
 # Import workspace controllers
-from .workspace import workspace, members, providers, account
+from .workspace import workspace, members, providers, model_providers, account, tool_providers, models
 
 # Import explore controllers
 from .explore import installed_app, recommended_app, completion, conversation, message, parameter, saved_message, audio
+
+# Import universal chat controllers
+from .universal_chat import chat, conversation, message, parameter, audio
+
+# Import webhook controllers
+from .webhook import stripe

api/controllers/console/admin.py

@@ -55,7 +55,7 @@ class InsertExploreAppListApi(Resource):
 
         app = App.query.filter(App.id == args['app_id']).first()
         if not app:
-            raise NotFound('App not found')
+            raise NotFound(f'App \'{args["app_id"]}\' is not found')
 
         site = app.site
         if not site:
@@ -63,10 +63,12 @@ class InsertExploreAppListApi(Resource):
             copy_right = args['copyright'] if args['copyright'] else ''
             privacy_policy = args['privacy_policy'] if args['privacy_policy'] else ''
         else:
-            desc = site.description if (site.description if not args['desc'] else args['desc']) else ''
-            copy_right = site.copyright if (site.copyright if not args['copyright'] else args['copyright']) else ''
-            privacy_policy = site.privacy_policy \
-                if (site.privacy_policy if not args['privacy_policy'] else args['privacy_policy']) else ''
+            desc = site.description if site.description else \
+                args['desc'] if args['desc'] else ''
+            copy_right = site.copyright if site.copyright else \
+                args['copyright'] if args['copyright'] else ''
+            privacy_policy = site.privacy_policy if site.privacy_policy else \
+                args['privacy_policy'] if args['privacy_policy']  else ''
 
         recommended_app = RecommendedApp.query.filter(RecommendedApp.app_id == args['app_id']).first()
 

api/controllers/console/app/app.py

@@ -2,16 +2,17 @@
 import json
 from datetime import datetime
 
-import flask
 from flask_login import login_required, current_user
 from flask_restful import Resource, reqparse, fields, marshal_with, abort, inputs
-from werkzeug.exceptions import Unauthorized, Forbidden
+from werkzeug.exceptions import Forbidden
 
 from constants.model_template import model_templates, demo_model_templates
 from controllers.console import api
-from controllers.console.app.error import AppNotFoundError
+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.model_factory import ModelFactory
+from core.model_providers.models.entity.model_params import ModelType
 from events.app_event import app_was_created, app_was_deleted
 from libs.helper import TimestampField
 from extensions.ext_database import db
@@ -24,6 +25,7 @@ model_config_fields = {
     '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,
@@ -96,7 +98,8 @@ class AppListApi(Resource):
         args = parser.parse_args()
 
         app_models = db.paginate(
-            db.select(App).where(App.tenant_id == current_user.current_tenant_id).order_by(App.created_at.desc()),
+            db.select(App).where(App.tenant_id == current_user.current_tenant_id,
+                                 App.is_universal == False).order_by(App.created_at.desc()),
             page=args['page'],
             per_page=args['limit'],
             error_out=False)
@@ -124,9 +127,9 @@ class AppListApi(Resource):
         if args['model_config'] is not None:
             # validate config
             model_configuration = AppModelConfigService.validate_configuration(
+                tenant_id=current_user.current_tenant_id,
                 account=current_user,
-                config=args['model_config'],
-                mode=args['mode']
+                config=args['model_config']
             )
 
             app = App(
@@ -147,6 +150,7 @@ class AppListApi(Resource):
                 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'],
@@ -161,6 +165,21 @@ class AppListApi(Resource):
             app = App(**model_config_template['app'])
             app_model_config = AppModelConfig(**model_config_template['model_config'])
 
+            default_model = ModelFactory.get_default_model(
+                tenant_id=current_user.current_tenant_id,
+                model_type=ModelType.TEXT_GENERATION
+            )
+
+            if default_model:
+                model_dict = app_model_config.model_dict
+                model_dict['provider'] = default_model.provider_name
+                model_dict['name'] = default_model.model_name
+                app_model_config.model = json.dumps(model_dict)
+            else:
+                raise ProviderNotInitializeError(
+                    f"No Text Generation Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+
         app.name = args['name']
         app.mode = args['mode']
         app.icon = args['icon']
@@ -275,6 +294,10 @@ class AppApi(Resource):
     def delete(self, app_id):
         """Delete app"""
         app_id = str(app_id)
+
+        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)
@@ -294,19 +317,13 @@ class AppNameApi(Resource):
     @account_initialization_required
     @marshal_with(app_detail_fields)
     def post(self, app_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()
+        app_id = str(app_id)
+        app = _get_app(app_id, current_user.current_tenant_id)
 
         parser = reqparse.RequestParser()
         parser.add_argument('name', type=str, required=True, location='json')
         args = parser.parse_args()
 
-        app = db.get_or_404(App, str(app_id))
-        if app.tenant_id != flask.session.get('tenant_id'):
-            raise Unauthorized()
-
         app.name = args.get('name')
         app.updated_at = datetime.utcnow()
         db.session.commit()
@@ -319,20 +336,14 @@ class AppIconApi(Resource):
     @account_initialization_required
     @marshal_with(app_detail_fields)
     def post(self, app_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()
+        app_id = str(app_id)
+        app = _get_app(app_id, current_user.current_tenant_id)
 
         parser = reqparse.RequestParser()
         parser.add_argument('icon', type=str, location='json')
         parser.add_argument('icon_background', type=str, location='json')
         args = parser.parse_args()
 
-        app = db.get_or_404(App, str(app_id))
-        if app.tenant_id != flask.session.get('tenant_id'):
-            raise Unauthorized()
-
         app.icon = args.get('icon')
         app.icon_background = args.get('icon_background')
         app.updated_at = datetime.utcnow()
@@ -386,29 +397,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):
@@ -438,6 +426,7 @@ class AppCopy(Resource):
             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,
@@ -470,21 +459,11 @@ 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>')
 api.add_resource(AppCopy, '/apps/<uuid:app_id>/copy')
 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

@@ -14,7 +14,7 @@ from controllers.console.app.error import AppUnavailableError, \
     UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from flask_restful import Resource
 from services.audio_service import AudioService

api/controllers/console/app/completion.py

@@ -17,7 +17,7 @@ from controllers.console.app.error import ConversationCompletedError, AppUnavail
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.conversation_message_task import PubHandler
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value
 from flask_restful import Resource, reqparse
@@ -41,8 +41,11 @@ class CompletionMessageApi(Resource):
         parser.add_argument('inputs', type=dict, required=True, location='json')
         parser.add_argument('query', type=str, location='json')
         parser.add_argument('model_config', type=dict, required=True, location='json')
+        parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         args = parser.parse_args()
 
+        streaming = args['response_mode'] != 'blocking'
+
         account = flask_login.current_user
 
         try:
@@ -51,7 +54,7 @@ class CompletionMessageApi(Resource):
                 user=account,
                 args=args,
                 from_source='console',
-                streaming=True,
+                streaming=streaming,
                 is_model_config_override=True
             )
 
@@ -111,8 +114,11 @@ class ChatMessageApi(Resource):
         parser.add_argument('query', type=str, required=True, location='json')
         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')
         args = parser.parse_args()
 
+        streaming = args['response_mode'] != 'blocking'
+
         account = flask_login.current_user
 
         try:
@@ -121,7 +127,7 @@ class ChatMessageApi(Resource):
                 user=account,
                 args=args,
                 from_source='console',
-                streaming=True,
+                streaming=streaming,
                 is_model_config_override=True
             )
 

api/controllers/console/app/conversation.py

@@ -95,6 +95,7 @@ class CompletionConversationApi(Resource):
         '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,
@@ -160,7 +161,7 @@ class CompletionConversationApi(Resource):
 
         if args['end']:
             end_datetime = datetime.strptime(args['end'], '%Y-%m-%d %H:%M')
-            end_datetime = end_datetime.replace(second=0)
+            end_datetime = end_datetime.replace(second=59)
 
             end_datetime_timezone = timezone.localize(end_datetime)
             end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
@@ -246,6 +247,7 @@ class ChatConversationApi(Resource):
         '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,
@@ -316,7 +318,7 @@ class ChatConversationApi(Resource):
 
         if args['end']:
             end_datetime = datetime.strptime(args['end'], '%Y-%m-%d %H:%M')
-            end_datetime = end_datetime.replace(second=0)
+            end_datetime = end_datetime.replace(second=59)
 
             end_datetime_timezone = timezone.localize(end_datetime)
             end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)

api/controllers/console/app/generator.py

@@ -7,7 +7,7 @@ from controllers.console.app.error import ProviderNotInitializeError, ProviderQu
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.generator.llm_generator import LLMGenerator
-from core.llm.error import ProviderTokenNotInitError, QuotaExceededError, LLMBadRequestError, LLMAPIConnectionError, \
+from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, LLMBadRequestError, LLMAPIConnectionError, \
     LLMAPIUnavailableError, LLMRateLimitError, LLMAuthorizationError, ModelCurrentlyNotSupportError
 
 

api/controllers/console/app/message.py

@@ -14,7 +14,7 @@ from controllers.console.app.error import CompletionRequestError, ProviderNotIni
     AppMoreLikeThisDisabledError, ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.llm.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value, TimestampField
 from libs.infinite_scroll_pagination import InfiniteScrollPagination

api/controllers/console/app/model_config.py

@@ -28,9 +28,9 @@ class ModelConfigResource(Resource):
 
         # validate config
         model_configuration = AppModelConfigService.validate_configuration(
+            tenant_id=current_user.current_tenant_id,
             account=current_user,
-            config=request.json,
-            mode=app_model.mode
+            config=request.json
         )
 
         new_app_model_config = AppModelConfig(
@@ -43,6 +43,7 @@ class ModelConfigResource(Resource):
             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'],

api/controllers/console/app/site.py

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

@@ -398,9 +398,74 @@ class AverageResponseTimeStatistic(Resource):
         })
 
 
+class TokensPerSecondStatistic(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, app_id):
+        account = current_user
+        app_id = str(app_id)
+        app_model = _get_app(app_id)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('start', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
+        parser.add_argument('end', type=datetime_string('%Y-%m-%d %H:%M'), location='args')
+        args = parser.parse_args()
+
+        sql_query = '''SELECT date(DATE_TRUNC('day', created_at AT TIME ZONE 'UTC' AT TIME ZONE :tz )) AS date, 
+    CASE 
+        WHEN SUM(provider_response_latency) = 0 THEN 0
+        ELSE (SUM(answer_tokens) / SUM(provider_response_latency))
+    END as tokens_per_second
+FROM messages
+WHERE app_id = :app_id'''
+        arg_dict = {'tz': account.timezone, 'app_id': app_model.id}
+
+        timezone = pytz.timezone(account.timezone)
+        utc_timezone = pytz.utc
+
+        if args['start']:
+            start_datetime = datetime.strptime(args['start'], '%Y-%m-%d %H:%M')
+            start_datetime = start_datetime.replace(second=0)
+
+            start_datetime_timezone = timezone.localize(start_datetime)
+            start_datetime_utc = start_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += ' and created_at >= :start'
+            arg_dict['start'] = start_datetime_utc
+
+        if args['end']:
+            end_datetime = datetime.strptime(args['end'], '%Y-%m-%d %H:%M')
+            end_datetime = end_datetime.replace(second=0)
+
+            end_datetime_timezone = timezone.localize(end_datetime)
+            end_datetime_utc = end_datetime_timezone.astimezone(utc_timezone)
+
+            sql_query += ' and created_at < :end'
+            arg_dict['end'] = end_datetime_utc
+
+        sql_query += ' GROUP BY date order by date'
+
+        with db.engine.begin() as conn:
+            rs = conn.execute(db.text(sql_query), arg_dict)
+
+        response_data = []
+
+        for i in rs:
+            response_data.append({
+                'date': str(i.date),
+                'tps': round(i.tokens_per_second, 4)
+            })
+
+        return jsonify({
+            'data': response_data
+        })
+
+
 api.add_resource(DailyConversationStatistic, '/apps/<uuid:app_id>/statistics/daily-conversations')
 api.add_resource(DailyTerminalsStatistic, '/apps/<uuid:app_id>/statistics/daily-end-users')
 api.add_resource(DailyTokenCostStatistic, '/apps/<uuid:app_id>/statistics/token-costs')
 api.add_resource(AverageSessionInteractionStatistic, '/apps/<uuid:app_id>/statistics/average-session-interactions')
 api.add_resource(UserSatisfactionRateStatistic, '/apps/<uuid:app_id>/statistics/user-satisfaction-rate')
 api.add_resource(AverageResponseTimeStatistic, '/apps/<uuid:app_id>/statistics/average-response-time')
+api.add_resource(TokensPerSecondStatistic, '/apps/<uuid:app_id>/statistics/tokens-per-second')

api/controllers/console/auth/login.py

@@ -35,7 +35,7 @@ class LoginApi(Resource):
         try:
             TenantService.switch_tenant(account)
         except Exception:
-            raise AccountNotLinkTenantError("Account not link tenant")
+            pass
 
         flask_login.login_user(account, remember=args['remember_me'])
         AccountService.update_last_login(account, request)

api/controllers/console/datasets/data_source.py

@@ -255,7 +255,7 @@ class DataSourceNotionApi(Resource):
         # validate args
         DocumentService.estimate_args_validate(args)
         indexing_runner = IndexingRunner()
-        response = indexing_runner.notion_indexing_estimate(args['notion_info_list'], args['process_rule'])
+        response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id, args['notion_info_list'], args['process_rule'])
         return response, 200
 
 

api/controllers/console/datasets/datasets.py

@@ -3,13 +3,15 @@ from flask import request
 from flask_login import login_required, current_user
 from flask_restful import Resource, reqparse, fields, marshal, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
-
 import services
 from controllers.console import api
+from controllers.console.app.error import ProviderNotInitializeError
 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.model_factory import ModelFactory
 from libs.helper import TimestampField
 from extensions.ext_database import db
 from models.dataset import DocumentSegment, Document
@@ -98,6 +100,15 @@ 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,
@@ -221,6 +232,7 @@ class DatasetIndexingEstimateApi(Resource):
         parser = reqparse.RequestParser()
         parser.add_argument('info_list', type=dict, required=True, nullable=True, location='json')
         parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
+        parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
         args = parser.parse_args()
         # validate args
         DocumentService.estimate_args_validate(args)
@@ -235,12 +247,26 @@ class DatasetIndexingEstimateApi(Resource):
                 raise NotFound("File not found.")
 
             indexing_runner = IndexingRunner()
-            response = indexing_runner.file_indexing_estimate(file_details, args['process_rule'])
+
+            try:
+                response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
+                                                                  args['process_rule'], args['doc_form'])
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
         elif args['info_list']['data_source_type'] == 'notion_import':
 
             indexing_runner = IndexingRunner()
-            response = indexing_runner.notion_indexing_estimate(args['info_list']['notion_info_list'],
-                                                                args['process_rule'])
+
+            try:
+                response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
+                                                                    args['info_list']['notion_info_list'],
+                                                                    args['process_rule'], args['doc_form'])
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
         else:
             raise ValueError('Data source type not support')
         return response, 200

api/controllers/console/datasets/datasets_document.py

@@ -18,7 +18,9 @@ from controllers.console.datasets.error import DocumentAlreadyFinishedError, Inv
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.indexing_runner import IndexingRunner
-from core.llm.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError, \
+    LLMBadRequestError
+from core.model_providers.model_factory import ModelFactory
 from extensions.ext_redis import redis_client
 from libs.helper import TimestampField
 from extensions.ext_database import db
@@ -60,6 +62,7 @@ document_fields = {
     'display_status': fields.String,
     'word_count': fields.Integer,
     'hit_count': fields.Integer,
+    'doc_form': fields.String,
 }
 
 document_with_segments_fields = {
@@ -86,6 +89,7 @@ document_with_segments_fields = {
     'total_segments': fields.Integer
 }
 
+
 class DocumentResource(Resource):
     def get_document(self, dataset_id: str, document_id: str) -> Document:
         dataset = DatasetService.get_dataset(dataset_id)
@@ -269,6 +273,7 @@ class DatasetDocumentListApi(Resource):
         parser.add_argument('process_rule', type=dict, required=False, location='json')
         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')
         args = parser.parse_args()
 
         if not dataset.indexing_technique and not args['indexing_technique']:
@@ -277,6 +282,15 @@ 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:
@@ -313,8 +327,18 @@ class DatasetInitApi(Resource):
                             nullable=False, location='json')
         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')
         args = parser.parse_args()
 
+        try:
+            ModelFactory.get_embedding_model(
+                tenant_id=current_user.current_tenant_id
+            )
+        except LLMBadRequestError:
+            raise ProviderNotInitializeError(
+                f"No Embedding Model available. Please configure a valid provider "
+                f"in the Settings -> Model Provider.")
+
         # validate args
         DocumentService.document_create_args_validate(args)
 
@@ -380,7 +404,13 @@ class DocumentIndexingEstimateApi(DocumentResource):
 
                 indexing_runner = IndexingRunner()
 
-                response = indexing_runner.file_indexing_estimate([file], data_process_rule_dict)
+                try:
+                    response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, [file],
+                                                                      data_process_rule_dict)
+                except LLMBadRequestError:
+                    raise ProviderNotInitializeError(
+                        f"No Embedding Model available. Please configure a valid provider "
+                        f"in the Settings -> Model Provider.")
 
         return response
 
@@ -441,12 +471,24 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
                 raise NotFound("File not found.")
 
             indexing_runner = IndexingRunner()
-            response = indexing_runner.file_indexing_estimate(file_details, data_process_rule_dict)
+            try:
+                response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
+                                                                  data_process_rule_dict)
+            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:
 
             indexing_runner = IndexingRunner()
-            response = indexing_runner.notion_indexing_estimate(info_list,
-                                                                data_process_rule_dict)
+            try:
+                response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
+                                                                    info_list,
+                                                                    data_process_rule_dict)
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
         else:
             raise ValueError('Data source type not support')
         return response
@@ -488,6 +530,8 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
                                                           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, self.document_status_fields))
         data = {
             'data': documents_status
@@ -583,7 +627,8 @@ class DocumentDetailApi(DocumentResource):
                 'segment_count': document.segment_count,
                 'average_segment_length': document.average_segment_length,
                 'hit_count': document.hit_count,
-                'display_status': document.display_status
+                'display_status': document.display_status,
+                'doc_form': document.doc_form
             }
         else:
             process_rules = DatasetService.get_process_rules(dataset_id)
@@ -614,7 +659,8 @@ class DocumentDetailApi(DocumentResource):
                 'segment_count': document.segment_count,
                 'average_segment_length': document.average_segment_length,
                 'hit_count': document.hit_count,
-                'display_status': document.display_status
+                'display_status': document.display_status,
+                'doc_form': document.doc_form
             }
 
         return response, 200

api/controllers/console/datasets/datasets_segments.py

@@ -15,8 +15,8 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 from libs.helper import TimestampField
-from services.dataset_service import DatasetService, DocumentService
-from tasks.add_segment_to_index_task import add_segment_to_index_task
+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
 
 segment_fields = {
@@ -24,6 +24,7 @@ segment_fields = {
     '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),
@@ -125,6 +126,7 @@ class DatasetDocumentSegmentListApi(Resource):
 
         return {
             'data': marshal(segments, segment_fields),
+            'doc_form': document.doc_form,
             'has_more': has_more,
             'limit': limit,
             'total': total
@@ -180,7 +182,7 @@ class DatasetDocumentSegmentApi(Resource):
             # Set cache to prevent indexing the same segment multiple times
             redis_client.setex(indexing_cache_key, 600, 1)
 
-            add_segment_to_index_task.delay(segment.id)
+            enable_segment_to_index_task.delay(segment.id)
 
             return {'result': 'success'}, 200
         elif action == "disable":
@@ -202,7 +204,91 @@ class DatasetDocumentSegmentApi(Resource):
             raise InvalidActionError()
 
 
+class DatasetDocumentSegmentAddApi(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.')
+        # 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))
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('content', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('answer', type=str, required=False, nullable=True, location='json')
+        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)
+        return {
+            'data': marshal(segment, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+
+class DatasetDocumentSegmentUpdateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def patch(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 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))
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('content', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('answer', type=str, required=False, nullable=True, location='json')
+        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)
+        return {
+            'data': marshal(segment, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+
 api.add_resource(DatasetDocumentSegmentListApi,
                  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments')
 api.add_resource(DatasetDocumentSegmentApi,
                  '/datasets/<uuid:dataset_id>/segments/<uuid:segment_id>/<string:action>')
+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/controllers/console/datasets/hit_testing.py

@@ -11,7 +11,7 @@ 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.llm.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
+from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import TimestampField
 from services.dataset_service import DatasetService
 from services.hit_testing_service import HitTestingService
@@ -28,6 +28,7 @@ segment_fields = {
     '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),
@@ -101,6 +102,8 @@ class HitTestingApi(Resource):
             raise ProviderQuotaExceededError()
         except ModelCurrentlyNotSupportError:
             raise ProviderModelCurrentlyNotSupportError()
+        except ValueError as e:
+            raise ValueError(str(e))
         except Exception as e:
             logging.exception("Hit testing failed.")
             raise InternalServerError(str(e))

api/controllers/console/explore/audio.py

@@ -11,7 +11,7 @@ from controllers.console.app.error import AppUnavailableError, ProviderNotInitia
     NoAudioUploadedError, AudioTooLargeError, \
     UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
 from controllers.console.explore.wraps import InstalledAppResource
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from services.audio_service import AudioService
 from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \

api/controllers/console/explore/completion.py

@@ -15,7 +15,7 @@ from controllers.console.app.error import ConversationCompletedError, AppUnavail
 from controllers.console.explore.error import NotCompletionAppError, NotChatAppError
 from controllers.console.explore.wraps import InstalledAppResource
 from core.conversation_message_task import PubHandler
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value
 from services.completion_service import CompletionService

api/controllers/console/explore/conversation.py

@@ -65,7 +65,10 @@ class ConversationApi(InstalledAppResource):
             raise NotChatAppError()
 
         conversation_id = str(c_id)
-        ConversationService.delete(app_model, conversation_id, current_user)
+        try:
+            ConversationService.delete(app_model, conversation_id, current_user)
+        except ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
         WebConversationService.unpin(app_model, conversation_id, current_user)
 
         return {"result": "success"}, 204

api/controllers/console/explore/message.py

@@ -15,7 +15,7 @@ from controllers.console.app.error import AppMoreLikeThisDisabledError, Provider
     ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, CompletionRequestError
 from controllers.console.explore.error import NotCompletionAppError, AppSuggestedQuestionsAfterAnswerDisabledError
 from controllers.console.explore.wraps import InstalledAppResource
-from core.llm.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value, TimestampField
 from services.completion_service import CompletionService

api/controllers/console/explore/parameter.py

@@ -4,6 +4,8 @@ from flask_restful import marshal_with, fields
 from controllers.console import api
 from controllers.console.explore.wraps import InstalledAppResource
 
+from models.model import InstalledApp
+
 
 class AppParameterApi(InstalledAppResource):
     """Resource for app variables."""
@@ -27,7 +29,7 @@ class AppParameterApi(InstalledAppResource):
     }
 
     @marshal_with(parameters_fields)
-    def get(self, installed_app):
+    def get(self, installed_app: InstalledApp):
         """Retrieve app parameters."""
         app_model = installed_app.app
         app_model_config = app_model.app_model_config

api/controllers/console/setup.py

@@ -19,15 +19,16 @@ from .wraps import only_edition_self_hosted
 
 class SetupApi(Resource):
 
-    @only_edition_self_hosted
     def get(self):
-        setup_status = get_setup_status()
-        if setup_status:
-            return {
-                'step': 'finished',
-                'setup_at': setup_status.setup_at.isoformat()
-            }  
-        return {'step': 'not_start'}
+        if current_app.config['EDITION'] == 'SELF_HOSTED':
+            setup_status = get_setup_status()
+            if setup_status:
+                return {
+                    'step': 'finished',
+                    'setup_at': setup_status.setup_at.isoformat()
+                }
+            return {'step': 'not_start'}
+        return {'step': 'finished'}
 
     @only_edition_self_hosted
     def post(self):

api/controllers/console/universal_chat/audio.py

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

api/controllers/console/universal_chat/chat.py

@@ -0,0 +1,138 @@
+import json
+import logging
+from typing import Generator, Union
+
+from flask import Response, stream_with_context
+from flask_login import current_user
+from flask_restful import reqparse
+from werkzeug.exceptions import InternalServerError, NotFound
+
+import services
+from controllers.console import api
+from controllers.console.app.error import ConversationCompletedError, AppUnavailableError, ProviderNotInitializeError, \
+    ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, CompletionRequestError
+from controllers.console.universal_chat.wraps import UniversalChatResource
+from core.conversation_message_task import PubHandler
+from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError, \
+    LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError, LLMRateLimitError, LLMAuthorizationError
+from libs.helper import uuid_value
+from services.completion_service import CompletionService
+
+
+class UniversalChatApi(UniversalChatResource):
+    def post(self, universal_app):
+        app_model = universal_app
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('query', type=str, required=True, location='json')
+        parser.add_argument('conversation_id', type=uuid_value, location='json')
+        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')
+        args = parser.parse_args()
+
+        app_model_config = app_model.app_model_config
+
+        # update app model config
+        args['model_config'] = app_model_config.to_dict()
+        args['model_config']['model']['name'] = args['model']
+        args['model_config']['model']['provider'] = args['provider']
+        args['model_config']['agent_mode']['tools'] = args['tools']
+
+        if not args['model_config']['agent_mode']['tools']:
+            args['model_config']['agent_mode']['tools'] = [
+                {
+                    "current_datetime": {
+                        "enabled": True
+                    }
+                }
+            ]
+        else:
+            args['model_config']['agent_mode']['tools'].append({
+                    "current_datetime": {
+                        "enabled": True
+                    }
+                })
+
+        args['inputs'] = {}
+
+        del args['model']
+        del args['tools']
+
+        try:
+            response = CompletionService.completion(
+                app_model=app_model,
+                user=current_user,
+                args=args,
+                from_source='console',
+                streaming=True,
+                is_model_config_override=True,
+            )
+
+            return compact_response(response)
+        except services.errors.conversation.ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
+        except services.errors.conversation.ConversationCompletedError:
+            raise ConversationCompletedError()
+        except services.errors.app_model_config.AppModelConfigBrokenError:
+            logging.exception("App model config broken.")
+            raise AppUnavailableError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except ValueError as e:
+            raise e
+        except Exception as e:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+
+
+class UniversalChatStopApi(UniversalChatResource):
+    def post(self, universal_app, task_id):
+        PubHandler.stop(current_user, task_id)
+
+        return {'result': 'success'}, 200
+
+
+def compact_response(response: Union[dict | Generator]) -> Response:
+    if isinstance(response, dict):
+        return Response(response=json.dumps(response), status=200, mimetype='application/json')
+    else:
+        def generate() -> Generator:
+            try:
+                for chunk in response:
+                    yield chunk
+            except services.errors.conversation.ConversationNotExistsError:
+                yield "data: " + json.dumps(api.handle_error(NotFound("Conversation Not Exists.")).get_json()) + "\n\n"
+            except services.errors.conversation.ConversationCompletedError:
+                yield "data: " + json.dumps(api.handle_error(ConversationCompletedError()).get_json()) + "\n\n"
+            except services.errors.app_model_config.AppModelConfigBrokenError:
+                logging.exception("App model config broken.")
+                yield "data: " + json.dumps(api.handle_error(AppUnavailableError()).get_json()) + "\n\n"
+            except ProviderTokenNotInitError:
+                yield "data: " + json.dumps(api.handle_error(ProviderNotInitializeError()).get_json()) + "\n\n"
+            except QuotaExceededError:
+                yield "data: " + json.dumps(api.handle_error(ProviderQuotaExceededError()).get_json()) + "\n\n"
+            except ModelCurrentlyNotSupportError:
+                yield "data: " + json.dumps(api.handle_error(ProviderModelCurrentlyNotSupportError()).get_json()) + "\n\n"
+            except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                    LLMRateLimitError, LLMAuthorizationError) as e:
+                yield "data: " + json.dumps(api.handle_error(CompletionRequestError(str(e))).get_json()) + "\n\n"
+            except ValueError as e:
+                yield "data: " + json.dumps(api.handle_error(e).get_json()) + "\n\n"
+            except Exception:
+                logging.exception("internal server error.")
+                yield "data: " + json.dumps(api.handle_error(InternalServerError()).get_json()) + "\n\n"
+
+        return Response(stream_with_context(generate()), status=200,
+                        mimetype='text/event-stream')
+
+
+api.add_resource(UniversalChatApi, '/universal-chat/messages')
+api.add_resource(UniversalChatStopApi, '/universal-chat/messages/<string:task_id>/stop')

api/controllers/console/universal_chat/conversation.py

@@ -0,0 +1,118 @@
+# -*- coding:utf-8 -*-
+from flask_login import current_user
+from flask_restful import fields, reqparse, marshal_with
+from flask_restful.inputs import int_range
+from werkzeug.exceptions import NotFound
+
+from controllers.console import api
+from controllers.console.universal_chat.wraps import UniversalChatResource
+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)
+    def get(self, universal_app):
+        app_model = universal_app
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('last_id', type=uuid_value, location='args')
+        parser.add_argument('limit', type=int_range(1, 100), required=False, default=20, location='args')
+        parser.add_argument('pinned', type=str, choices=['true', 'false', None], location='args')
+        args = parser.parse_args()
+
+        pinned = None
+        if 'pinned' in args and args['pinned'] is not None:
+            pinned = True if args['pinned'] == 'true' else False
+
+        try:
+            return WebConversationService.pagination_by_last_id(
+                app_model=app_model,
+                user=current_user,
+                last_id=args['last_id'],
+                limit=args['limit'],
+                pinned=pinned
+            )
+        except LastConversationNotExistsError:
+            raise NotFound("Last Conversation Not Exists.")
+
+
+class UniversalChatConversationApi(UniversalChatResource):
+    def delete(self, universal_app, c_id):
+        app_model = universal_app
+        conversation_id = str(c_id)
+
+        try:
+            ConversationService.delete(app_model, conversation_id, current_user)
+        except ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
+
+        WebConversationService.unpin(app_model, conversation_id, current_user)
+
+        return {"result": "success"}, 204
+
+
+class UniversalChatConversationRenameApi(UniversalChatResource):
+
+    @marshal_with(conversation_fields)
+    def post(self, universal_app, c_id):
+        app_model = universal_app
+        conversation_id = str(c_id)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=True, location='json')
+        args = parser.parse_args()
+
+        try:
+            return ConversationService.rename(app_model, conversation_id, current_user, args['name'])
+        except ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
+
+
+class UniversalChatConversationPinApi(UniversalChatResource):
+
+    def patch(self, universal_app, c_id):
+        app_model = universal_app
+        conversation_id = str(c_id)
+
+        try:
+            WebConversationService.pin(app_model, conversation_id, current_user)
+        except ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
+
+        return {"result": "success"}
+
+
+class UniversalChatConversationUnPinApi(UniversalChatResource):
+    def patch(self, universal_app, c_id):
+        app_model = universal_app
+        conversation_id = str(c_id)
+        WebConversationService.unpin(app_model, conversation_id, current_user)
+
+        return {"result": "success"}
+
+
+api.add_resource(UniversalChatConversationRenameApi, '/universal-chat/conversations/<uuid:c_id>/name')
+api.add_resource(UniversalChatConversationListApi, '/universal-chat/conversations')
+api.add_resource(UniversalChatConversationApi, '/universal-chat/conversations/<uuid:c_id>')
+api.add_resource(UniversalChatConversationPinApi, '/universal-chat/conversations/<uuid:c_id>/pin')
+api.add_resource(UniversalChatConversationUnPinApi, '/universal-chat/conversations/<uuid:c_id>/unpin')

api/controllers/console/universal_chat/message.py

@@ -0,0 +1,127 @@
+# -*- coding:utf-8 -*-
+import logging
+
+from flask_login import current_user
+from flask_restful import reqparse, fields, marshal_with
+from flask_restful.inputs import int_range
+from werkzeug.exceptions import NotFound, InternalServerError
+
+import services
+from controllers.console import api
+from controllers.console.app.error import ProviderNotInitializeError, \
+    ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, CompletionRequestError
+from controllers.console.explore.error import AppSuggestedQuestionsAfterAnswerDisabledError
+from controllers.console.universal_chat.wraps import UniversalChatResource
+from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
+    ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
+from libs.helper import uuid_value, TimestampField
+from services.errors.conversation import ConversationNotExistsError
+from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError
+from services.message_service import MessageService
+
+
+class UniversalChatMessageListApi(UniversalChatResource):
+    feedback_fields = {
+        'rating': fields.String
+    }
+
+    agent_thought_fields = {
+        'id': fields.String,
+        'chain_id': fields.String,
+        'message_id': fields.String,
+        'position': fields.Integer,
+        'thought': fields.String,
+        'tool': fields.String,
+        'tool_input': fields.String,
+        'created_at': TimestampField
+    }
+
+    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,
+        'agent_thoughts': fields.List(fields.Nested(agent_thought_fields))
+    }
+
+    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, universal_app):
+        app_model = universal_app
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('conversation_id', required=True, type=uuid_value, location='args')
+        parser.add_argument('first_id', type=uuid_value, location='args')
+        parser.add_argument('limit', type=int_range(1, 100), required=False, default=20, location='args')
+        args = parser.parse_args()
+
+        try:
+            return MessageService.pagination_by_first_id(app_model, current_user,
+                                                     args['conversation_id'], args['first_id'], args['limit'])
+        except services.errors.conversation.ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
+        except services.errors.message.FirstMessageNotExistsError:
+            raise NotFound("First Message Not Exists.")
+
+
+class UniversalChatMessageFeedbackApi(UniversalChatResource):
+    def post(self, universal_app, message_id):
+        app_model = universal_app
+        message_id = str(message_id)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('rating', type=str, choices=['like', 'dislike', None], location='json')
+        args = parser.parse_args()
+
+        try:
+            MessageService.create_feedback(app_model, message_id, current_user, args['rating'])
+        except services.errors.message.MessageNotExistsError:
+            raise NotFound("Message Not Exists.")
+
+        return {'result': 'success'}
+
+
+class UniversalChatMessageSuggestedQuestionApi(UniversalChatResource):
+    def get(self, universal_app, message_id):
+        app_model = universal_app
+        message_id = str(message_id)
+
+        try:
+            questions = MessageService.get_suggested_questions_after_answer(
+                app_model=app_model,
+                user=current_user,
+                message_id=message_id
+            )
+        except MessageNotExistsError:
+            raise NotFound("Message not found")
+        except ConversationNotExistsError:
+            raise NotFound("Conversation not found")
+        except SuggestedQuestionsAfterAnswerDisabledError:
+            raise AppSuggestedQuestionsAfterAnswerDisabledError()
+        except ProviderTokenNotInitError:
+            raise ProviderNotInitializeError()
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except (LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError,
+                LLMRateLimitError, LLMAuthorizationError) as e:
+            raise CompletionRequestError(str(e))
+        except Exception:
+            logging.exception("internal server error.")
+            raise InternalServerError()
+
+        return {'data': questions}
+
+
+api.add_resource(UniversalChatMessageListApi, '/universal-chat/messages')
+api.add_resource(UniversalChatMessageFeedbackApi, '/universal-chat/messages/<uuid:message_id>/feedbacks')
+api.add_resource(UniversalChatMessageSuggestedQuestionApi, '/universal-chat/messages/<uuid:message_id>/suggested-questions')

api/controllers/console/universal_chat/parameter.py

@@ -0,0 +1,33 @@
+# -*- coding:utf-8 -*-
+from flask_restful import marshal_with, fields
+
+from controllers.console import api
+from controllers.console.universal_chat.wraps import UniversalChatResource
+
+from models.model import App
+
+
+class UniversalChatParameterApi(UniversalChatResource):
+    """Resource for app variables."""
+    parameters_fields = {
+        'opening_statement': fields.String,
+        'suggested_questions': fields.Raw,
+        'suggested_questions_after_answer': fields.Raw,
+        'speech_to_text': fields.Raw,
+    }
+
+    @marshal_with(parameters_fields)
+    def get(self, universal_app: App):
+        """Retrieve app parameters."""
+        app_model = universal_app
+        app_model_config = app_model.app_model_config
+
+        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,
+        }
+
+
+api.add_resource(UniversalChatParameterApi, '/universal-chat/parameters')

api/controllers/console/universal_chat/wraps.py

@@ -0,0 +1,84 @@
+import json
+from functools import wraps
+
+from flask_login import login_required, current_user
+from flask_restful import Resource
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from extensions.ext_database import db
+from models.model import App, AppModelConfig
+
+
+def universal_chat_app_required(view=None):
+    def decorator(view):
+        @wraps(view)
+        def decorated(*args, **kwargs):
+            # get universal chat app
+            universal_app = db.session.query(App).filter(
+                App.tenant_id == current_user.current_tenant_id,
+                App.is_universal == True
+            ).first()
+
+            if universal_app is None:
+                # create universal app if not exists
+                universal_app = App(
+                    tenant_id=current_user.current_tenant_id,
+                    name='Universal Chat',
+                    mode='chat',
+                    is_universal=True,
+                    icon='',
+                    icon_background='',
+                    api_rpm=0,
+                    api_rph=0,
+                    enable_site=False,
+                    enable_api=False,
+                    status='normal'
+                )
+
+                db.session.add(universal_app)
+                db.session.flush()
+
+                app_model_config = AppModelConfig(
+                    provider="",
+                    model_id="",
+                    configs={},
+                    opening_statement='',
+                    suggested_questions=json.dumps([]),
+                    suggested_questions_after_answer=json.dumps({'enabled': True}),
+                    speech_to_text=json.dumps({'enabled': True}),
+                    more_like_this=None,
+                    sensitive_word_avoidance=None,
+                    model=json.dumps({
+                        "provider": "openai",
+                        "name": "gpt-3.5-turbo-16k",
+                        "completion_params": {
+                            "max_tokens": 800,
+                            "temperature": 0.8,
+                            "top_p": 1,
+                            "presence_penalty": 0,
+                            "frequency_penalty": 0
+                        }
+                    }),
+                    user_input_form=json.dumps([]),
+                    pre_prompt='',
+                    agent_mode=json.dumps({"enabled": True, "strategy": "function_call", "tools": []}),
+                )
+
+                app_model_config.app_id = universal_app.id
+                db.session.add(app_model_config)
+                db.session.flush()
+
+                universal_app.app_model_config_id = app_model_config.id
+                db.session.commit()
+
+            return view(universal_app, *args, **kwargs)
+        return decorated
+
+    if view:
+        return decorator(view)
+    return decorator
+
+
+class UniversalChatResource(Resource):
+    # must be reversed if there are multiple decorators
+    method_decorators = [universal_chat_app_required, account_initialization_required, login_required, setup_required]

api/controllers/console/webhook/stripe.py

@@ -0,0 +1,53 @@
+import logging
+
+import stripe
+from flask import request, current_app
+from flask_restful import Resource
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import only_edition_cloud
+from services.provider_checkout_service import ProviderCheckoutService
+
+
+class StripeWebhookApi(Resource):
+    @setup_required
+    @only_edition_cloud
+    def post(self):
+        payload = request.data
+        sig_header = request.headers.get('STRIPE_SIGNATURE')
+        webhook_secret = current_app.config.get('STRIPE_WEBHOOK_SECRET')
+
+        try:
+            event = stripe.Webhook.construct_event(
+                payload, sig_header, webhook_secret
+            )
+        except ValueError as e:
+            # Invalid payload
+            return 'Invalid payload', 400
+        except stripe.error.SignatureVerificationError as e:
+            # Invalid signature
+            return 'Invalid signature', 400
+
+        # Handle the checkout.session.completed event
+        if event['type'] == 'checkout.session.completed':
+            logging.debug(event['data']['object']['id'])
+            logging.debug(event['data']['object']['amount_subtotal'])
+            logging.debug(event['data']['object']['currency'])
+            logging.debug(event['data']['object']['payment_intent'])
+            logging.debug(event['data']['object']['payment_status'])
+            logging.debug(event['data']['object']['metadata'])
+
+            # Fulfill the purchase...
+            provider_checkout_service = ProviderCheckoutService()
+
+            try:
+                provider_checkout_service.fulfill_provider_order(event)
+            except Exception as e:
+                logging.debug(str(e))
+                return 'success', 200
+
+        return 'success', 200
+
+
+api.add_resource(StripeWebhookApi, '/webhook/stripe')

api/controllers/console/workspace/model_providers.py

@@ -0,0 +1,301 @@
+from flask_login import login_required, current_user
+from flask_restful import Resource, reqparse
+from werkzeug.exceptions import Forbidden
+
+from controllers.console import api
+from controllers.console.app.error import ProviderNotInitializeError
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import CredentialsValidateFailedError
+from services.provider_checkout_service import ProviderCheckoutService
+from services.provider_service import ProviderService
+
+
+class ModelProviderListApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        tenant_id = current_user.current_tenant_id
+
+        provider_service = ProviderService()
+        provider_list = provider_service.get_provider_list(tenant_id)
+
+        return provider_list
+
+
+class ModelProviderValidateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('config', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+
+        result = True
+        error = None
+
+        try:
+            provider_service.custom_provider_config_validate(
+                provider_name=provider_name,
+                config=args['config']
+            )
+        except CredentialsValidateFailedError as ex:
+            result = False
+            error = str(ex)
+
+        response = {'result': 'success' if result else 'error'}
+
+        if not result:
+            response['error'] = error
+
+        return response
+
+
+class ModelProviderUpdateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('config', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+
+        try:
+            provider_service.save_custom_provider_config(
+                tenant_id=current_user.current_tenant_id,
+                provider_name=provider_name,
+                config=args['config']
+            )
+        except CredentialsValidateFailedError as ex:
+            raise ValueError(str(ex))
+
+        return {'result': 'success'}, 201
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, provider_name: str):
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        provider_service = ProviderService()
+        provider_service.delete_custom_provider(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name
+        )
+
+        return {'result': 'success'}, 204
+
+
+class ModelProviderModelValidateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_name', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=['text-generation', 'embeddings', 'speech2text'], location='json')
+        parser.add_argument('config', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+
+        result = True
+        error = None
+
+        try:
+            provider_service.custom_provider_model_config_validate(
+                provider_name=provider_name,
+                model_name=args['model_name'],
+                model_type=args['model_type'],
+                config=args['config']
+            )
+        except CredentialsValidateFailedError as ex:
+            result = False
+            error = str(ex)
+
+        response = {'result': 'success' if result else 'error'}
+
+        if not result:
+            response['error'] = error
+
+        return response
+
+
+class ModelProviderModelUpdateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_name', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=['text-generation', 'embeddings', 'speech2text'], location='json')
+        parser.add_argument('config', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+
+        try:
+            provider_service.add_or_save_custom_provider_model_config(
+                tenant_id=current_user.current_tenant_id,
+                provider_name=provider_name,
+                model_name=args['model_name'],
+                model_type=args['model_type'],
+                config=args['config']
+            )
+        except CredentialsValidateFailedError as ex:
+            raise ValueError(str(ex))
+
+        return {'result': 'success'}, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, provider_name: str):
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_name', type=str, required=True, nullable=False, location='args')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=['text-generation', 'embeddings', 'speech2text'], location='args')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+        provider_service.delete_custom_provider_model(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name,
+            model_name=args['model_name'],
+            model_type=args['model_type']
+        )
+
+        return {'result': 'success'}, 204
+
+
+class PreferredProviderTypeUpdateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('preferred_provider_type', type=str, required=True, nullable=False,
+                            choices=['system', 'custom'], location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+        provider_service.switch_preferred_provider(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name,
+            preferred_provider_type=args['preferred_provider_type']
+        )
+
+        return {'result': 'success'}
+
+
+class ModelProviderModelParameterRuleApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_name: str):
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_name', type=str, required=True, nullable=False, location='args')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+
+        try:
+            parameter_rules = provider_service.get_model_parameter_rules(
+                tenant_id=current_user.current_tenant_id,
+                model_provider_name=provider_name,
+                model_name=args['model_name'],
+                model_type='text-generation'
+            )
+        except LLMBadRequestError:
+            raise ProviderNotInitializeError(
+                f"Current Text Generation Model is invalid. Please switch to the available model.")
+
+        rules = {
+            k: {
+                'enabled': v.enabled,
+                'min': v.min,
+                'max': v.max,
+                'default': v.default
+            }
+            for k, v in vars(parameter_rules).items()
+        }
+
+        return rules
+
+
+class ModelProviderPaymentCheckoutUrlApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_name: str):
+        provider_service = ProviderCheckoutService()
+        provider_checkout = provider_service.create_checkout(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name,
+            account=current_user
+        )
+
+        return {
+            'url': provider_checkout.get_checkout_url()
+        }
+
+
+class ModelProviderFreeQuotaSubmitApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider_name: str):
+        provider_service = ProviderService()
+        result = provider_service.free_quota_submit(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name
+        )
+
+        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>')
+api.add_resource(ModelProviderModelValidateApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/models/validate')
+api.add_resource(ModelProviderModelUpdateApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/models')
+api.add_resource(PreferredProviderTypeUpdateApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/preferred-provider-type')
+api.add_resource(ModelProviderModelParameterRuleApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/models/parameter-rules')
+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/controllers/console/workspace/models.py

@@ -0,0 +1,108 @@
+from flask_login import login_required, current_user
+from flask_restful import Resource, reqparse
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.model_providers.model_provider_factory import ModelProviderFactory
+from core.model_providers.models.entity.model_params import ModelType
+from models.provider import ProviderType
+from services.provider_service import ProviderService
+
+
+class DefaultModelApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=['text-generation', 'embeddings', 'speech2text'], location='args')
+        args = parser.parse_args()
+
+        tenant_id = current_user.current_tenant_id
+
+        provider_service = ProviderService()
+        default_model = provider_service.get_default_model_of_model_type(
+            tenant_id=tenant_id,
+            model_type=args['model_type']
+        )
+
+        if not default_model:
+            return None
+
+        model_provider = ModelProviderFactory.get_preferred_model_provider(
+            tenant_id,
+            default_model.provider_name
+        )
+
+        if not model_provider:
+            return {
+                'model_name': default_model.model_name,
+                'model_type': default_model.model_type,
+                'model_provider': {
+                    'provider_name': default_model.provider_name
+                }
+            }
+
+        provider = model_provider.provider
+        rst = {
+            'model_name': default_model.model_name,
+            'model_type': default_model.model_type,
+            'model_provider': {
+                'provider_name': provider.provider_name,
+                'provider_type': provider.provider_type
+            }
+        }
+
+        model_provider_rules = ModelProviderFactory.get_provider_rule(default_model.provider_name)
+        if provider.provider_type == ProviderType.SYSTEM.value:
+            rst['model_provider']['quota_type'] = provider.quota_type
+            rst['model_provider']['quota_unit'] = model_provider_rules['system_config']['quota_unit']
+            rst['model_provider']['quota_limit'] = provider.quota_limit
+            rst['model_provider']['quota_used'] = provider.quota_used
+
+        return rst
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('model_name', type=str, required=True, nullable=False, location='json')
+        parser.add_argument('model_type', type=str, required=True, nullable=False,
+                            choices=['text-generation', 'embeddings', 'speech2text'], location='json')
+        parser.add_argument('provider_name', type=str, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        provider_service = ProviderService()
+        provider_service.update_default_model_of_model_type(
+            tenant_id=current_user.current_tenant_id,
+            model_type=args['model_type'],
+            provider_name=args['provider_name'],
+            model_name=args['model_name']
+        )
+
+        return {'result': 'success'}
+
+
+class ValidModelApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, model_type):
+        ModelType.value_of(model_type)
+
+        provider_service = ProviderService()
+        valid_models = provider_service.get_valid_model_list(
+            tenant_id=current_user.current_tenant_id,
+            model_type=model_type
+        )
+
+        return valid_models
+
+
+api.add_resource(DefaultModelApi, '/workspaces/current/default-model')
+api.add_resource(ValidModelApi, '/workspaces/current/models/model-type/<string:model_type>')

api/controllers/console/workspace/providers.py

@@ -1,20 +1,13 @@
 # -*- coding:utf-8 -*-
-import base64
-import json
-import logging
-
-from flask import current_app
 from flask_login import login_required, current_user
-from flask_restful import Resource, reqparse, abort
+from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 
 from controllers.console import api
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.llm.provider.errors import ValidateFailedError
-from extensions.ext_database import db
-from libs import rsa
-from models.provider import Provider, ProviderType, ProviderName
+from core.model_providers.providers.base import CredentialsValidateFailedError
+from models.provider import ProviderType
 from services.provider_service import ProviderService
 
 
@@ -35,26 +28,26 @@ class ProviderListApi(Resource):
         plaintext, the rest is replaced by * and the last two bits are displayed in plaintext
         """
 
-        ProviderService.init_supported_provider(current_user.current_tenant)
-        providers = Provider.query.filter_by(tenant_id=tenant_id).all()
+        provider_service = ProviderService()
+        provider_info_list = provider_service.get_provider_list(tenant_id)
 
         provider_list = [
             {
-                'provider_name': p.provider_name,
-                'provider_type': p.provider_type,
-                'is_valid': p.is_valid,
-                'last_used': p.last_used,
-                'is_enabled': p.is_enabled,
+                'provider_name': p['provider_name'],
+                'provider_type': p['provider_type'],
+                'is_valid': p['is_valid'],
+                'last_used': p['last_used'],
+                'is_enabled': p['is_valid'],
                 **({
-                       'quota_type': p.quota_type,
-                       'quota_limit': p.quota_limit,
-                       'quota_used': p.quota_used
-                   } if p.provider_type == ProviderType.SYSTEM.value else {}),
-                'token': ProviderService.get_obfuscated_api_key(current_user.current_tenant,
-                                                                ProviderName(p.provider_name), only_custom=True)
-                if p.provider_type == ProviderType.CUSTOM.value else None
+                       'quota_type': p['quota_type'],
+                       'quota_limit': p['quota_limit'],
+                       'quota_used': p['quota_used']
+                   } if p['provider_type'] == ProviderType.SYSTEM.value else {}),
+                'token': (p['config'] if p['provider_name'] != 'openai' else p['config']['openai_api_key'])
+                        if p['config'] else None
             }
-            for p in providers
+            for name, provider_info in provider_info_list.items()
+            for p in provider_info['providers']
         ]
 
         return provider_list
@@ -66,79 +59,28 @@ class ProviderTokenApi(Resource):
     @login_required
     @account_initialization_required
     def post(self, provider):
-        if provider not in [p.value for p in ProviderName]:
-            abort(404)
-
         # 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']:
-            logging.log(logging.ERROR,
-                        f'User {current_user.id} is not authorized to update provider token, current_role is {current_user.current_tenant.current_role}')
             raise Forbidden()
 
         parser = reqparse.RequestParser()
-
-        parser.add_argument('token', type=ProviderService.get_token_type(
-            tenant=current_user.current_tenant,
-            provider_name=ProviderName(provider)
-        ), required=True, nullable=False, location='json')
-
+        parser.add_argument('token', required=True, nullable=False, location='json')
         args = parser.parse_args()
 
-        if args['token']:
-            try:
-                ProviderService.validate_provider_configs(
-                    tenant=current_user.current_tenant,
-                    provider_name=ProviderName(provider),
-                    configs=args['token']
-                )
-                token_is_valid = True
-            except ValidateFailedError as ex:
-                raise ValueError(str(ex))
+        if provider == 'openai':
+            args['token'] = {
+                'openai_api_key': args['token']
+            }
 
-            base64_encrypted_token = ProviderService.get_encrypted_token(
-                tenant=current_user.current_tenant,
-                provider_name=ProviderName(provider),
-                configs=args['token']
+        provider_service = ProviderService()
+        try:
+            provider_service.save_custom_provider_config(
+                tenant_id=current_user.current_tenant_id,
+                provider_name=provider,
+                config=args['token']
             )
-        else:
-            base64_encrypted_token = None
-            token_is_valid = False
-
-        tenant = current_user.current_tenant
-
-        provider_model = db.session.query(Provider).filter(
-                Provider.tenant_id == tenant.id,
-                Provider.provider_name == provider,
-                Provider.provider_type == ProviderType.CUSTOM.value
-            ).first()
-
-        # Only allow updating token for CUSTOM provider type
-        if provider_model:
-            provider_model.encrypted_config = base64_encrypted_token
-            provider_model.is_valid = token_is_valid
-        else:
-            provider_model = Provider(tenant_id=tenant.id, provider_name=provider,
-                                      provider_type=ProviderType.CUSTOM.value,
-                                      encrypted_config=base64_encrypted_token,
-                                      is_valid=token_is_valid)
-            db.session.add(provider_model)
-
-        if provider in [ProviderName.OPENAI.value, ProviderName.AZURE_OPENAI.value] and provider_model.is_valid:
-            other_providers = db.session.query(Provider).filter(
-                Provider.tenant_id == tenant.id,
-                Provider.provider_name.in_([ProviderName.OPENAI.value, ProviderName.AZURE_OPENAI.value]),
-                Provider.provider_name != provider,
-                Provider.provider_type == ProviderType.CUSTOM.value
-            ).all()
-
-            for other_provider in other_providers:
-                other_provider.is_valid = False
-
-        db.session.commit()
-
-        if provider in [ProviderName.AZURE_OPENAI.value, ProviderName.COHERE.value,
-                        ProviderName.HUGGINGFACEHUB.value]:
-            return {'result': 'success', 'warning': 'MOCK: This provider is not supported yet.'}, 201
+        except CredentialsValidateFailedError as ex:
+            raise ValueError(str(ex))
 
         return {'result': 'success'}, 201
 
@@ -149,33 +91,28 @@ class ProviderTokenValidateApi(Resource):
     @login_required
     @account_initialization_required
     def post(self, provider):
-        if provider not in [p.value for p in ProviderName]:
-            abort(404)
-
         parser = reqparse.RequestParser()
-        parser.add_argument('token', type=ProviderService.get_token_type(
-            tenant=current_user.current_tenant,
-            provider_name=ProviderName(provider)
-        ), required=True, nullable=False, location='json')
+        parser.add_argument('token', required=True, nullable=False, location='json')
         args = parser.parse_args()
 
-        # todo: remove this when the provider is supported
-        if provider in [ProviderName.COHERE.value,
-                        ProviderName.HUGGINGFACEHUB.value]:
-            return {'result': 'success', 'warning': 'MOCK: This provider is not supported yet.'}
+        provider_service = ProviderService()
+
+        if provider == 'openai':
+            args['token'] = {
+                'openai_api_key': args['token']
+            }
 
         result = True
         error = None
 
         try:
-            ProviderService.validate_provider_configs(
-                tenant=current_user.current_tenant,
-                provider_name=ProviderName(provider),
-                configs=args['token']
+            provider_service.custom_provider_config_validate(
+                provider_name=provider,
+                config=args['token']
             )
-        except ValidateFailedError as e:
+        except CredentialsValidateFailedError as ex:
             result = False
-            error = str(e)
+            error = str(ex)
 
         response = {'result': 'success' if result else 'error'}
 
@@ -185,91 +122,9 @@ class ProviderTokenValidateApi(Resource):
         return response
 
 
-class ProviderSystemApi(Resource):
-
-    @setup_required
-    @login_required
-    @account_initialization_required
-    def put(self, provider):
-        if provider not in [p.value for p in ProviderName]:
-            abort(404)
-
-        parser = reqparse.RequestParser()
-        parser.add_argument('is_enabled', type=bool, required=True, location='json')
-        args = parser.parse_args()
-
-        tenant = current_user.current_tenant_id
-
-        provider_model = Provider.query.filter_by(tenant_id=tenant.id, provider_name=provider).first()
-
-        if provider_model and provider_model.provider_type == ProviderType.SYSTEM.value:
-            provider_model.is_valid = args['is_enabled']
-            db.session.commit()
-        elif not provider_model:
-            if provider == ProviderName.OPENAI.value:
-                quota_limit = current_app.config['OPENAI_HOSTED_QUOTA_LIMIT']
-            elif provider == ProviderName.ANTHROPIC.value:
-                quota_limit = current_app.config['ANTHROPIC_HOSTED_QUOTA_LIMIT']
-            else:
-                quota_limit = 0
-
-            ProviderService.create_system_provider(
-                tenant,
-                provider,
-                quota_limit,
-                args['is_enabled']
-            )
-        else:
-            abort(403)
-
-        return {'result': 'success'}
-
-    @setup_required
-    @login_required
-    @account_initialization_required
-    def get(self, provider):
-        if provider not in [p.value for p in ProviderName]:
-            abort(404)
-
-        # 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()
-
-        provider_model = db.session.query(Provider).filter(Provider.tenant_id == current_user.current_tenant_id,
-                                                           Provider.provider_name == provider,
-                                                           Provider.provider_type == ProviderType.SYSTEM.value).first()
-
-        system_model = None
-        if provider_model:
-            system_model = {
-                'result': 'success',
-                'provider': {
-                    'provider_name': provider_model.provider_name,
-                    'provider_type': provider_model.provider_type,
-                    'is_valid': provider_model.is_valid,
-                    'last_used': provider_model.last_used,
-                    'is_enabled': provider_model.is_enabled,
-                    'quota_type': provider_model.quota_type,
-                    'quota_limit': provider_model.quota_limit,
-                    'quota_used': provider_model.quota_used
-                }
-            }
-        else:
-            abort(404)
-
-        return system_model
-
-
-api.add_resource(ProviderTokenApi, '/providers/<provider>/token',
-                 endpoint='current_providers_token')  # Deprecated
-api.add_resource(ProviderTokenValidateApi, '/providers/<provider>/token-validate',
-                 endpoint='current_providers_token_validate')  # Deprecated
-
 api.add_resource(ProviderTokenApi, '/workspaces/current/providers/<provider>/token',
                  endpoint='workspaces_current_providers_token')  # PUT for updating provider token
 api.add_resource(ProviderTokenValidateApi, '/workspaces/current/providers/<provider>/token-validate',
                  endpoint='workspaces_current_providers_token_validate')  # POST for validating provider token
 
 api.add_resource(ProviderListApi, '/workspaces/current/providers')  # GET for getting providers list
-api.add_resource(ProviderSystemApi, '/workspaces/current/providers/<provider>/system',
-                 endpoint='workspaces_current_providers_system')  # GET for getting provider quota, PUT for updating provider status

api/controllers/console/workspace/tool_providers.py

@@ -0,0 +1,136 @@
+import json
+
+from flask_login import login_required, current_user
+from flask_restful import Resource, abort, reqparse
+from werkzeug.exceptions import Forbidden
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.tool.provider.errors import ToolValidateFailedError
+from core.tool.provider.tool_provider_service import ToolProviderService
+from extensions.ext_database import db
+from models.tool import ToolProvider, ToolProviderName
+
+
+class ToolProviderListApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        tenant_id = current_user.current_tenant_id
+
+        tool_credential_dict = {}
+        for tool_name in ToolProviderName:
+            tool_credential_dict[tool_name.value] = {
+                'tool_name': tool_name.value,
+                'is_enabled': False,
+                'credentials': None
+            }
+
+        tool_providers = db.session.query(ToolProvider).filter(ToolProvider.tenant_id == tenant_id).all()
+
+        for p in tool_providers:
+            if p.is_enabled:
+                tool_credential_dict[p.tool_name] = {
+                    'tool_name': p.tool_name,
+                    'is_enabled': p.is_enabled,
+                    'credentials': ToolProviderService(tenant_id, p.tool_name).get_credentials(obfuscated=True)
+                }
+
+        return list(tool_credential_dict.values())
+
+
+class ToolProviderCredentialsApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider):
+        if provider not in [p.value for p in ToolProviderName]:
+            abort(404)
+
+        # 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(f'User {current_user.id} is not authorized to update provider token, '
+                            f'current_role is {current_user.current_tenant.current_role}')
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        tenant_id = current_user.current_tenant_id
+
+        tool_provider_service = ToolProviderService(tenant_id, provider)
+
+        try:
+            tool_provider_service.credentials_validate(args['credentials'])
+        except ToolValidateFailedError as ex:
+            raise ValueError(str(ex))
+
+        encrypted_credentials = json.dumps(tool_provider_service.encrypt_credentials(args['credentials']))
+
+        tenant = current_user.current_tenant
+
+        tool_provider_model = db.session.query(ToolProvider).filter(
+                ToolProvider.tenant_id == tenant.id,
+                ToolProvider.tool_name == provider,
+            ).first()
+
+        # Only allow updating token for CUSTOM provider type
+        if tool_provider_model:
+            tool_provider_model.encrypted_credentials = encrypted_credentials
+            tool_provider_model.is_enabled = True
+        else:
+            tool_provider_model = ToolProvider(
+                tenant_id=tenant.id,
+                tool_name=provider,
+                encrypted_credentials=encrypted_credentials,
+                is_enabled=True
+            )
+            db.session.add(tool_provider_model)
+
+        db.session.commit()
+
+        return {'result': 'success'}, 201
+
+
+class ToolProviderCredentialsValidateApi(Resource):
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider):
+        if provider not in [p.value for p in ToolProviderName]:
+            abort(404)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+        args = parser.parse_args()
+
+        result = True
+        error = None
+
+        tenant_id = current_user.current_tenant_id
+
+        tool_provider_service = ToolProviderService(tenant_id, provider)
+
+        try:
+            tool_provider_service.credentials_validate(args['credentials'])
+        except ToolValidateFailedError as ex:
+            result = False
+            error = str(ex)
+
+        response = {'result': 'success' if result else 'error'}
+
+        if not result:
+            response['error'] = error
+
+        return response
+
+
+api.add_resource(ToolProviderListApi, '/workspaces/current/tool-providers')
+api.add_resource(ToolProviderCredentialsApi, '/workspaces/current/tool-providers/<provider>/credentials')
+api.add_resource(ToolProviderCredentialsValidateApi,
+                 '/workspaces/current/tool-providers/<provider>/credentials-validate')

api/controllers/console/workspace/workspace.py

@@ -30,7 +30,7 @@ tenant_fields = {
     'created_at': TimestampField,
     'role': fields.String,
     'providers': fields.List(fields.Nested(provider_fields)),
-    'in_trail': fields.Boolean,
+    'in_trial': fields.Boolean,
     'trial_end_reason': fields.String,
 }
 

api/controllers/service_api/app/app.py

@@ -4,6 +4,8 @@ from flask_restful import fields, marshal_with
 from controllers.service_api import api
 from controllers.service_api.wraps import AppApiResource
 
+from models.model import App
+
 
 class AppParameterApi(AppApiResource):
     """Resource for app variables."""
@@ -28,7 +30,7 @@ class AppParameterApi(AppApiResource):
     }
 
     @marshal_with(parameters_fields)
-    def get(self, app_model, end_user):
+    def get(self, app_model: App, end_user):
         """Retrieve app parameters."""
         app_model_config = app_model.app_model_config
 

api/controllers/service_api/app/audio.py

@@ -9,7 +9,7 @@ from controllers.service_api.app.error import AppUnavailableError, ProviderNotIn
     ProviderModelCurrentlyNotSupportError, NoAudioUploadedError, AudioTooLargeError, UnsupportedAudioTypeError, \
     ProviderNotSupportSpeechToTextError
 from controllers.service_api.wraps import AppApiResource
-from core.llm.error import LLMBadRequestError, LLMAuthorizationError, LLMAPIUnavailableError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAuthorizationError, LLMAPIUnavailableError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from models.model import App, AppModelConfig
 from services.audio_service import AudioService

api/controllers/service_api/app/completion.py

@@ -14,7 +14,7 @@ from controllers.service_api.app.error import AppUnavailableError, ProviderNotIn
     ProviderModelCurrentlyNotSupportError
 from controllers.service_api.wraps import AppApiResource
 from core.conversation_message_task import PubHandler
-from core.llm.error import LLMBadRequestError, LLMAuthorizationError, LLMAPIUnavailableError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAuthorizationError, LLMAPIUnavailableError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value
 from services.completion_service import CompletionService

api/controllers/service_api/app/conversation.py

@@ -64,9 +64,9 @@ class ConversationDetailApi(AppApiResource):
 
         try:
             ConversationService.delete(app_model, conversation_id, end_user)
-            return {"result": "success"}
         except services.errors.conversation.ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")
+        return {"result": "success"}, 204
 
 class ConversationRenameApi(AppApiResource):
 

api/controllers/service_api/dataset/document.py

@@ -11,7 +11,7 @@ from controllers.service_api.app.error import ProviderNotInitializeError
 from controllers.service_api.dataset.error import ArchivedDocumentImmutableError, DocumentIndexingError, \
     DatasetNotInitedError
 from controllers.service_api.wraps import DatasetApiResource
-from core.llm.error import ProviderTokenNotInitError
+from core.model_providers.error import ProviderTokenNotInitError
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models.model import UploadFile

api/controllers/web/app.py

@@ -4,6 +4,8 @@ from flask_restful import marshal_with, fields
 from controllers.web import api
 from controllers.web.wraps import WebApiResource
 
+from models.model import App
+
 
 class AppParameterApi(WebApiResource):
     """Resource for app variables."""
@@ -27,7 +29,7 @@ class AppParameterApi(WebApiResource):
     }
 
     @marshal_with(parameters_fields)
-    def get(self, app_model, end_user):
+    def get(self, app_model: App, end_user):
         """Retrieve app parameters."""
         app_model_config = app_model.app_model_config
 

api/controllers/web/audio.py

@@ -10,7 +10,7 @@ from controllers.web.error import AppUnavailableError, ProviderNotInitializeErro
     ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError, NoAudioUploadedError, AudioTooLargeError, \
     UnsupportedAudioTypeError, ProviderNotSupportSpeechToTextError
 from controllers.web.wraps import WebApiResource
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from services.audio_service import AudioService
 from services.errors.audio import NoAudioUploadedServiceError, AudioTooLargeServiceError, \

api/controllers/web/completion.py

@@ -14,7 +14,7 @@ from controllers.web.error import AppUnavailableError, ConversationCompletedErro
     ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError
 from controllers.web.wraps import WebApiResource
 from core.conversation_message_task import PubHandler
-from core.llm.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMBadRequestError, LLMAPIUnavailableError, LLMAuthorizationError, LLMAPIConnectionError, \
     LLMRateLimitError, ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value
 from services.completion_service import CompletionService

api/controllers/web/conversation.py

@@ -62,7 +62,10 @@ class ConversationApi(WebApiResource):
             raise NotChatAppError()
 
         conversation_id = str(c_id)
-        ConversationService.delete(app_model, conversation_id, end_user)
+        try:
+            ConversationService.delete(app_model, conversation_id, end_user)
+        except ConversationNotExistsError:
+            raise NotFound("Conversation Not Exists.")
         WebConversationService.unpin(app_model, conversation_id, end_user)
 
         return {"result": "success"}, 204

api/controllers/web/message.py

@@ -14,7 +14,7 @@ from controllers.web.error import NotChatAppError, CompletionRequestError, Provi
     AppMoreLikeThisDisabledError, NotCompletionAppError, AppSuggestedQuestionsAfterAnswerDisabledError, \
     ProviderQuotaExceededError, ProviderModelCurrentlyNotSupportError
 from controllers.web.wraps import WebApiResource
-from core.llm.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
+from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
 from libs.helper import uuid_value, TimestampField
 from services.completion_service import CompletionService

api/controllers/web/passport.py

@@ -11,13 +11,13 @@ from libs.passport import PassportService
 class PassportResource(Resource):
     """Base resource for passport."""
     def get(self):
-        app_id = request.headers.get('X-App-Code')
-        if app_id is None:
+        app_code = request.headers.get('X-App-Code')
+        if app_code is None:
             raise Unauthorized('X-App-Code header is missing.')
 
         # get site from db and check if it is normal
         site = db.session.query(Site).filter(
-            Site.code == app_id,
+            Site.code == app_code,
             Site.status == 'normal'
         ).first()
         if not site:
@@ -41,6 +41,7 @@ class PassportResource(Resource):
             "iss": site.app_id,
             'sub': 'Web API Passport',
             'app_id': site.app_id,
+            'app_code': app_code,
             'end_user_id': end_user.id,
         }
 

api/controllers/web/wraps.py

@@ -6,7 +6,7 @@ from flask_restful import Resource
 from werkzeug.exceptions import NotFound, Unauthorized
 
 from extensions.ext_database import db
-from models.model import App, EndUser
+from models.model import App, EndUser, Site
 from libs.passport import PassportService
 
 def validate_jwt_token(view=None):
@@ -35,9 +35,15 @@ def decode_jwt_token():
     if auth_scheme != 'bearer':
         raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
     decoded = PassportService().verify(tk)
+    app_code = decoded.get('app_code')
     app_model = db.session.query(App).filter(App.id == decoded['app_id']).first()
+    site = db.session.query(Site).filter(Site.code == app_code).first()
     if not app_model:
         raise NotFound()
+    if not app_code and not site:
+        raise Unauthorized('Site URL is no longer valid.')
+    if app_model.enable_site is False:
+        raise Unauthorized('Site is disabled.')
     end_user = db.session.query(EndUser).filter(EndUser.id == decoded['end_user_id']).first()
     if not end_user:
         raise NotFound()

api/core/__init__.py

@@ -1,36 +0,0 @@
-import os
-from typing import Optional
-
-import langchain
-from flask import Flask
-from pydantic import BaseModel
-
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.prompt.prompt_template import OneLineFormatter
-
-
-class HostedOpenAICredential(BaseModel):
-    api_key: str
-
-
-class HostedAnthropicCredential(BaseModel):
-    api_key: str
-
-
-class HostedLLMCredentials(BaseModel):
-    openai: Optional[HostedOpenAICredential] = None
-    anthropic: Optional[HostedAnthropicCredential] = None
-
-
-hosted_llm_credentials = HostedLLMCredentials()
-
-
-def init_app(app: Flask):
-    if os.environ.get("DEBUG") and os.environ.get("DEBUG").lower() == 'true':
-        langchain.verbose = True
-
-    if app.config.get("OPENAI_API_KEY"):
-        hosted_llm_credentials.openai = HostedOpenAICredential(api_key=app.config.get("OPENAI_API_KEY"))
-
-    if app.config.get("ANTHROPIC_API_KEY"):
-        hosted_llm_credentials.anthropic = HostedAnthropicCredential(api_key=app.config.get("ANTHROPIC_API_KEY"))

api/core/agent/agent/calc_token_mixin.py

@@ -0,0 +1,31 @@
+from typing import List
+
+from langchain.schema import BaseMessage
+
+from core.model_providers.models.entity.message import to_prompt_messages
+from core.model_providers.models.llm.base import BaseLLM
+
+
+class CalcTokenMixin:
+
+    def get_num_tokens_from_messages(self, model_instance: BaseLLM, messages: List[BaseMessage], **kwargs) -> int:
+        return model_instance.get_num_tokens(to_prompt_messages(messages))
+
+    def get_message_rest_tokens(self, model_instance: BaseLLM, messages: List[BaseMessage], **kwargs) -> int:
+        """
+        Got the rest tokens available for the model after excluding messages tokens and completion max tokens
+
+        :param llm:
+        :param messages:
+        :return:
+        """
+        llm_max_tokens = model_instance.model_rules.max_tokens.max
+        completion_max_tokens = model_instance.model_kwargs.max_tokens
+        used_tokens = self.get_num_tokens_from_messages(model_instance, messages, **kwargs)
+        rest_tokens = llm_max_tokens - completion_max_tokens - used_tokens
+
+        return rest_tokens
+
+
+class ExceededLLMTokensLimitError(Exception):
+    pass

api/core/agent/agent/multi_dataset_router_agent.py

@@ -0,0 +1,91 @@
+from typing import Tuple, List, Any, Union, Sequence, Optional, cast
+
+from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, SystemMessage
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.tools import BaseTool
+
+from core.model_providers.models.llm.base import BaseLLM
+from core.tool.dataset_retriever_tool import DatasetRetrieverTool
+
+
+class MultiDatasetRouterAgent(OpenAIFunctionsAgent):
+    """
+    An Multi Dataset Retrieve Agent driven by Router.
+    """
+    model_instance: BaseLLM
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        return True
+
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        if len(self.tools) == 0:
+            return AgentFinish(return_values={"output": ''}, log='')
+        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']})
+            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)
+
+    async def aplan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        raise NotImplementedError()
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseSingleActionAgent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=system_message,
+            **kwargs,
+        )

api/core/agent/agent/openai_function_call.py

@@ -0,0 +1,113 @@
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+from langchain.agents import OpenAIFunctionsAgent, BaseSingleActionAgent
+from langchain.agents.openai_functions_agent.base import _parse_ai_message, \
+    _format_intermediate_steps
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, SystemMessage
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.tools import BaseTool
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError
+from core.agent.agent.openai_function_call_summarize_mixin import OpenAIFunctionCallSummarizeMixin
+
+
+class AutoSummarizingOpenAIFunctionCallAgent(OpenAIFunctionsAgent, OpenAIFunctionCallSummarizeMixin):
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseSingleActionAgent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=cls.get_system_message(),
+            **kwargs,
+        )
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        original_max_tokens = self.llm.max_tokens
+        self.llm.max_tokens = 15
+
+        prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
+        messages = prompt.to_messages()
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=None
+        )
+
+        function_call = predicted_message.additional_kwargs.get("function_call", {})
+
+        self.llm.max_tokens = original_max_tokens
+
+        return True if function_call else False
+
+    def plan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
+        selected_inputs = {
+            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
+        }
+        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
+        prompt = self.prompt.format_prompt(**full_inputs)
+        messages = prompt.to_messages()
+
+        # summarize messages if rest_tokens < 0
+        try:
+            messages = self.summarize_messages_if_needed(messages, functions=self.functions)
+        except ExceededLLMTokensLimitError as e:
+            return AgentFinish(return_values={"output": str(e)}, log=str(e))
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=callbacks
+        )
+        agent_decision = _parse_ai_message(predicted_message)
+        return agent_decision
+
+    @classmethod
+    def get_system_message(cls):
+        return SystemMessage(content="You are a helpful AI assistant.\n"
+                                     "The current date or current time you know is wrong.\n"
+                                     "Respond directly if appropriate.")
+
+    def return_stopped_response(
+            self,
+            early_stopping_method: str,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            **kwargs: Any,
+    ) -> AgentFinish:
+        try:
+            return super().return_stopped_response(early_stopping_method, intermediate_steps, **kwargs)
+        except ValueError:
+            return AgentFinish({"output": "I'm sorry, I don't know how to respond to that."}, "")

api/core/agent/agent/openai_function_call_summarize_mixin.py

@@ -0,0 +1,140 @@
+from typing import cast, List
+
+from langchain.chat_models import ChatOpenAI
+from langchain.chat_models.openai import _convert_message_to_dict
+from langchain.memory.summary import SummarizerMixin
+from langchain.schema import SystemMessage, HumanMessage, BaseMessage, AIMessage
+from langchain.schema.language_model import BaseLanguageModel
+from pydantic import BaseModel
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError, CalcTokenMixin
+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
+    model_instance: BaseLLM
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
+    def summarize_messages_if_needed(self, messages: List[BaseMessage], **kwargs) -> List[BaseMessage]:
+        # calculate rest tokens and summarize previous function observation messages if rest_tokens < 0
+        rest_tokens = self.get_message_rest_tokens(self.model_instance, messages, **kwargs)
+        rest_tokens = rest_tokens - 20  # to deal with the inaccuracy of rest_tokens
+        if rest_tokens >= 0:
+            return messages
+
+        system_message = None
+        human_message = None
+        should_summary_messages = []
+        for message in messages:
+            if isinstance(message, SystemMessage):
+                system_message = message
+            elif isinstance(message, HumanMessage):
+                human_message = message
+            else:
+                should_summary_messages.append(message)
+
+        if len(should_summary_messages) > 2:
+            ai_message = should_summary_messages[-2]
+            function_message = should_summary_messages[-1]
+            should_summary_messages = should_summary_messages[self.moving_summary_index:-2]
+            self.moving_summary_index = len(should_summary_messages)
+        else:
+            error_msg = "Exceeded LLM tokens limit, stopped."
+            raise ExceededLLMTokensLimitError(error_msg)
+
+        new_messages = [system_message, human_message]
+
+        if self.moving_summary_index == 0:
+            should_summary_messages.insert(0, human_message)
+
+        summary_handler = SummarizerMixin(llm=self.summary_llm)
+        self.moving_summary_buffer = summary_handler.predict_new_summary(
+            messages=should_summary_messages,
+            existing_summary=self.moving_summary_buffer
+        )
+
+        new_messages.append(AIMessage(content=self.moving_summary_buffer))
+        new_messages.append(ai_message)
+        new_messages.append(function_message)
+
+        return new_messages
+
+    def get_num_tokens_from_messages(self, llm: BaseLanguageModel, 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)
+        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
+            tokens_per_message = 4
+            # if there's a name, the role is omitted
+            tokens_per_name = -1
+        elif model.startswith("gpt-4"):
+            tokens_per_message = 3
+            tokens_per_name = 1
+        else:
+            raise NotImplementedError(
+                f"get_num_tokens_from_messages() is not presently implemented "
+                f"for model {model}."
+                "See https://github.com/openai/openai-python/blob/main/chatml.md for "
+                "information on how messages are converted to tokens."
+            )
+        num_tokens = 0
+        for m in messages:
+            message = _convert_message_to_dict(m)
+            num_tokens += tokens_per_message
+            for key, value in message.items():
+                if key == "function_call":
+                    for f_key, f_value in value.items():
+                        num_tokens += len(encoding.encode(f_key))
+                        num_tokens += len(encoding.encode(f_value))
+                else:
+                    num_tokens += len(encoding.encode(value))
+
+                if key == "name":
+                    num_tokens += tokens_per_name
+        # every reply is primed with <im_start>assistant
+        num_tokens += 3
+
+        if kwargs.get('functions'):
+            for function in kwargs.get('functions'):
+                num_tokens += len(encoding.encode('name'))
+                num_tokens += len(encoding.encode(function.get("name")))
+                num_tokens += len(encoding.encode('description'))
+                num_tokens += len(encoding.encode(function.get("description")))
+                parameters = function.get("parameters")
+                num_tokens += len(encoding.encode('parameters'))
+                if 'title' in parameters:
+                    num_tokens += len(encoding.encode('title'))
+                    num_tokens += len(encoding.encode(parameters.get("title")))
+                num_tokens += len(encoding.encode('type'))
+                num_tokens += len(encoding.encode(parameters.get("type")))
+                if 'properties' in parameters:
+                    num_tokens += len(encoding.encode('properties'))
+                    for key, value in parameters.get('properties').items():
+                        num_tokens += len(encoding.encode(key))
+                        for field_key, field_value in value.items():
+                            num_tokens += len(encoding.encode(field_key))
+                            if field_key == 'enum':
+                                for enum_field in field_value:
+                                    num_tokens += 3
+                                    num_tokens += len(encoding.encode(enum_field))
+                            else:
+                                num_tokens += len(encoding.encode(field_key))
+                                num_tokens += len(encoding.encode(str(field_value)))
+                if 'required' in parameters:
+                    num_tokens += len(encoding.encode('required'))
+                    for required_field in parameters['required']:
+                        num_tokens += 3
+                        num_tokens += len(encoding.encode(required_field))
+
+        return num_tokens

api/core/agent/agent/openai_multi_function_call.py

@@ -0,0 +1,103 @@
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+from langchain.agents import BaseMultiActionAgent
+from langchain.agents.openai_functions_multi_agent.base import OpenAIMultiFunctionsAgent, _format_intermediate_steps, \
+    _parse_ai_message
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts.chat import BaseMessagePromptTemplate
+from langchain.schema import AgentAction, AgentFinish, SystemMessage
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.tools import BaseTool
+
+from core.agent.agent.calc_token_mixin import ExceededLLMTokensLimitError
+from core.agent.agent.openai_function_call_summarize_mixin import OpenAIFunctionCallSummarizeMixin
+
+
+class AutoSummarizingOpenMultiAIFunctionCallAgent(OpenAIMultiFunctionsAgent, OpenAIFunctionCallSummarizeMixin):
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
+            system_message: Optional[SystemMessage] = SystemMessage(
+                content="You are a helpful AI assistant."
+            ),
+            **kwargs: Any,
+    ) -> BaseMultiActionAgent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            extra_prompt_messages=extra_prompt_messages,
+            system_message=cls.get_system_message(),
+            **kwargs,
+        )
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+
+        :param query:
+        :return:
+        """
+        original_max_tokens = self.llm.max_tokens
+        self.llm.max_tokens = 15
+
+        prompt = self.prompt.format_prompt(input=query, agent_scratchpad=[])
+        messages = prompt.to_messages()
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=None
+        )
+
+        function_call = predicted_message.additional_kwargs.get("function_call", {})
+
+        self.llm.max_tokens = original_max_tokens
+
+        return True if function_call else False
+
+    def plan(
+            self,
+            intermediate_steps: List[Tuple[AgentAction, str]],
+            callbacks: Callbacks = None,
+            **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date, along with observations
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
+        selected_inputs = {
+            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
+        }
+        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
+        prompt = self.prompt.format_prompt(**full_inputs)
+        messages = prompt.to_messages()
+
+        # summarize messages if rest_tokens < 0
+        try:
+            messages = self.summarize_messages_if_needed(messages, functions=self.functions)
+        except ExceededLLMTokensLimitError as e:
+            return AgentFinish(return_values={"output": str(e)}, log=str(e))
+
+        predicted_message = self.llm.predict_messages(
+            messages, functions=self.functions, callbacks=callbacks
+        )
+        agent_decision = _parse_ai_message(predicted_message)
+        return agent_decision
+
+    @classmethod
+    def get_system_message(cls):
+        # get current time
+        return SystemMessage(content="You are a helpful AI assistant.\n"
+                                     "The current date or current time you know is wrong.\n"
+                                     "Respond directly if appropriate.")

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

@@ -0,0 +1,29 @@
+import json
+import re
+from typing import Union
+
+from langchain.agents.structured_chat.output_parser import StructuredChatOutputParser as LCStructuredChatOutputParser, \
+    logger
+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)
+            if action_match is not None:
+                response = json.loads(action_match.group(2).strip(), strict=False)
+                if isinstance(response, list):
+                    # gpt turbo frequently ignores the directive to emit a single action
+                    logger.warning("Got multiple action responses: %s", response)
+                    response = response[0]
+                if response["action"] == "Final Answer":
+                    return AgentFinish({"output": response["action_input"]}, text)
+                else:
+                    return AgentAction(
+                        response["action"], response.get("action_input", {}), text
+                    )
+            else:
+                return AgentFinish({"output": text}, text)
+        except Exception as e:
+            raise OutputParserException(f"Could not parse LLM output: {text}") from e

api/core/agent/agent/structed_multi_dataset_router_agent.py

@@ -0,0 +1,162 @@
+import re
+from typing import List, Tuple, Any, Union, Sequence, Optional, cast
+
+from langchain import BasePromptTemplate
+from langchain.agents import StructuredChatAgent, AgentOutputParser, Agent
+from langchain.agents.structured_chat.base import HUMAN_MESSAGE_TEMPLATE
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+from langchain.tools import BaseTool
+from langchain.agents.structured_chat.prompt import PREFIX, SUFFIX
+
+from core.model_providers.models.llm.base import BaseLLM
+from core.tool.dataset_retriever_tool import DatasetRetrieverTool
+
+FORMAT_INSTRUCTIONS = """Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
+The nouns in the format of "Thought", "Action", "Action Input", "Final Answer" must be expressed in English.
+Valid "action" values: "Final Answer" or {tool_names}
+
+Provide only ONE action per $JSON_BLOB, as shown:
+
+```
+{{{{
+  "action": $TOOL_NAME,
+  "action_input": $INPUT
+}}}}
+```
+
+Follow this format:
+
+Question: input question to answer
+Thought: consider previous and subsequent steps
+Action:
+```
+$JSON_BLOB
+```
+Observation: action result
+... (repeat Thought/Action/Observation N times)
+Thought: I know what to respond
+Action:
+```
+{{{{
+  "action": "Final Answer",
+  "action_input": "Final response to human"
+}}}}
+```"""
+
+
+class StructuredMultiDatasetRouterAgent(StructuredChatAgent):
+    model_instance: BaseLLM
+    dataset_tools: Sequence[BaseTool]
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+        Using the ReACT mode to determine whether an agent is needed is costly,
+        so it's better to just use an Agent for reasoning, which is cheaper.
+
+        :param query:
+        :return:
+        """
+        return True
+
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date,
+                along with observations
+            callbacks: Callbacks to run.
+            **kwargs: User inputs.
+
+        Returns:
+            Action specifying what tool to use.
+        """
+        if len(self.dataset_tools) == 0:
+            return AgentFinish(return_values={"output": ''}, log='')
+        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']})
+            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)
+        except OutputParserException:
+            return AgentFinish({"output": "I'm sorry, the answer of model is invalid, "
+                                          "I don't know how to respond to that."}, "")
+    @classmethod
+    def create_prompt(
+            cls,
+            tools: Sequence[BaseTool],
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+    ) -> BasePromptTemplate:
+        tool_strings = []
+        for tool in tools:
+            args_schema = re.sub("}", "}}}}", re.sub("{", "{{{{", str(tool.args)))
+            tool_strings.append(f"{tool.name}: {tool.description}, args: {args_schema}")
+        formatted_tools = "\n".join(tool_strings)
+        unique_tool_names = set(tool.name for tool in tools)
+        tool_names = ", ".join('"' + name + '"' for name in unique_tool_names)
+        format_instructions = format_instructions.format(tool_names=tool_names)
+        template = "\n\n".join([prefix, formatted_tools, format_instructions, suffix])
+        if input_variables is None:
+            input_variables = ["input", "agent_scratchpad"]
+        _memory_prompts = memory_prompts or []
+        messages = [
+            SystemMessagePromptTemplate.from_template(template),
+            *_memory_prompts,
+            HumanMessagePromptTemplate.from_template(human_message_template),
+        ]
+        return ChatPromptTemplate(input_variables=input_variables, messages=messages)
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            output_parser: Optional[AgentOutputParser] = None,
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+            **kwargs: Any,
+    ) -> Agent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            output_parser=output_parser,
+            prefix=prefix,
+            suffix=suffix,
+            human_message_template=human_message_template,
+            format_instructions=format_instructions,
+            input_variables=input_variables,
+            memory_prompts=memory_prompts,
+            dataset_tools=tools,
+            **kwargs,
+        )

api/core/agent/agent/structured_chat.py

@@ -0,0 +1,193 @@
+import re
+from typing import List, Tuple, Any, Union, Sequence, Optional
+
+from langchain import BasePromptTemplate
+from langchain.agents import StructuredChatAgent, AgentOutputParser, Agent
+from langchain.agents.structured_chat.base import HUMAN_MESSAGE_TEMPLATE
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.base import BaseCallbackManager
+from langchain.callbacks.manager import Callbacks
+from langchain.memory.summary import SummarizerMixin
+from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain.schema import AgentAction, AgentFinish, AIMessage, HumanMessage, OutputParserException
+from langchain.tools import BaseTool
+from langchain.agents.structured_chat.prompt import PREFIX, SUFFIX
+
+from core.agent.agent.calc_token_mixin import CalcTokenMixin, ExceededLLMTokensLimitError
+from core.model_providers.models.llm.base import BaseLLM
+
+FORMAT_INSTRUCTIONS = """Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
+The nouns in the format of "Thought", "Action", "Action Input", "Final Answer" must be expressed in English.
+Valid "action" values: "Final Answer" or {tool_names}
+
+Provide only ONE action per $JSON_BLOB, as shown:
+
+```
+{{{{
+  "action": $TOOL_NAME,
+  "action_input": $INPUT
+}}}}
+```
+
+Follow this format:
+
+Question: input question to answer
+Thought: consider previous and subsequent steps
+Action:
+```
+$JSON_BLOB
+```
+Observation: action result
+... (repeat Thought/Action/Observation N times)
+Thought: I know what to respond
+Action:
+```
+{{{{
+  "action": "Final Answer",
+  "action_input": "Final response to human"
+}}}}
+```"""
+
+
+class AutoSummarizingStructuredChatAgent(StructuredChatAgent, CalcTokenMixin):
+    moving_summary_buffer: str = ""
+    moving_summary_index: int = 0
+    summary_llm: BaseLanguageModel
+    model_instance: BaseLLM
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        arbitrary_types_allowed = True
+
+    def should_use_agent(self, query: str):
+        """
+        return should use agent
+        Using the ReACT mode to determine whether an agent is needed is costly,
+        so it's better to just use an Agent for reasoning, which is cheaper.
+
+        :param query:
+        :return:
+        """
+        return True
+
+    def plan(
+        self,
+        intermediate_steps: List[Tuple[AgentAction, str]],
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> Union[AgentAction, AgentFinish]:
+        """Given input, decided what to do.
+
+        Args:
+            intermediate_steps: Steps the LLM has taken to date,
+                along with observations
+            callbacks: Callbacks to run.
+            **kwargs: User inputs.
+
+        Returns:
+            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()
+
+        rest_tokens = self.get_message_rest_tokens(self.model_instance, messages)
+        if rest_tokens < 0:
+            full_inputs = self.summarize_messages(intermediate_steps, **kwargs)
+
+        full_output = self.llm_chain.predict(callbacks=callbacks, **full_inputs)
+
+        try:
+            return self.output_parser.parse(full_output)
+        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:
+            should_summary_intermediate_steps = intermediate_steps[self.moving_summary_index:-1]
+            should_summary_messages = [AIMessage(content=observation)
+                                       for _, observation in should_summary_intermediate_steps]
+            if self.moving_summary_index == 0:
+                should_summary_messages.insert(0, HumanMessage(content=kwargs.get("input")))
+
+            self.moving_summary_index = len(intermediate_steps)
+        else:
+            error_msg = "Exceeded LLM tokens limit, stopped."
+            raise ExceededLLMTokensLimitError(error_msg)
+
+        summary_handler = SummarizerMixin(llm=self.summary_llm)
+        if self.moving_summary_buffer and 'chat_history' in kwargs:
+            kwargs["chat_history"].pop()
+
+        self.moving_summary_buffer = summary_handler.predict_new_summary(
+            messages=should_summary_messages,
+            existing_summary=self.moving_summary_buffer
+        )
+
+        if 'chat_history' in kwargs:
+            kwargs["chat_history"].append(AIMessage(content=self.moving_summary_buffer))
+
+        return self.get_full_inputs([intermediate_steps[-1]], **kwargs)
+
+    @classmethod
+    def create_prompt(
+            cls,
+            tools: Sequence[BaseTool],
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+    ) -> BasePromptTemplate:
+        tool_strings = []
+        for tool in tools:
+            args_schema = re.sub("}", "}}}}", re.sub("{", "{{{{", str(tool.args)))
+            tool_strings.append(f"{tool.name}: {tool.description}, args: {args_schema}")
+        formatted_tools = "\n".join(tool_strings)
+        tool_names = ", ".join([('"' + tool.name + '"') for tool in tools])
+        format_instructions = format_instructions.format(tool_names=tool_names)
+        template = "\n\n".join([prefix, formatted_tools, format_instructions, suffix])
+        if input_variables is None:
+            input_variables = ["input", "agent_scratchpad"]
+        _memory_prompts = memory_prompts or []
+        messages = [
+            SystemMessagePromptTemplate.from_template(template),
+            *_memory_prompts,
+            HumanMessagePromptTemplate.from_template(human_message_template),
+        ]
+        return ChatPromptTemplate(input_variables=input_variables, messages=messages)
+
+    @classmethod
+    def from_llm_and_tools(
+            cls,
+            llm: BaseLanguageModel,
+            tools: Sequence[BaseTool],
+            callback_manager: Optional[BaseCallbackManager] = None,
+            output_parser: Optional[AgentOutputParser] = None,
+            prefix: str = PREFIX,
+            suffix: str = SUFFIX,
+            human_message_template: str = HUMAN_MESSAGE_TEMPLATE,
+            format_instructions: str = FORMAT_INSTRUCTIONS,
+            input_variables: Optional[List[str]] = None,
+            memory_prompts: Optional[List[BasePromptTemplate]] = None,
+            **kwargs: Any,
+    ) -> Agent:
+        return super().from_llm_and_tools(
+            llm=llm,
+            tools=tools,
+            callback_manager=callback_manager,
+            output_parser=output_parser,
+            prefix=prefix,
+            suffix=suffix,
+            human_message_template=human_message_template,
+            format_instructions=format_instructions,
+            input_variables=input_variables,
+            memory_prompts=memory_prompts,
+            **kwargs,
+        )

api/core/agent/agent_builder.py

@@ -1,86 +0,0 @@
-from typing import Optional
-
-from langchain import LLMChain
-from langchain.agents import ZeroShotAgent, AgentExecutor, ConversationalAgent
-from langchain.callbacks.manager import CallbackManager
-from langchain.memory.chat_memory import BaseChatMemory
-
-from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
-from core.callback_handler.dataset_tool_callback_handler import DatasetToolCallbackHandler
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.llm.llm_builder import LLMBuilder
-
-
-class AgentBuilder:
-    @classmethod
-    def to_agent_chain(cls, tenant_id: str, tools, memory: Optional[BaseChatMemory],
-                       dataset_tool_callback_handler: DatasetToolCallbackHandler,
-                       agent_loop_gather_callback_handler: AgentLoopGatherCallbackHandler):
-        llm = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name=agent_loop_gather_callback_handler.model_name,
-            temperature=0,
-            max_tokens=1024,
-            callbacks=[agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()]
-        )
-
-        for tool in tools:
-            tool.callbacks = [
-                agent_loop_gather_callback_handler,
-                dataset_tool_callback_handler,
-                DifyStdOutCallbackHandler()
-            ]
-
-        prompt = cls.build_agent_prompt_template(
-            tools=tools,
-            memory=memory,
-        )
-
-        agent_llm_chain = LLMChain(
-            llm=llm,
-            prompt=prompt,
-        )
-
-        agent = cls.build_agent(agent_llm_chain=agent_llm_chain, memory=memory)
-
-        agent_callback_manager = CallbackManager(
-            [agent_loop_gather_callback_handler, DifyStdOutCallbackHandler()]
-        )
-
-        agent_chain = AgentExecutor.from_agent_and_tools(
-            tools=tools,
-            agent=agent,
-            memory=memory,
-            callbacks=agent_callback_manager,
-            max_iterations=6,
-            early_stopping_method="generate",
-            # `generate` will continue to complete the last inference after reaching the iteration limit or request time limit
-        )
-
-        return agent_chain
-
-    @classmethod
-    def build_agent_prompt_template(cls, tools, memory: Optional[BaseChatMemory]):
-        if memory:
-            prompt = ConversationalAgent.create_prompt(
-                tools=tools,
-            )
-        else:
-            prompt = ZeroShotAgent.create_prompt(
-                tools=tools,
-            )
-
-        return prompt
-
-    @classmethod
-    def build_agent(cls, agent_llm_chain: LLMChain, memory: Optional[BaseChatMemory]):
-        if memory:
-            agent = ConversationalAgent(
-                llm_chain=agent_llm_chain
-            )
-        else:
-            agent = ZeroShotAgent(
-                llm_chain=agent_llm_chain
-            )
-
-        return agent

api/core/agent/agent_executor.py

@@ -0,0 +1,136 @@
+import enum
+import logging
+from typing import Union, Optional
+
+from langchain.agents import BaseSingleActionAgent, BaseMultiActionAgent
+from langchain.callbacks.manager import Callbacks
+from langchain.memory.chat_memory import BaseChatMemory
+from langchain.tools import BaseTool
+from pydantic import BaseModel, Extra
+
+from core.agent.agent.multi_dataset_router_agent import MultiDatasetRouterAgent
+from core.agent.agent.openai_function_call import AutoSummarizingOpenAIFunctionCallAgent
+from core.agent.agent.openai_multi_function_call import AutoSummarizingOpenMultiAIFunctionCallAgent
+from core.agent.agent.output_parser.structured_chat import StructuredChatOutputParser
+from core.agent.agent.structed_multi_dataset_router_agent import StructuredMultiDatasetRouterAgent
+from core.agent.agent.structured_chat import AutoSummarizingStructuredChatAgent
+from langchain.agents import AgentExecutor as LCAgentExecutor
+
+from core.model_providers.models.llm.base import BaseLLM
+from core.tool.dataset_retriever_tool import DatasetRetrieverTool
+
+
+class PlanningStrategy(str, enum.Enum):
+    ROUTER = 'router'
+    REACT_ROUTER = 'react_router'
+    REACT = 'react'
+    FUNCTION_CALL = 'function_call'
+    MULTI_FUNCTION_CALL = 'multi_function_call'
+
+
+class AgentConfiguration(BaseModel):
+    strategy: PlanningStrategy
+    model_instance: BaseLLM
+    tools: list[BaseTool]
+    summary_model_instance: BaseLLM
+    memory: Optional[BaseChatMemory] = None
+    callbacks: Callbacks = None
+    max_iterations: int = 6
+    max_execution_time: Optional[float] = None
+    early_stopping_method: str = "generate"
+    # `generate` will continue to complete the last inference after reaching the iteration limit or request time limit
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+
+class AgentExecuteResult(BaseModel):
+    strategy: PlanningStrategy
+    output: Optional[str]
+    configuration: AgentConfiguration
+
+
+class AgentExecutor:
+    def __init__(self, configuration: AgentConfiguration):
+        self.configuration = configuration
+        self.agent = self._init_agent()
+
+    def _init_agent(self) -> Union[BaseSingleActionAgent | BaseMultiActionAgent]:
+        if self.configuration.strategy == PlanningStrategy.REACT:
+            agent = AutoSummarizingStructuredChatAgent.from_llm_and_tools(
+                model_instance=self.configuration.model_instance,
+                llm=self.configuration.model_instance.client,
+                tools=self.configuration.tools,
+                output_parser=StructuredChatOutputParser(),
+                summary_llm=self.configuration.summary_model_instance.client,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.FUNCTION_CALL:
+            agent = AutoSummarizingOpenAIFunctionCallAgent.from_llm_and_tools(
+                model_instance=self.configuration.model_instance,
+                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,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.MULTI_FUNCTION_CALL:
+            agent = AutoSummarizingOpenMultiAIFunctionCallAgent.from_llm_and_tools(
+                model_instance=self.configuration.model_instance,
+                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,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.ROUTER:
+            self.configuration.tools = [t for t in self.configuration.tools if isinstance(t, DatasetRetrieverTool)]
+            agent = MultiDatasetRouterAgent.from_llm_and_tools(
+                model_instance=self.configuration.model_instance,
+                llm=self.configuration.model_instance.client,
+                tools=self.configuration.tools,
+                extra_prompt_messages=self.configuration.memory.buffer if self.configuration.memory else None,
+                verbose=True
+            )
+        elif self.configuration.strategy == PlanningStrategy.REACT_ROUTER:
+            self.configuration.tools = [t for t in self.configuration.tools if isinstance(t, DatasetRetrieverTool)]
+            agent = StructuredMultiDatasetRouterAgent.from_llm_and_tools(
+                model_instance=self.configuration.model_instance,
+                llm=self.configuration.model_instance.client,
+                tools=self.configuration.tools,
+                output_parser=StructuredChatOutputParser(),
+                verbose=True
+            )
+        else:
+            raise NotImplementedError(f"Unknown Agent Strategy: {self.configuration.strategy}")
+
+        return agent
+
+    def should_use_agent(self, query: str) -> bool:
+        return self.agent.should_use_agent(query)
+
+    def run(self, query: str) -> AgentExecuteResult:
+        agent_executor = LCAgentExecutor.from_agent_and_tools(
+            agent=self.agent,
+            tools=self.configuration.tools,
+            memory=self.configuration.memory,
+            max_iterations=self.configuration.max_iterations,
+            max_execution_time=self.configuration.max_execution_time,
+            early_stopping_method=self.configuration.early_stopping_method,
+            callbacks=self.configuration.callbacks
+        )
+
+        try:
+            output = agent_executor.run(query)
+        except Exception:
+            logging.exception("agent_executor run failed")
+            output = None
+
+        return AgentExecuteResult(
+            output=output,
+            strategy=self.configuration.strategy,
+            configuration=self.configuration
+        )

api/core/callback_handler/agent_loop_gather_callback_handler.py

@@ -1,25 +1,29 @@
+import json
 import logging
 import time
 
 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
+from langchain.schema import AgentAction, AgentFinish, LLMResult, ChatGeneration
 
 from core.callback_handler.entity.agent_loop import AgentLoop
 from core.conversation_message_task import ConversationMessageTask
+from core.model_providers.models.llm.base import BaseLLM
 
 
 class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
     """Callback Handler that prints to std out."""
     raise_error: bool = True
 
-    def __init__(self, model_name, conversation_message_task: ConversationMessageTask) -> None:
+    def __init__(self, model_instant: BaseLLM, conversation_message_task: ConversationMessageTask) -> None:
         """Initialize callback handler."""
-        self.model_name = model_name
+        self.model_instant = model_instant
         self.conversation_message_task = conversation_message_task
         self._agent_loops = []
         self._current_loop = None
+        self._message_agent_thought = None
         self.current_chain = None
 
     @property
@@ -29,6 +33,7 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
     def clear_agent_loops(self) -> None:
         self._agent_loops = []
         self._current_loop = None
+        self._message_agent_thought = None
 
     @property
     def always_verbose(self) -> bool:
@@ -61,16 +66,29 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
         # kwargs={}
         if self._current_loop and self._current_loop.status == 'llm_started':
             self._current_loop.status = 'llm_end'
-            self._current_loop.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
-            self._current_loop.completion = response.generations[0][0].text
-            self._current_loop.completion_tokens = response.llm_output['token_usage']['completion_tokens']
+            if response.llm_output:
+                self._current_loop.prompt_tokens = response.llm_output['token_usage']['prompt_tokens']
+            completion_generation = response.generations[0][0]
+            if isinstance(completion_generation, ChatGeneration):
+                completion_message = completion_generation.message
+                if 'function_call' in completion_message.additional_kwargs:
+                    self._current_loop.completion \
+                        = json.dumps({'function_call': completion_message.additional_kwargs['function_call']})
+                else:
+                    self._current_loop.completion = response.generations[0][0].text
+            else:
+                self._current_loop.completion = completion_generation.text
+
+            if response.llm_output:
+                self._current_loop.completion_tokens = response.llm_output['token_usage']['completion_tokens']
 
     def on_llm_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
-        logging.error(error)
+        logging.exception(error)
         self._agent_loops = []
         self._current_loop = None
+        self._message_agent_thought = None
 
     def on_tool_start(
         self,
@@ -89,15 +107,29 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
     ) -> Any:
         """Run on agent action."""
         tool = action.tool
-        tool_input = action.tool_input
-        action_name_position = action.log.index("\nAction:") + 1 if action.log else -1
-        thought = action.log[:action_name_position].strip() if action.log else ''
+        tool_input = json.dumps({"query": action.tool_input}
+                                if isinstance(action.tool_input, str) else action.tool_input)
+        completion = None
+        if isinstance(action, openai_functions_agent.base._FunctionsAgentAction) \
+                or isinstance(action, openai_functions_multi_agent.base._FunctionsAgentAction):
+            thought = action.log.strip()
+            completion = json.dumps({'function_call': action.message_log[0].additional_kwargs['function_call']})
+        else:
+            action_name_position = action.log.index("Action:") if action.log else -1
+            thought = action.log[:action_name_position].strip() if action.log else ''
 
         if self._current_loop and self._current_loop.status == 'llm_end':
             self._current_loop.status = 'agent_action'
             self._current_loop.thought = thought
             self._current_loop.tool_name = tool
             self._current_loop.tool_input = tool_input
+            if completion is not None:
+                self._current_loop.completion = completion
+
+            self._message_agent_thought = self.conversation_message_task.on_agent_start(
+                self.current_chain,
+                self._current_loop
+            )
 
     def on_tool_end(
         self,
@@ -120,18 +152,22 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             self._current_loop.completed_at = time.perf_counter()
             self._current_loop.latency = self._current_loop.completed_at - self._current_loop.started_at
 
-            self.conversation_message_task.on_agent_end(self.current_chain, self.model_name, self._current_loop)
+            self.conversation_message_task.on_agent_end(
+                self._message_agent_thought, self.model_instant, self._current_loop
+            )
 
             self._agent_loops.append(self._current_loop)
             self._current_loop = None
+            self._message_agent_thought = None
 
     def on_tool_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
         """Do nothing."""
-        logging.error(error)
+        logging.exception(error)
         self._agent_loops = []
         self._current_loop = None
+        self._message_agent_thought = None
 
     def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
         """Run on agent end."""
@@ -141,10 +177,18 @@ class AgentLoopGatherCallbackHandler(BaseCallbackHandler):
             self._current_loop.completed = True
             self._current_loop.completed_at = time.perf_counter()
             self._current_loop.latency = self._current_loop.completed_at - self._current_loop.started_at
+            self._current_loop.thought = '[DONE]'
+            self._message_agent_thought = self.conversation_message_task.on_agent_start(
+                self.current_chain,
+                self._current_loop
+            )
 
-            self.conversation_message_task.on_agent_end(self.current_chain, self.model_name, self._current_loop)
+            self.conversation_message_task.on_agent_end(
+                self._message_agent_thought, self.model_instant, self._current_loop
+            )
 
             self._agent_loops.append(self._current_loop)
             self._current_loop = None
+            self._message_agent_thought = None
         elif not self._current_loop and self._agent_loops:
             self._agent_loops[-1].status = 'agent_finish'

api/core/callback_handler/dataset_tool_callback_handler.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from typing import Any, Dict, List, Union, Optional
@@ -43,9 +44,11 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
         input_str: str,
         **kwargs: Any,
     ) -> None:
-        tool_name = serialized.get('name')
-        dataset_id = tool_name[len("dataset-"):]
-        self.conversation_message_task.on_dataset_query_end(DatasetQueryObj(dataset_id=dataset_id, query=input_str))
+        # tool_name = serialized.get('name')
+        input_dict = json.loads(input_str.replace("'", "\""))
+        dataset_id = input_dict.get('dataset_id')
+        query = input_dict.get('query')
+        self.conversation_message_task.on_dataset_query_end(DatasetQueryObj(dataset_id=dataset_id, query=query))
 
     def on_tool_end(
         self,
@@ -65,4 +68,4 @@ class DatasetToolCallbackHandler(BaseCallbackHandler):
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
         """Do nothing."""
-        logging.error(error)
+        logging.exception(error)

api/core/callback_handler/entity/agent_loop.py

@@ -10,9 +10,9 @@ class AgentLoop(BaseModel):
     tool_output: str = None
 
     prompt: str = None
-    prompt_tokens: int = None
+    prompt_tokens: int = 0
     completion: str = None
-    completion_tokens: int = None
+    completion_tokens: int = 0
 
     latency: float = None
 

api/core/callback_handler/llm_callback_handler.py

@@ -1,22 +1,22 @@
 import logging
 import time
-from typing import Any, Dict, List, Union, Optional
+from typing import Any, Dict, List, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult, HumanMessage, AIMessage, SystemMessage, BaseMessage
+from langchain.schema import LLMResult, BaseMessage
 
 from core.callback_handler.entity.llm_message import LLMMessage
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
-from core.llm.streamable_chat_open_ai import StreamableChatOpenAI
-from core.llm.streamable_open_ai import StreamableOpenAI
+from core.model_providers.models.entity.message import to_prompt_messages, PromptMessage
+from core.model_providers.models.llm.base import BaseLLM
 
 
 class LLMCallbackHandler(BaseCallbackHandler):
     raise_error: bool = True
 
-    def __init__(self, llm: Union[StreamableOpenAI, StreamableChatOpenAI],
+    def __init__(self, model_instance: BaseLLM,
                  conversation_message_task: ConversationMessageTask):
-        self.llm = llm
+        self.model_instance = model_instance
         self.llm_message = LLMMessage()
         self.start_at = None
         self.conversation_message_task = conversation_message_task
@@ -48,7 +48,7 @@ class LLMCallbackHandler(BaseCallbackHandler):
             })
 
         self.llm_message.prompt = real_prompts
-        self.llm_message.prompt_tokens = self.llm.get_num_tokens_from_messages(messages[0])
+        self.llm_message.prompt_tokens = self.model_instance.get_num_tokens(to_prompt_messages(messages[0]))
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
@@ -60,7 +60,7 @@ class LLMCallbackHandler(BaseCallbackHandler):
             "text": prompts[0]
         }]
 
-        self.llm_message.prompt_tokens = self.llm.get_num_tokens(prompts[0])
+        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()
@@ -70,7 +70,7 @@ class LLMCallbackHandler(BaseCallbackHandler):
             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.llm.get_num_tokens(self.llm_message.completion)
+        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)
 
@@ -91,7 +91,9 @@ class LLMCallbackHandler(BaseCallbackHandler):
             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.llm.get_num_tokens(self.llm_message.completion)
+                self.llm_message.completion_tokens = self.model_instance.get_num_tokens(
+                    [PromptMessage(content=self.llm_message.completion)]
+                )
                 self.conversation_message_task.save_message(llm_message=self.llm_message, by_stopped=True)
         else:
-            logging.error(error)
+            logging.debug("on_llm_error: %s", error)

api/core/callback_handler/main_chain_gather_callback_handler.py

@@ -5,9 +5,7 @@ from typing import Any, Dict, Union
 
 from langchain.callbacks.base import BaseCallbackHandler
 
-from core.callback_handler.agent_loop_gather_callback_handler import AgentLoopGatherCallbackHandler
 from core.callback_handler.entity.chain_result import ChainResult
-from core.constant import llm_constant
 from core.conversation_message_task import ConversationMessageTask
 
 
@@ -20,15 +18,13 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
         self._current_chain_result = None
         self._current_chain_message = None
         self.conversation_message_task = conversation_message_task
-        self.agent_loop_gather_callback_handler = AgentLoopGatherCallbackHandler(
-            llm_constant.agent_model_name,
-            conversation_message_task
-        )
+        self.agent_callback = None
 
     def clear_chain_results(self) -> None:
         self._current_chain_result = None
         self._current_chain_message = None
-        self.agent_loop_gather_callback_handler.current_chain = None
+        if self.agent_callback:
+            self.agent_callback.current_chain = None
 
     @property
     def always_verbose(self) -> bool:
@@ -58,7 +54,8 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
                     started_at=time.perf_counter()
                 )
                 self._current_chain_message = self.conversation_message_task.init_chain(self._current_chain_result)
-                self.agent_loop_gather_callback_handler.current_chain = self._current_chain_message
+                if self.agent_callback:
+                    self.agent_callback.current_chain = self._current_chain_message
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
         """Print out that we finished a chain."""
@@ -75,5 +72,5 @@ class MainChainGatherCallbackHandler(BaseCallbackHandler):
     def on_chain_error(
         self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
     ) -> None:
-        logging.error(error)
-        self.clear_chain_results()
+        logging.exception(error)
+        self.clear_chain_results()

api/core/chain/chain_builder.py

@@ -1,32 +0,0 @@
-from typing import Optional
-
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.chain.sensitive_word_avoidance_chain import SensitiveWordAvoidanceChain
-from core.chain.tool_chain import ToolChain
-
-
-class ChainBuilder:
-    @classmethod
-    def to_tool_chain(cls, tool, **kwargs) -> ToolChain:
-        return ToolChain(
-            tool=tool,
-            input_key=kwargs.get('input_key', 'input'),
-            output_key=kwargs.get('output_key', 'tool_output'),
-            callbacks=[DifyStdOutCallbackHandler()]
-        )
-
-    @classmethod
-    def to_sensitive_word_avoidance_chain(cls, tool_config: dict, **kwargs) -> Optional[
-        SensitiveWordAvoidanceChain]:
-        sensitive_words = tool_config.get("words", "")
-        if tool_config.get("enabled", False) \
-                and sensitive_words:
-            return SensitiveWordAvoidanceChain(
-                sensitive_words=sensitive_words.split(","),
-                canned_response=tool_config.get("canned_response", ''),
-                output_key="sensitive_word_avoidance_output",
-                callbacks=[DifyStdOutCallbackHandler()],
-                **kwargs
-            )
-
-        return None

api/core/chain/llm_router_chain.py

@@ -1,111 +0,0 @@
-"""Base classes for LLM-powered router chains."""
-from __future__ import annotations
-
-from typing import Any, Dict, List, Optional, Type, cast, NamedTuple
-
-from langchain.base_language import BaseLanguageModel
-from langchain.callbacks.manager import CallbackManagerForChainRun
-from langchain.chains.base import Chain
-from pydantic import root_validator
-
-from langchain.chains import LLMChain
-from langchain.prompts import BasePromptTemplate
-from langchain.schema import BaseOutputParser, OutputParserException
-
-from libs.json_in_md_parser import parse_and_check_json_markdown
-
-
-class Route(NamedTuple):
-    destination: Optional[str]
-    next_inputs: Dict[str, Any]
-
-
-class LLMRouterChain(Chain):
-    """A router chain that uses an LLM chain to perform routing."""
-
-    llm_chain: LLMChain
-    """LLM chain used to perform routing"""
-
-    @root_validator()
-    def validate_prompt(cls, values: dict) -> dict:
-        prompt = values["llm_chain"].prompt
-        if prompt.output_parser is None:
-            raise ValueError(
-                "LLMRouterChain requires base llm_chain prompt to have an output"
-                " parser that converts LLM text output to a dictionary with keys"
-                " 'destination' and 'next_inputs'. Received a prompt with no output"
-                " parser."
-            )
-        return values
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Will be whatever keys the LLM chain prompt expects.
-
-        :meta private:
-        """
-        return self.llm_chain.input_keys
-
-    def _validate_outputs(self, outputs: Dict[str, Any]) -> None:
-        super()._validate_outputs(outputs)
-        if not isinstance(outputs["next_inputs"], dict):
-            raise ValueError
-
-    def _call(
-            self,
-            inputs: Dict[str, Any],
-            run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        output = cast(
-            Dict[str, Any],
-            self.llm_chain.predict_and_parse(**inputs),
-        )
-        return output
-
-    @classmethod
-    def from_llm(
-        cls, llm: BaseLanguageModel, prompt: BasePromptTemplate, **kwargs: Any
-    ) -> LLMRouterChain:
-        """Convenience constructor."""
-        llm_chain = LLMChain(llm=llm, prompt=prompt)
-        return cls(llm_chain=llm_chain, **kwargs)
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["destination", "next_inputs"]
-
-    def route(self, inputs: Dict[str, Any]) -> Route:
-        result = self(inputs)
-        return Route(result["destination"], result["next_inputs"])
-
-
-class RouterOutputParser(BaseOutputParser[Dict[str, str]]):
-    """Parser for output of router chain int he multi-prompt chain."""
-
-    default_destination: str = "DEFAULT"
-    next_inputs_type: Type = str
-    next_inputs_inner_key: str = "input"
-
-    def parse(self, text: str) -> Dict[str, Any]:
-        try:
-            expected_keys = ["destination", "next_inputs"]
-            parsed = parse_and_check_json_markdown(text, expected_keys)
-            if not isinstance(parsed["destination"], str):
-                raise ValueError("Expected 'destination' to be a string.")
-            if not isinstance(parsed["next_inputs"], self.next_inputs_type):
-                raise ValueError(
-                    f"Expected 'next_inputs' to be {self.next_inputs_type}."
-                )
-            parsed["next_inputs"] = {self.next_inputs_inner_key: parsed["next_inputs"]}
-            if (
-                parsed["destination"].strip().lower()
-                == self.default_destination.lower()
-            ):
-                parsed["destination"] = None
-            else:
-                parsed["destination"] = parsed["destination"].strip()
-            return parsed
-        except Exception as e:
-            raise OutputParserException(
-                f"Parsing text\n{text}\n of llm router raised following error:\n{e}"
-            )

api/core/chain/main_chain_builder.py

@@ -1,110 +0,0 @@
-from typing import Optional, List, cast
-
-from langchain.chains import SequentialChain
-from langchain.chains.base import Chain
-from langchain.memory.chat_memory import BaseChatMemory
-
-from core.callback_handler.main_chain_gather_callback_handler import MainChainGatherCallbackHandler
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.chain.chain_builder import ChainBuilder
-from core.chain.multi_dataset_router_chain import MultiDatasetRouterChain
-from core.conversation_message_task import ConversationMessageTask
-from extensions.ext_database import db
-from models.dataset import Dataset
-
-
-class MainChainBuilder:
-    @classmethod
-    def to_langchain_components(cls, tenant_id: str, agent_mode: dict, memory: Optional[BaseChatMemory],
-                                rest_tokens: int,
-                                conversation_message_task: ConversationMessageTask):
-        first_input_key = "input"
-        final_output_key = "output"
-
-        chains = []
-
-        chain_callback_handler = MainChainGatherCallbackHandler(conversation_message_task)
-
-        # agent mode
-        tool_chains, chains_output_key = cls.get_agent_chains(
-            tenant_id=tenant_id,
-            agent_mode=agent_mode,
-            rest_tokens=rest_tokens,
-            memory=memory,
-            conversation_message_task=conversation_message_task
-        )
-        chains += tool_chains
-
-        if chains_output_key:
-            final_output_key = chains_output_key
-
-        if len(chains) == 0:
-            return None
-
-        for chain in chains:
-            chain = cast(Chain, chain)
-            chain.callbacks.append(chain_callback_handler)
-
-        # build main chain
-        overall_chain = SequentialChain(
-            chains=chains,
-            input_variables=[first_input_key],
-            output_variables=[final_output_key],
-            memory=memory,  # only for use the memory prompt input key
-        )
-
-        return overall_chain
-
-    @classmethod
-    def get_agent_chains(cls, tenant_id: str, agent_mode: dict,
-                         rest_tokens: int,
-                         memory: Optional[BaseChatMemory],
-                         conversation_message_task: ConversationMessageTask):
-        # agent mode
-        chains = []
-        if agent_mode and agent_mode.get('enabled'):
-            tools = agent_mode.get('tools', [])
-
-            pre_fixed_chains = []
-            # agent_tools = []
-            datasets = []
-            for tool in tools:
-                tool_type = list(tool.keys())[0]
-                tool_config = list(tool.values())[0]
-                if tool_type == 'sensitive-word-avoidance':
-                    chain = ChainBuilder.to_sensitive_word_avoidance_chain(tool_config)
-                    if chain:
-                        pre_fixed_chains.append(chain)
-                elif tool_type == "dataset":
-                    # get dataset from dataset id
-                    dataset = db.session.query(Dataset).filter(
-                        Dataset.tenant_id == tenant_id,
-                        Dataset.id == tool_config.get("id")
-                    ).first()
-
-                    if dataset:
-                        datasets.append(dataset)
-
-            # add pre-fixed chains
-            chains += pre_fixed_chains
-
-            if len(datasets) > 0:
-                # tool to chain
-                multi_dataset_router_chain = MultiDatasetRouterChain.from_datasets(
-                    tenant_id=tenant_id,
-                    datasets=datasets,
-                    conversation_message_task=conversation_message_task,
-                    rest_tokens=rest_tokens,
-                    callbacks=[DifyStdOutCallbackHandler()]
-                )
-                chains.append(multi_dataset_router_chain)
-
-        final_output_key = cls.get_chains_output_key(chains)
-
-        return chains, final_output_key
-
-    @classmethod
-    def get_chains_output_key(cls, chains: List[Chain]):
-        if len(chains) > 0:
-            return chains[-1].output_keys[0]
-        return None

api/core/chain/multi_dataset_router_chain.py

@@ -1,198 +0,0 @@
-import math
-import re
-from typing import Mapping, List, Dict, Any, Optional
-
-from langchain import PromptTemplate
-from langchain.callbacks.manager import CallbackManagerForChainRun
-from langchain.chains.base import Chain
-from pydantic import Extra
-
-from core.callback_handler.dataset_tool_callback_handler import DatasetToolCallbackHandler
-from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
-from core.chain.llm_router_chain import LLMRouterChain, RouterOutputParser
-from core.conversation_message_task import ConversationMessageTask
-from core.llm.llm_builder import LLMBuilder
-from core.tool.dataset_index_tool import DatasetTool
-from models.dataset import Dataset, DatasetProcessRule
-
-DEFAULT_K = 2
-CONTEXT_TOKENS_PERCENT = 0.3
-MULTI_PROMPT_ROUTER_TEMPLATE = """
-Given a raw text input to a language model select the model prompt best suited for \
-the input. You will be given the names of the available prompts and a description of \
-what the prompt is best suited for. You may also revise the original input if you \
-think that revising it will ultimately lead to a better response from the language \
-model.
-
-<< FORMATTING >>
-Return a markdown code snippet with a JSON object formatted to look like, \
-no any other string out of markdown code snippet:
-```json
-{{{{
-    "destination": string \\ name of the prompt to use or "DEFAULT"
-    "next_inputs": string \\ a potentially modified version of the original input
-}}}}
-```
-
-REMEMBER: "destination" MUST be one of the candidate prompt names specified below OR \
-it can be "DEFAULT" if the input is not well suited for any of the candidate prompts.
-REMEMBER: "next_inputs" can just be the original input if you don't think any \
-modifications are needed.
-
-<< CANDIDATE PROMPTS >>
-{destinations}
-
-<< INPUT >>
-{{input}}
-
-<< OUTPUT >>
-"""
-
-
-class MultiDatasetRouterChain(Chain):
-    """Use a single chain to route an input to one of multiple candidate chains."""
-
-    router_chain: LLMRouterChain
-    """Chain for deciding a destination chain and the input to it."""
-    dataset_tools: Mapping[str, DatasetTool]
-    """Map of name to candidate chains that inputs can be routed to."""
-
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Will be whatever keys the router chain prompt expects.
-
-        :meta private:
-        """
-        return self.router_chain.input_keys
-
-    @property
-    def output_keys(self) -> List[str]:
-        return ["text"]
-
-    @classmethod
-    def from_datasets(
-            cls,
-            tenant_id: str,
-            datasets: List[Dataset],
-            conversation_message_task: ConversationMessageTask,
-            rest_tokens: int,
-            **kwargs: Any,
-    ):
-        """Convenience constructor for instantiating from destination prompts."""
-        llm = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name='gpt-3.5-turbo',
-            temperature=0,
-            max_tokens=1024,
-            callbacks=[DifyStdOutCallbackHandler()]
-        )
-
-        destinations = ["[[{}]]: {}".format(d.id, d.description.replace('\n', ' ') if d.description
-                        else ('useful for when you want to answer queries about the ' + d.name))
-                        for d in datasets]
-        destinations_str = "\n".join(destinations)
-        router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
-            destinations=destinations_str
-        )
-
-        router_prompt = PromptTemplate(
-            template=router_template,
-            input_variables=["input"],
-            output_parser=RouterOutputParser(),
-        )
-
-        router_chain = LLMRouterChain.from_llm(llm, router_prompt)
-        dataset_tools = {}
-        for dataset in datasets:
-            # fulfill description when it is empty
-            if dataset.available_document_count == 0 or dataset.available_document_count == 0:
-                continue
-
-            description = dataset.description
-            if not description:
-                description = 'useful for when you want to answer queries about the ' + dataset.name
-
-            k = cls._dynamic_calc_retrieve_k(dataset, rest_tokens)
-            if k == 0:
-                continue
-
-            dataset_tool = DatasetTool(
-                name=f"dataset-{dataset.id}",
-                description=description,
-                k=k,
-                dataset=dataset,
-                callbacks=[DatasetToolCallbackHandler(conversation_message_task), DifyStdOutCallbackHandler()]
-            )
-
-            dataset_tools[str(dataset.id)] = dataset_tool
-
-        return cls(
-            router_chain=router_chain,
-            dataset_tools=dataset_tools,
-            **kwargs,
-        )
-
-    @classmethod
-    def _dynamic_calc_retrieve_k(cls, dataset: Dataset, rest_tokens: int) -> int:
-        processing_rule = dataset.latest_process_rule
-        if not processing_rule:
-            return DEFAULT_K
-
-        if processing_rule.mode == "custom":
-            rules = processing_rule.rules_dict
-            if not rules:
-                return DEFAULT_K
-
-            segmentation = rules["segmentation"]
-            segment_max_tokens = segmentation["max_tokens"]
-        else:
-            segment_max_tokens = DatasetProcessRule.AUTOMATIC_RULES['segmentation']['max_tokens']
-
-        # when rest_tokens is less than default context tokens
-        if rest_tokens < segment_max_tokens * DEFAULT_K:
-            return rest_tokens // segment_max_tokens
-
-        context_limit_tokens = math.floor(rest_tokens * CONTEXT_TOKENS_PERCENT)
-
-        # when context_limit_tokens is less than default context tokens, use default_k
-        if context_limit_tokens <= segment_max_tokens * DEFAULT_K:
-            return DEFAULT_K
-
-        # Expand the k value when there's still some room left in the 30% rest tokens space
-        return context_limit_tokens // segment_max_tokens
-
-    def _call(
-        self,
-        inputs: Dict[str, Any],
-        run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        if len(self.dataset_tools) == 0:
-            return {"text": ''}
-        elif len(self.dataset_tools) == 1:
-            return {"text": next(iter(self.dataset_tools.values())).run(inputs['input'])}
-
-        route = self.router_chain.route(inputs)
-
-        destination = ''
-        if route.destination:
-            pattern = r'\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b'
-            match = re.search(pattern, route.destination, re.IGNORECASE)
-            if match:
-                destination = match.group()
-
-        if not destination:
-            return {"text": ''}
-        elif destination in self.dataset_tools:
-            return {"text": self.dataset_tools[destination].run(
-                route.next_inputs['input']
-            )}
-        else:
-            raise ValueError(
-                f"Received invalid destination chain name '{destination}'"
-            )

api/core/chain/tool_chain.py

@@ -1,51 +0,0 @@
-from typing import List, Dict, Optional, Any
-
-from langchain.callbacks.manager import CallbackManagerForChainRun, AsyncCallbackManagerForChainRun
-from langchain.chains.base import Chain
-from langchain.tools import BaseTool
-
-
-class ToolChain(Chain):
-    input_key: str = "input"  #: :meta private:
-    output_key: str = "output"  #: :meta private:
-
-    tool: BaseTool
-
-    @property
-    def _chain_type(self) -> str:
-        return "tool_chain"
-
-    @property
-    def input_keys(self) -> List[str]:
-        """Expect input key.
-
-        :meta private:
-        """
-        return [self.input_key]
-
-    @property
-    def output_keys(self) -> List[str]:
-        """Return output key.
-
-        :meta private:
-        """
-        return [self.output_key]
-
-    def _call(
-            self,
-            inputs: Dict[str, Any],
-            run_manager: Optional[CallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        input = inputs[self.input_key]
-        output = self.tool.run(input, self.verbose)
-        return {self.output_key: output}
-
-    async def _acall(
-            self,
-            inputs: Dict[str, Any],
-            run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
-    ) -> Dict[str, Any]:
-        """Run the logic of this chain and return the output."""
-        input = inputs[self.input_key]
-        output = await self.tool.arun(input, self.verbose)
-        return {self.output_key: output}

api/core/completion.py

@@ -1,37 +1,31 @@
 import logging
+import re
 from typing import Optional, List, Union, Tuple
 
-from langchain.base_language import BaseLanguageModel
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.chat_models.base import BaseChatModel
-from langchain.llms import BaseLLM
-from langchain.schema import BaseMessage, HumanMessage
+from langchain.schema import BaseMessage
 from requests.exceptions import ChunkedEncodingError
 
-from core.constant import llm_constant
+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.callback_handler.std_out_callback_handler import DifyStreamingStdOutCallbackHandler, \
-    DifyStdOutCallbackHandler
 from core.conversation_message_task import ConversationMessageTask, ConversationTaskStoppedException
-from core.llm.error import LLMBadRequestError
-from core.llm.llm_builder import LLMBuilder
-from core.chain.main_chain_builder import MainChainBuilder
-from core.llm.streamable_chat_open_ai import StreamableChatOpenAI
-from core.llm.streamable_open_ai import StreamableOpenAI
+from core.model_providers.error import LLMBadRequestError
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
     ReadOnlyConversationTokenDBBufferSharedMemory
-from core.memory.read_only_conversation_token_db_string_buffer_shared_memory import \
-    ReadOnlyConversationTokenDBStringBufferSharedMemory
+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.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.model import App, AppModelConfig, Account, Conversation, Message
+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: Account, conversation: Optional[Conversation], streaming: bool, is_override: bool = False):
+                 user: Union[Account, EndUser], conversation: Optional[Conversation], streaming: bool, is_override: bool = False):
         """
         errors: ProviderTokenNotInitError
         """
@@ -49,12 +43,10 @@ class Completion:
 
             inputs = conversation.inputs
 
-        rest_tokens_for_context_and_memory = cls.get_validate_rest_tokens(
-            mode=app.mode,
+        final_model_instance = ModelFactory.get_text_generation_model_from_model_config(
             tenant_id=app.tenant_id,
-            app_model_config=app_model_config,
-            query=query,
-            inputs=inputs
+            model_config=app_model_config.model_dict,
+            streaming=streaming
         )
 
         conversation_message_task = ConversationMessageTask(
@@ -66,34 +58,56 @@ class Completion:
             is_override=is_override,
             inputs=inputs,
             query=query,
-            streaming=streaming
+            streaming=streaming,
+            model_instance=final_model_instance
         )
 
-        # build main chain include agent
-        main_chain = MainChainBuilder.to_langchain_components(
+        rest_tokens_for_context_and_memory = cls.get_validate_rest_tokens(
+            mode=app.mode,
+            model_instance=final_model_instance,
+            app_model_config=app_model_config,
+            query=query,
+            inputs=inputs
+        )
+
+        # init orchestrator rule parser
+        orchestrator_rule_parser = OrchestratorRuleParser(
             tenant_id=app.tenant_id,
-            agent_mode=app_model_config.agent_mode_dict,
-            rest_tokens=rest_tokens_for_context_and_memory,
-            memory=ReadOnlyConversationTokenDBStringBufferSharedMemory(memory=memory) if memory else None,
-            conversation_message_task=conversation_message_task
+            app_model_config=app_model_config
         )
 
-        chain_output = ''
-        if main_chain:
-            chain_output = main_chain.run(query)
+        # 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:
             cls.run_final_llm(
-                tenant_id=app.tenant_id,
+                model_instance=final_model_instance,
                 mode=app.mode,
                 app_model_config=app_model_config,
                 query=query,
                 inputs=inputs,
-                chain_output=chain_output,
+                agent_execute_result=agent_execute_result,
                 conversation_message_task=conversation_message_task,
-                memory=memory,
-                streaming=streaming
+                memory=memory
             )
         except ConversationTaskStoppedException:
             return
@@ -104,54 +118,48 @@ class Completion:
             return
 
     @classmethod
-    def run_final_llm(cls, tenant_id: str, mode: str, app_model_config: AppModelConfig, query: str, inputs: dict,
-                      chain_output: str,
+    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], streaming: bool):
-        final_llm = LLMBuilder.to_llm_from_model(
-            tenant_id=tenant_id,
-            model=app_model_config.model_dict,
-            streaming=streaming
-        )
+                      memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]):
+        # 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 != PlanningStrategy.ROUTER:
+            fake_response = agent_execute_result.output
 
         # get llm prompt
-        prompt, stop_words = cls.get_main_llm_prompt(
+        prompt_messages, stop_words = cls.get_main_llm_prompt(
             mode=mode,
-            llm=final_llm,
             model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
             query=query,
             inputs=inputs,
-            chain_output=chain_output,
+            agent_execute_result=agent_execute_result,
             memory=memory
         )
 
-        final_llm.callbacks = cls.get_llm_callbacks(final_llm, streaming, conversation_message_task)
-
         cls.recale_llm_max_tokens(
-            final_llm=final_llm,
-            model=app_model_config.model_dict,
-            prompt=prompt,
-            mode=mode
+            model_instance=model_instance,
+            prompt_messages=prompt_messages,
         )
 
-        response = final_llm.generate([prompt], stop_words)
+        response = model_instance.run(
+            messages=prompt_messages,
+            stop=stop_words,
+            callbacks=[LLMCallbackHandler(model_instance, conversation_message_task)],
+            fake_response=fake_response
+        )
 
         return response
 
     @classmethod
-    def get_main_llm_prompt(cls, mode: str, llm: BaseLanguageModel, model: dict,
+    def get_main_llm_prompt(cls, mode: str, model: dict,
                             pre_prompt: str, query: str, inputs: dict,
-                            chain_output: Optional[str],
+                            agent_execute_result: Optional[AgentExecuteResult],
                             memory: Optional[ReadOnlyConversationTokenDBBufferSharedMemory]) -> \
-            Tuple[Union[str | List[BaseMessage]], Optional[List[str]]]:
-        # disable template string in query
-        # query_params = JinjaPromptTemplate.from_template(template=query).input_variables
-        # if query_params:
-        #     for query_param in query_params:
-        #         if query_param not in inputs:
-        #             inputs[query_param] = '{{' + query_param + '}}'
-
+            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.
@@ -165,18 +173,13 @@ When answer to user:
 - 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 chain_output else "")
+""" if agent_execute_result else "")
                          + (pre_prompt + "\n" if pre_prompt else "")
                          + "{{query}}\n"
             )
 
-            if chain_output:
-                inputs['context'] = chain_output
-                # context_params = JinjaPromptTemplate.from_template(template=chain_output).input_variables
-                # if context_params:
-                #     for context_param in context_params:
-                #         if context_param not in inputs:
-                #             inputs[context_param] = '{{' + context_param + '}}'
+            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(
@@ -184,11 +187,7 @@ And answer according to the language of the user's question.
                 **prompt_inputs
             )
 
-            if isinstance(llm, BaseChatModel):
-                # use chat llm as completion model
-                return [HumanMessage(content=prompt_content)], None
-            else:
-                return prompt_content, None
+            return [PromptMessage(content=prompt_content)], None
         else:
             messages: List[BaseMessage] = []
 
@@ -206,8 +205,8 @@ And answer according to the language of the user's question.
                 if pre_prompt_inputs:
                     human_inputs.update(pre_prompt_inputs)
 
-            if chain_output:
-                human_inputs['context'] = chain_output
+            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>
@@ -233,25 +232,19 @@ And answer according to the language of the user's question.
                     inputs=human_inputs
                 )
 
-                curr_message_tokens = memory.llm.get_num_tokens_from_messages([tmp_human_message])
-                model_name = model['name']
-                max_tokens = model.get("completion_params").get('max_tokens')
-                rest_tokens = llm_constant.max_context_token_length[model_name] \
-                              - max_tokens - curr_message_tokens
-                rest_tokens = max(rest_tokens, 0)
+                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)
-
-                # disable template string in query
-                # histories_params = JinjaPromptTemplate.from_template(template=histories).input_variables
-                # if histories_params:
-                #     for histories_param in histories_params:
-                #         if histories_param not in human_inputs:
-                #             human_inputs[histories_param] = '{{' + histories_param + '}}'
-
                 human_message_prompt += "\n\n" 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>"
-                human_message_prompt += histories + "</histories>"
+                                        "inside <histories></histories> XML tags.\n\n<histories>\n"
+                human_message_prompt += histories + "\n</histories>"
 
             human_message_prompt += query_prompt
 
@@ -263,22 +256,14 @@ And answer according to the language of the user's question.
 
             messages.append(human_message)
 
-            return messages, ['\nHuman:']
+            for message in messages:
+                message.content = re.sub(r'<\|.*?\|>', '', message.content)
 
-    @classmethod
-    def get_llm_callbacks(cls, llm: Union[StreamableOpenAI, StreamableChatOpenAI],
-                          streaming: bool,
-                          conversation_message_task: ConversationMessageTask) -> List[BaseCallbackHandler]:
-        llm_callback_handler = LLMCallbackHandler(llm, conversation_message_task)
-        if streaming:
-            return [llm_callback_handler, DifyStreamingStdOutCallbackHandler()]
-        else:
-            return [llm_callback_handler, DifyStdOutCallbackHandler()]
+            return to_prompt_messages(messages), ['\nHuman:', '</histories>']
 
     @classmethod
     def get_history_messages_from_memory(cls, memory: ReadOnlyConversationTokenDBBufferSharedMemory,
-                                         max_token_limit: int) -> \
-            str:
+                                         max_token_limit: int) -> str:
         """Get memory messages."""
         memory.max_token_limit = max_token_limit
         memory_key = memory.memory_variables[0]
@@ -290,15 +275,15 @@ And answer according to the language of the user's question.
                                      conversation: Conversation,
                                      **kwargs) -> ReadOnlyConversationTokenDBBufferSharedMemory:
         # only for calc token in memory
-        memory_llm = LLMBuilder.to_llm_from_model(
+        memory_model_instance = ModelFactory.get_text_generation_model_from_model_config(
             tenant_id=tenant_id,
-            model=app_model_config.model_dict
+            model_config=app_model_config.model_dict
         )
 
         # use llm config from conversation
         memory = ReadOnlyConversationTokenDBBufferSharedMemory(
             conversation=conversation,
-            llm=memory_llm,
+            model_instance=memory_model_instance,
             max_token_limit=kwargs.get("max_token_limit", 2048),
             memory_key=kwargs.get("memory_key", "chat_history"),
             return_messages=kwargs.get("return_messages", True),
@@ -310,32 +295,29 @@ And answer according to the language of the user's question.
         return memory
 
     @classmethod
-    def get_validate_rest_tokens(cls, mode: str, tenant_id: str, app_model_config: AppModelConfig,
+    def get_validate_rest_tokens(cls, mode: str, model_instance: BaseLLM, app_model_config: AppModelConfig,
                                  query: str, inputs: dict) -> int:
-        llm = LLMBuilder.to_llm_from_model(
-            tenant_id=tenant_id,
-            model=app_model_config.model_dict
-        )
+        model_limited_tokens = model_instance.model_rules.max_tokens.max
+        max_tokens = model_instance.get_model_kwargs().max_tokens
 
-        model_name = app_model_config.model_dict.get("name")
-        model_limited_tokens = llm_constant.max_context_token_length[model_name]
-        max_tokens = app_model_config.model_dict.get("completion_params").get('max_tokens')
+        if model_limited_tokens is None:
+            return -1
+
+        if max_tokens is None:
+            max_tokens = 0
 
         # get prompt without memory and context
-        prompt, _ = cls.get_main_llm_prompt(
+        prompt_messages, _ = cls.get_main_llm_prompt(
             mode=mode,
-            llm=llm,
             model=app_model_config.model_dict,
             pre_prompt=app_model_config.pre_prompt,
             query=query,
             inputs=inputs,
-            chain_output=None,
+            agent_execute_result=None,
             memory=None
         )
 
-        prompt_tokens = llm.get_num_tokens(prompt) if isinstance(prompt, str) \
-            else llm.get_num_tokens_from_messages(prompt)
-
+        prompt_tokens = model_instance.get_num_tokens(prompt_messages)
         rest_tokens = model_limited_tokens - max_tokens - prompt_tokens
         if rest_tokens < 0:
             raise LLMBadRequestError("Query or prefix prompt is too long, you can reduce the prefix prompt, "
@@ -344,51 +326,54 @@ And answer according to the language of the user's question.
         return rest_tokens
 
     @classmethod
-    def recale_llm_max_tokens(cls, final_llm: BaseLanguageModel, model: dict,
-                              prompt: Union[str, List[BaseMessage]], mode: str):
+    def recale_llm_max_tokens(cls, model_instance: BaseLLM, prompt_messages: List[PromptMessage]):
         # recalc max_tokens if sum(prompt_token +  max_tokens) over model token limit
-        model_name = model.get("name")
-        model_limited_tokens = llm_constant.max_context_token_length[model_name]
-        max_tokens = model.get("completion_params").get('max_tokens')
+        model_limited_tokens = model_instance.model_rules.max_tokens.max
+        max_tokens = model_instance.get_model_kwargs().max_tokens
 
-        if mode == 'completion' and isinstance(final_llm, BaseLLM):
-            prompt_tokens = final_llm.get_num_tokens(prompt)
-        else:
-            prompt_tokens = final_llm.get_num_tokens_from_messages(prompt)
+        if model_limited_tokens is None:
+            return
+
+        if max_tokens is None:
+            max_tokens = 0
+
+        prompt_tokens = model_instance.get_num_tokens(prompt_messages)
 
         if prompt_tokens + max_tokens > model_limited_tokens:
             max_tokens = max(model_limited_tokens - prompt_tokens, 16)
-            final_llm.max_tokens = max_tokens
+
+            # update model instance max tokens
+            model_kwargs = model_instance.get_model_kwargs()
+            model_kwargs.max_tokens = max_tokens
+            model_instance.set_model_kwargs(model_kwargs)
 
     @classmethod
     def generate_more_like_this(cls, task_id: str, app: App, message: Message, pre_prompt: str,
                                 app_model_config: AppModelConfig, user: Account, streaming: bool):
 
-        llm = LLMBuilder.to_llm_from_model(
+        final_model_instance = ModelFactory.get_text_generation_model_from_model_config(
             tenant_id=app.tenant_id,
-            model=app_model_config.model_dict,
+            model_config=app_model_config.model_dict,
             streaming=streaming
         )
 
         # get llm prompt
-        original_prompt, _ = cls.get_main_llm_prompt(
+        old_prompt_messages, _ = cls.get_main_llm_prompt(
             mode="completion",
-            llm=llm,
             model=app_model_config.model_dict,
             pre_prompt=pre_prompt,
             query=message.query,
             inputs=message.inputs,
-            chain_output=None,
+            agent_execute_result=None,
             memory=None
         )
 
         original_completion = message.answer.strip()
 
         prompt = MORE_LIKE_THIS_GENERATE_PROMPT
-        prompt = prompt.format(prompt=original_prompt, original_completion=original_completion)
+        prompt = prompt.format(prompt=old_prompt_messages[0].content, original_completion=original_completion)
 
-        if isinstance(llm, BaseChatModel):
-            prompt = [HumanMessage(content=prompt)]
+        prompt_messages = [PromptMessage(content=prompt)]
 
         conversation_message_task = ConversationMessageTask(
             task_id=task_id,
@@ -398,16 +383,16 @@ And answer according to the language of the user's question.
             inputs=message.inputs,
             query=message.query,
             is_override=True if message.override_model_configs else False,
-            streaming=streaming
+            streaming=streaming,
+            model_instance=final_model_instance
         )
 
-        llm.callbacks = cls.get_llm_callbacks(llm, streaming, conversation_message_task)
-
         cls.recale_llm_max_tokens(
-            final_llm=llm,
-            model=app_model_config.model_dict,
-            prompt=prompt,
-            mode='completion'
+            model_instance=final_model_instance,
+            prompt_messages=prompt_messages
         )
 
-        llm.generate([prompt])
+        final_model_instance.run(
+            messages=prompt_messages,
+            callbacks=[LLMCallbackHandler(final_model_instance, conversation_message_task)]
+        )

api/core/constant/llm_constant.py

@@ -1,109 +0,0 @@
-from _decimal import Decimal
-
-models = {
-    'claude-instant-1': 'anthropic',  # 100,000 tokens
-    'claude-2': 'anthropic',  # 100,000 tokens
-    'gpt-4': 'openai',  # 8,192 tokens
-    'gpt-4-32k': 'openai',  # 32,768 tokens
-    'gpt-3.5-turbo': 'openai',  # 4,096 tokens
-    'gpt-3.5-turbo-16k': 'openai',  # 16384 tokens
-    'text-davinci-003': 'openai',  # 4,097 tokens
-    'text-davinci-002': 'openai',  # 4,097 tokens
-    'text-curie-001': 'openai',  # 2,049 tokens
-    'text-babbage-001': 'openai',  # 2,049 tokens
-    'text-ada-001': 'openai',  # 2,049 tokens
-    'text-embedding-ada-002': 'openai',  # 8191 tokens, 1536 dimensions
-    'whisper-1': 'openai'
-}
-
-max_context_token_length = {
-    'claude-instant-1': 100000,
-    'claude-2': 100000,
-    'gpt-4': 8192,
-    'gpt-4-32k': 32768,
-    'gpt-3.5-turbo': 4096,
-    'gpt-3.5-turbo-16k': 16384,
-    'text-davinci-003': 4097,
-    'text-davinci-002': 4097,
-    'text-curie-001': 2049,
-    'text-babbage-001': 2049,
-    'text-ada-001': 2049,
-    'text-embedding-ada-002': 8191,
-}
-
-models_by_mode = {
-    'chat': [
-        'claude-instant-1',  # 100,000 tokens
-        'claude-2',  # 100,000 tokens
-        'gpt-4',  # 8,192 tokens
-        'gpt-4-32k',  # 32,768 tokens
-        'gpt-3.5-turbo',  # 4,096 tokens
-        'gpt-3.5-turbo-16k',  # 16,384 tokens
-    ],
-    'completion': [
-        'claude-instant-1',  # 100,000 tokens
-        'claude-2',  # 100,000 tokens
-        'gpt-4',  # 8,192 tokens
-        'gpt-4-32k',  # 32,768 tokens
-        'gpt-3.5-turbo',  # 4,096 tokens
-        'gpt-3.5-turbo-16k',  # 16,384 tokens
-        'text-davinci-003',  # 4,097 tokens
-        'text-davinci-002'  # 4,097 tokens
-        'text-curie-001',  # 2,049 tokens
-        'text-babbage-001',  # 2,049 tokens
-        'text-ada-001'  # 2,049 tokens
-    ],
-    'embedding': [
-        'text-embedding-ada-002'  # 8191 tokens, 1536 dimensions
-    ]
-}
-
-model_currency = 'USD'
-
-model_prices = {
-    'claude-instant-1': {
-        'prompt': Decimal('0.00163'),
-        'completion': Decimal('0.00551'),
-    },
-    'claude-2': {
-        'prompt': Decimal('0.01102'),
-        'completion': Decimal('0.03268'),
-    },
-    'gpt-4': {
-        'prompt': Decimal('0.03'),
-        'completion': Decimal('0.06'),
-    },
-    'gpt-4-32k': {
-        'prompt': Decimal('0.06'),
-        'completion': Decimal('0.12')
-    },
-    'gpt-3.5-turbo': {
-        'prompt': Decimal('0.0015'),
-        'completion': Decimal('0.002')
-    },
-    'gpt-3.5-turbo-16k': {
-        'prompt': Decimal('0.003'),
-        'completion': Decimal('0.004')
-    },
-    'text-davinci-003': {
-        'prompt': Decimal('0.02'),
-        'completion': Decimal('0.02')
-    },
-    'text-curie-001': {
-        'prompt': Decimal('0.002'),
-        'completion': Decimal('0.002')
-    },
-    'text-babbage-001': {
-        'prompt': Decimal('0.0005'),
-        'completion': Decimal('0.0005')
-    },
-    'text-ada-001': {
-        'prompt': Decimal('0.0004'),
-        'completion': Decimal('0.0004')
-    },
-    'text-embedding-ada-002': {
-        'usage': Decimal('0.0001'),
-    }
-}
-
-agent_model_name = 'text-davinci-003'

api/core/conversation_message_task.py

@@ -6,9 +6,9 @@ from core.callback_handler.entity.agent_loop import AgentLoop
 from core.callback_handler.entity.dataset_query import DatasetQueryObj
 from core.callback_handler.entity.llm_message import LLMMessage
 from core.callback_handler.entity.chain_result import ChainResult
-from core.constant import llm_constant
-from core.llm.llm_builder import LLMBuilder
-from core.llm.provider.llm_provider_service import LLMProviderService
+from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.entity.message import to_prompt_messages, MessageType
+from core.model_providers.models.llm.base import BaseLLM
 from core.prompt.prompt_builder import PromptBuilder
 from core.prompt.prompt_template import JinjaPromptTemplate
 from events.message_event import message_was_created
@@ -16,12 +16,11 @@ 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.provider import ProviderType, Provider
 
 
 class ConversationMessageTask:
     def __init__(self, task_id: str, app: App, app_model_config: AppModelConfig, user: Account,
-                 inputs: dict, query: str, streaming: bool,
+                 inputs: dict, query: str, streaming: bool, model_instance: BaseLLM,
                  conversation: Optional[Conversation] = None, is_override: bool = False):
         self.task_id = task_id
 
@@ -38,9 +37,12 @@ class ConversationMessageTask:
         self.conversation = conversation
         self.is_new_conversation = False
 
+        self.model_instance = model_instance
+
         self.message = 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')
         self.mode = app.mode
 
@@ -52,13 +54,10 @@ class ConversationMessageTask:
             message=self.message,
             conversation=self.conversation,
             chain_pub=False,  # disabled currently
-            agent_thought_pub=False  # disabled currently
+            agent_thought_pub=True
         )
 
     def init(self):
-        provider_name = LLMBuilder.get_default_provider(self.app.tenant_id, self.model_name)
-        self.model_dict['provider'] = provider_name
-
         override_model_configs = None
         if self.is_override:
             override_model_configs = {
@@ -69,6 +68,7 @@ class ConversationMessageTask:
                 "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,
             }
 
@@ -88,15 +88,19 @@ class ConversationMessageTask:
             if self.app_model_config.pre_prompt:
                 system_message = PromptBuilder.to_system_message(self.app_model_config.pre_prompt, self.inputs)
                 system_instruction = system_message.content
-                llm = LLMBuilder.to_llm(self.tenant_id, self.model_name)
-                system_instruction_tokens = llm.get_num_tokens_from_messages([system_message])
+                model_instance = ModelFactory.get_text_generation_model(
+                    tenant_id=self.tenant_id,
+                    model_provider_name=self.provider_name,
+                    model_name=self.model_name
+                )
+                system_instruction_tokens = model_instance.get_num_tokens(to_prompt_messages([system_message]))
 
         if not self.conversation:
             self.is_new_conversation = True
             self.conversation = Conversation(
                 app_id=self.app_model_config.app_id,
                 app_model_config_id=self.app_model_config.id,
-                model_provider=self.model_dict.get('provider'),
+                model_provider=self.provider_name,
                 model_id=self.model_name,
                 override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
                 mode=self.mode,
@@ -116,7 +120,7 @@ class ConversationMessageTask:
 
         self.message = Message(
             app_id=self.app_model_config.app_id,
-            model_provider=self.model_dict.get('provider'),
+            model_provider=self.provider_name,
             model_id=self.model_name,
             override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
             conversation_id=self.conversation.id,
@@ -130,7 +134,7 @@ class ConversationMessageTask:
             answer_unit_price=0,
             provider_response_latency=0,
             total_price=0,
-            currency=llm_constant.model_currency,
+            currency=self.model_instance.get_currency(),
             from_source=('console' if isinstance(self.user, Account) else 'api'),
             from_end_user_id=(self.user.id if isinstance(self.user, EndUser) else None),
             from_account_id=(self.user.id if isinstance(self.user, Account) else None),
@@ -144,12 +148,10 @@ class ConversationMessageTask:
         self._pub_handler.pub_text(text)
 
     def save_message(self, llm_message: LLMMessage, by_stopped: bool = False):
-        model_name = self.app_model_config.model_dict.get('name')
-
         message_tokens = llm_message.prompt_tokens
         answer_tokens = llm_message.completion_tokens
-        message_unit_price = llm_constant.model_prices[model_name]['prompt']
-        answer_unit_price = llm_constant.model_prices[model_name]['completion']
+        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)
 
@@ -162,8 +164,6 @@ class ConversationMessageTask:
         self.message.provider_response_latency = llm_message.latency
         self.message.total_price = total_price
 
-        self.update_provider_quota()
-
         db.session.commit()
 
         message_was_created.send(
@@ -175,20 +175,6 @@ class ConversationMessageTask:
         if not by_stopped:
             self.end()
 
-    def update_provider_quota(self):
-        llm_provider_service = LLMProviderService(
-            tenant_id=self.app.tenant_id,
-            provider_name=self.message.model_provider,
-        )
-
-        provider = llm_provider_service.get_provider_db_record()
-        if provider and provider.provider_type == ProviderType.SYSTEM.value:
-            db.session.query(Provider).filter(
-                Provider.tenant_id == self.app.tenant_id,
-                Provider.provider_name == provider.provider_name,
-                Provider.quota_limit > Provider.quota_used
-            ).update({'quota_used': Provider.quota_used + 1})
-
     def init_chain(self, chain_result: ChainResult):
         message_chain = MessageChain(
             message_id=self.message.id,
@@ -207,10 +193,31 @@ class ConversationMessageTask:
 
         self._pub_handler.pub_chain(message_chain)
 
-    def on_agent_end(self, message_chain: MessageChain, agent_model_name: str,
+    def on_agent_start(self, message_chain: MessageChain, agent_loop: AgentLoop) -> MessageAgentThought:
+        message_agent_thought = MessageAgentThought(
+            message_id=self.message.id,
+            message_chain_id=message_chain.id,
+            position=agent_loop.position,
+            thought=agent_loop.thought,
+            tool=agent_loop.tool_name,
+            tool_input=agent_loop.tool_input,
+            message=agent_loop.prompt,
+            answer=agent_loop.completion,
+            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()
+
+        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,
                      agent_loop: AgentLoop):
-        agent_message_unit_price = llm_constant.model_prices[agent_model_name]['prompt']
-        agent_answer_unit_price = llm_constant.model_prices[agent_model_name]['completion']
+        agent_message_unit_price = agent_model_instant.get_token_price(1, MessageType.HUMAN)
+        agent_answer_unit_price = agent_model_instant.get_token_price(1, MessageType.ASSISTANT)
 
         loop_message_tokens = agent_loop.prompt_tokens
         loop_answer_tokens = agent_loop.completion_tokens
@@ -222,34 +229,18 @@ class ConversationMessageTask:
             agent_answer_unit_price
         )
 
-        message_agent_loop = MessageAgentThought(
-            message_id=self.message.id,
-            message_chain_id=message_chain.id,
-            position=agent_loop.position,
-            thought=agent_loop.thought,
-            tool=agent_loop.tool_name,
-            tool_input=agent_loop.tool_input,
-            observation=agent_loop.tool_output,
-            tool_process_data='',  # currently not support
-            message=agent_loop.prompt,
-            message_token=loop_message_tokens,
-            message_unit_price=agent_message_unit_price,
-            answer=agent_loop.completion,
-            answer_token=loop_answer_tokens,
-            answer_unit_price=agent_answer_unit_price,
-            latency=agent_loop.latency,
-            tokens=agent_loop.prompt_tokens + agent_loop.completion_tokens,
-            total_price=loop_total_price,
-            currency=llm_constant.model_currency,
-            created_by_role=('account' if isinstance(self.user, Account) else 'end_user'),
-            created_by=self.user.id
-        )
-
-        db.session.add(message_agent_loop)
+        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.answer_token = loop_answer_tokens
+        message_agent_thought.answer_unit_price = agent_answer_unit_price
+        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()
         db.session.flush()
 
-        self._pub_handler.pub_agent_thought(message_agent_loop)
-
     def on_dataset_query_end(self, dataset_query_obj: DatasetQueryObj):
         dataset_query = DatasetQuery(
             dataset_id=dataset_query_obj.dataset_id,
@@ -346,16 +337,14 @@ class PubHandler:
             content = {
                 'event': 'agent_thought',
                 'data': {
+                    'id': message_agent_thought.id,
                     'task_id': self._task_id,
                     'message_id': self._message.id,
                     'chain_id': message_agent_thought.message_chain_id,
-                    'agent_thought_id': message_agent_thought.id,
                     'position': message_agent_thought.position,
                     'thought': message_agent_thought.thought,
                     'tool': message_agent_thought.tool,
                     'tool_input': message_agent_thought.tool_input,
-                    'observation': message_agent_thought.observation,
-                    'answer': message_agent_thought.answer,
                     'mode': self._conversation.mode,
                     'conversation_id': self._conversation.id
                 }
@@ -388,6 +377,15 @@ class PubHandler:
     def _is_stopped(self):
         return redis_client.get(self._stopped_cache_key) is not None
 
+    @classmethod
+    def ping(cls, user: Union[Account | EndUser], task_id: str):
+        content = {
+            'event': 'ping'
+        }
+
+        channel = cls.generate_channel_name(user, task_id)
+        redis_client.publish(channel, json.dumps(content))
+
     @classmethod
     def stop(cls, user: Union[Account | EndUser], task_id: str):
         stopped_cache_key = cls.generate_stopped_cache_key(user, task_id)

api/core/data_loader/file_extractor.py

@@ -1,7 +1,8 @@
 import tempfile
 from pathlib import Path
-from typing import List, Union
+from typing import List, Union, Optional
 
+import requests
 from langchain.document_loaders import TextLoader, Docx2txtLoader
 from langchain.schema import Document
 
@@ -13,6 +14,9 @@ from core.data_loader.loader.pdf import PdfLoader
 from extensions.ext_storage import storage
 from models.model import UploadFile
 
+SUPPORT_URL_CONTENT_TYPES = ['application/pdf', 'text/plain']
+USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+
 
 class FileExtractor:
     @classmethod
@@ -22,22 +26,41 @@ class FileExtractor:
             file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
             storage.download(upload_file.key, file_path)
 
-            input_file = Path(file_path)
-            delimiter = '\n'
-            if input_file.suffix == '.xlsx':
-                loader = ExcelLoader(file_path)
-            elif input_file.suffix == '.pdf':
-                loader = PdfLoader(file_path, upload_file=upload_file)
-            elif input_file.suffix in ['.md', '.markdown']:
-                loader = MarkdownLoader(file_path, autodetect_encoding=True)
-            elif input_file.suffix in ['.htm', '.html']:
-                loader = HTMLLoader(file_path)
-            elif input_file.suffix == '.docx':
-                loader = Docx2txtLoader(file_path)
-            elif input_file.suffix == '.csv':
-                loader = CSVLoader(file_path, autodetect_encoding=True)
-            else:
-                # txt
-                loader = TextLoader(file_path, autodetect_encoding=True)
+            return cls.load_from_file(file_path, return_text, upload_file)
 
-            return delimiter.join([document.page_content for document in loader.load()]) if return_text else loader.load()
+    @classmethod
+    def load_from_url(cls, url: str, return_text: bool = False) -> Union[List[Document] | str]:
+        response = requests.get(url, headers={
+            "User-Agent": USER_AGENT
+        })
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            suffix = Path(url).suffix
+            file_path = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
+            with open(file_path, 'wb') as file:
+                file.write(response.content)
+
+            return cls.load_from_file(file_path, return_text)
+
+    @classmethod
+    def load_from_file(cls, file_path: str, return_text: bool = False,
+                       upload_file: Optional[UploadFile] = None) -> Union[List[Document] | str]:
+        input_file = Path(file_path)
+        delimiter = '\n'
+        if input_file.suffix == '.xlsx':
+            loader = ExcelLoader(file_path)
+        elif input_file.suffix == '.pdf':
+            loader = PdfLoader(file_path, upload_file=upload_file)
+        elif input_file.suffix in ['.md', '.markdown']:
+            loader = MarkdownLoader(file_path, autodetect_encoding=True)
+        elif input_file.suffix in ['.htm', '.html']:
+            loader = HTMLLoader(file_path)
+        elif input_file.suffix == '.docx':
+            loader = Docx2txtLoader(file_path)
+        elif input_file.suffix == '.csv':
+            loader = CSVLoader(file_path, autodetect_encoding=True)
+        else:
+            # txt
+            loader = TextLoader(file_path, autodetect_encoding=True)
+
+        return delimiter.join([document.page_content for document in loader.load()]) if return_text else loader.load()

api/core/data_loader/loader/excel.py

@@ -39,7 +39,7 @@ class ExcelLoader(BaseLoader):
                     row_dict = dict(zip(keys, list(map(str, row))))
                     row_dict = {k: v for k, v in row_dict.items() if v}
                     item = ''.join(f'{k}:{v}\n' for k, v in row_dict.items())
-                    document = Document(page_content=item)
+                    document = Document(page_content=item, metadata={'source': self._file_path})
                     data.append(document)
 
         return data

api/core/docstore/dataset_docstore.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, Optional, Sequence
 from langchain.schema import Document
 from sqlalchemy import func
 
-from core.llm.token_calculator import TokenCalculator
+from core.model_providers.model_factory import ModelFactory
 from extensions.ext_database import db
 from models.dataset import Dataset, DocumentSegment
 
@@ -13,12 +13,10 @@ class DatesetDocumentStore:
         self,
         dataset: Dataset,
         user_id: str,
-        embedding_model_name: str,
         document_id: Optional[str] = None,
     ):
         self._dataset = dataset
         self._user_id = user_id
-        self._embedding_model_name = embedding_model_name
         self._document_id = document_id
 
     @classmethod
@@ -39,10 +37,6 @@ class DatesetDocumentStore:
     def user_id(self) -> Any:
         return self._user_id
 
-    @property
-    def embedding_model_name(self) -> Any:
-        return self._embedding_model_name
-
     @property
     def docs(self) -> Dict[str, Document]:
         document_segments = db.session.query(DocumentSegment).filter(
@@ -68,12 +62,16 @@ class DatesetDocumentStore:
         self, docs: Sequence[Document], allow_update: bool = True
     ) -> None:
         max_position = db.session.query(func.max(DocumentSegment.position)).filter(
-            DocumentSegment.document == self._document_id
+            DocumentSegment.document_id == self._document_id
         ).scalar()
 
         if max_position is None:
             max_position = 0
 
+        embedding_model = ModelFactory.get_embedding_model(
+            tenant_id=self._dataset.tenant_id
+        )
+
         for doc in docs:
             if not isinstance(doc, Document):
                 raise ValueError("doc must be a Document")
@@ -88,7 +86,7 @@ class DatesetDocumentStore:
                 )
 
             # calc embedding use tokens
-            tokens = TokenCalculator.get_num_tokens(self._embedding_model_name, doc.page_content)
+            tokens = embedding_model.get_num_tokens(doc.page_content)
 
             if not segment_document:
                 max_position += 1
@@ -105,9 +103,14 @@ class DatesetDocumentStore:
                     tokens=tokens,
                     created_by=self._user_id,
                 )
+                if 'answer' in doc.metadata and doc.metadata['answer']:
+                    segment_document.answer = doc.metadata.pop('answer', '')
+
                 db.session.add(segment_document)
             else:
                 segment_document.content = doc.page_content
+                if 'answer' in doc.metadata and doc.metadata['answer']:
+                    segment_document.answer = doc.metadata.pop('answer', '')
                 segment_document.index_node_hash = doc.metadata['doc_hash']
                 segment_document.word_count = len(doc.page_content)
                 segment_document.tokens = tokens

api/core/embedding/cached_embedding.py

@@ -4,14 +4,14 @@ from typing import List
 from langchain.embeddings.base import Embeddings
 from sqlalchemy.exc import IntegrityError
 
-from core.llm.wrappers.openai_wrapper import handle_openai_exceptions
+from core.model_providers.models.embedding.base import BaseEmbedding
 from extensions.ext_database import db
 from libs import helper
 from models.dataset import Embedding
 
 
 class CacheEmbedding(Embeddings):
-    def __init__(self, embeddings: Embeddings):
+    def __init__(self, embeddings: BaseEmbedding):
         self._embeddings = embeddings
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
@@ -21,48 +21,54 @@ class CacheEmbedding(Embeddings):
         embedding_queue_texts = []
         for text in texts:
             hash = helper.generate_text_hash(text)
-            embedding = db.session.query(Embedding).filter_by(hash=hash).first()
+            embedding = db.session.query(Embedding).filter_by(model_name=self._embeddings.name, hash=hash).first()
             if embedding:
                 text_embeddings.append(embedding.get_embedding())
             else:
                 embedding_queue_texts.append(text)
 
-        embedding_results = self._embeddings.embed_documents(embedding_queue_texts)
-
-        i = 0
-        for text in embedding_queue_texts:
-            hash = helper.generate_text_hash(text)
-
+        if embedding_queue_texts:
             try:
-                embedding = Embedding(hash=hash)
-                embedding.set_embedding(embedding_results[i])
-                db.session.add(embedding)
-                db.session.commit()
-            except IntegrityError:
-                db.session.rollback()
-                continue
-            except:
-                logging.exception('Failed to add embedding to db')
-                continue
+                embedding_results = self._embeddings.client.embed_documents(embedding_queue_texts)
+            except Exception as ex:
+                raise self._embeddings.handle_exceptions(ex)
 
-            i += 1
+            i = 0
+            for text in embedding_queue_texts:
+                hash = helper.generate_text_hash(text)
 
-        text_embeddings.extend(embedding_results)
+                try:
+                    embedding = Embedding(model_name=self._embeddings.name, hash=hash)
+                    embedding.set_embedding(embedding_results[i])
+                    db.session.add(embedding)
+                    db.session.commit()
+                except IntegrityError:
+                    db.session.rollback()
+                    continue
+                except:
+                    logging.exception('Failed to add embedding to db')
+                    continue
+                finally:
+                    i += 1
+
+            text_embeddings.extend(embedding_results)
         return text_embeddings
 
-    @handle_openai_exceptions
     def embed_query(self, text: str) -> List[float]:
         """Embed query text."""
         # use doc embedding cache or store if not exists
         hash = helper.generate_text_hash(text)
-        embedding = db.session.query(Embedding).filter_by(hash=hash).first()
+        embedding = db.session.query(Embedding).filter_by(model_name=self._embeddings.name, hash=hash).first()
         if embedding:
             return embedding.get_embedding()
 
-        embedding_results = self._embeddings.embed_query(text)
+        try:
+            embedding_results = self._embeddings.client.embed_query(text)
+        except Exception as ex:
+            raise self._embeddings.handle_exceptions(ex)
 
         try:
-            embedding = Embedding(hash=hash)
+            embedding = Embedding(model_name=self._embeddings.name, hash=hash)
             embedding.set_embedding(embedding_results)
             db.session.add(embedding)
             db.session.commit()
@@ -72,3 +78,5 @@ class CacheEmbedding(Embeddings):
             logging.exception('Failed to add embedding to db')
 
         return embedding_results
+
+

api/core/generator/llm_generator.py

@@ -1,22 +1,17 @@
 import logging
 
-from langchain import PromptTemplate
-from langchain.chat_models.base import BaseChatModel
-from langchain.schema import HumanMessage, OutputParserException, BaseMessage
+from langchain.schema import OutputParserException
 
-from core.constant import llm_constant
-from core.llm.llm_builder import LLMBuilder
-from core.llm.streamable_open_ai import StreamableOpenAI
-from core.llm.token_calculator import TokenCalculator
+from core.model_providers.error import LLMError
+from core.model_providers.model_factory import ModelFactory
+from core.model_providers.models.entity.message import PromptMessage, MessageType
+from core.model_providers.models.entity.model_params import ModelKwargs
 from core.prompt.output_parser.rule_config_generator import RuleConfigGeneratorOutputParser
 
 from core.prompt.output_parser.suggested_questions_after_answer import SuggestedQuestionsAfterAnswerOutputParser
 from core.prompt.prompt_template import JinjaPromptTemplate, OutLinePromptTemplate
-from core.prompt.prompts import CONVERSATION_TITLE_PROMPT, CONVERSATION_SUMMARY_PROMPT, INTRODUCTION_GENERATE_PROMPT
-
-
-# gpt-3.5-turbo works not well
-generate_base_model = 'text-davinci-003'
+from core.prompt.prompts import CONVERSATION_TITLE_PROMPT, CONVERSATION_SUMMARY_PROMPT, INTRODUCTION_GENERATE_PROMPT, \
+    GENERATOR_QA_PROMPT
 
 
 class LLMGenerator:
@@ -28,28 +23,35 @@ class LLMGenerator:
             query = query[:300] + "...[TRUNCATED]..." + query[-300:]
 
         prompt = prompt.format(query=query)
-        llm: StreamableOpenAI = LLMBuilder.to_llm(
+
+        model_instance = ModelFactory.get_text_generation_model(
             tenant_id=tenant_id,
-            model_name='gpt-3.5-turbo',
-            max_tokens=50
+            model_kwargs=ModelKwargs(
+                max_tokens=50
+            )
         )
 
-        if isinstance(llm, BaseChatModel):
-            prompt = [HumanMessage(content=prompt)]
-
-        response = llm.generate([prompt])
-        answer = response.generations[0][0].text
+        prompts = [PromptMessage(content=prompt)]
+        response = model_instance.run(prompts)
+        answer = response.content
         return answer.strip()
 
     @classmethod
     def generate_conversation_summary(cls, tenant_id: str, messages):
         max_tokens = 200
-        model = 'gpt-3.5-turbo'
+
+        model_instance = ModelFactory.get_text_generation_model(
+            tenant_id=tenant_id,
+            model_kwargs=ModelKwargs(
+                max_tokens=max_tokens
+            )
+        )
 
         prompt = CONVERSATION_SUMMARY_PROMPT
         prompt_with_empty_context = prompt.format(context='')
-        prompt_tokens = TokenCalculator.get_num_tokens(model, prompt_with_empty_context)
-        rest_tokens = llm_constant.max_context_token_length[model] - prompt_tokens - max_tokens - 1
+        prompt_tokens = model_instance.get_num_tokens([PromptMessage(content=prompt_with_empty_context)])
+        max_context_token_length = model_instance.model_rules.max_tokens.max
+        rest_tokens = max_context_token_length - prompt_tokens - max_tokens - 1
 
         context = ''
         for message in messages:
@@ -67,25 +69,16 @@ class LLMGenerator:
                 answer = message.answer
 
             message_qa_text = "\n\nHuman:" + query + "\n\nAssistant:" + answer
-            if rest_tokens - TokenCalculator.get_num_tokens(model, context + message_qa_text) > 0:
+            if rest_tokens - model_instance.get_num_tokens([PromptMessage(content=context + message_qa_text)]) > 0:
                 context += message_qa_text
 
         if not context:
             return '[message too long, no summary]'
 
         prompt = prompt.format(context=context)
-
-        llm: StreamableOpenAI = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name=model,
-            max_tokens=max_tokens
-        )
-
-        if isinstance(llm, BaseChatModel):
-            prompt = [HumanMessage(content=prompt)]
-
-        response = llm.generate([prompt])
-        answer = response.generations[0][0].text
+        prompts = [PromptMessage(content=prompt)]
+        response = model_instance.run(prompts)
+        answer = response.content
         return answer.strip()
 
     @classmethod
@@ -93,16 +86,13 @@ class LLMGenerator:
         prompt = INTRODUCTION_GENERATE_PROMPT
         prompt = prompt.format(prompt=pre_prompt)
 
-        llm: StreamableOpenAI = LLMBuilder.to_llm(
-            tenant_id=tenant_id,
-            model_name=generate_base_model,
+        model_instance = ModelFactory.get_text_generation_model(
+            tenant_id=tenant_id
         )
 
-        if isinstance(llm, BaseChatModel):
-            prompt = [HumanMessage(content=prompt)]
-
-        response = llm.generate([prompt])
-        answer = response.generations[0][0].text
+        prompts = [PromptMessage(content=prompt)]
+        response = model_instance.run(prompts)
+        answer = response.content
         return answer.strip()
 
     @classmethod
@@ -118,25 +108,23 @@ class LLMGenerator:
 
         _input = prompt.format_prompt(histories=histories)
 
-        llm: StreamableOpenAI = LLMBuilder.to_llm(
+        model_instance = ModelFactory.get_text_generation_model(
             tenant_id=tenant_id,
-            model_name='gpt-3.5-turbo',
-            temperature=0,
-            max_tokens=256
+            model_kwargs=ModelKwargs(
+                max_tokens=256,
+                temperature=0
+            )
         )
 
-        if isinstance(llm, BaseChatModel):
-            query = [HumanMessage(content=_input.to_string())]
-        else:
-            query = _input.to_string()
+        prompts = [PromptMessage(content=_input.to_string())]
 
         try:
-            output = llm(query)
-            if isinstance(output, BaseMessage):
-                output = output.content
-            questions = output_parser.parse(output)
-        except Exception:
-            logging.exception("Error generating suggested questions after answer")
+            output = model_instance.run(prompts)
+            questions = output_parser.parse(output.content)
+        except LLMError:
+            questions = []
+        except Exception as e:
+            logging.exception(e)
             questions = []
 
         return questions
@@ -159,25 +147,25 @@ class LLMGenerator:
 
         _input = prompt.format_prompt(audiences=audiences, hoping_to_solve=hoping_to_solve)
 
-        llm: StreamableOpenAI = LLMBuilder.to_llm(
+        model_instance = ModelFactory.get_text_generation_model(
             tenant_id=tenant_id,
-            model_name=generate_base_model,
-            temperature=0,
-            max_tokens=512
+            model_kwargs=ModelKwargs(
+                max_tokens=512,
+                temperature=0
+            )
         )
 
-        if isinstance(llm, BaseChatModel):
-            query = [HumanMessage(content=_input.to_string())]
-        else:
-            query = _input.to_string()
+        prompts = [PromptMessage(content=_input.to_string())]
 
         try:
-            output = llm(query)
-            rule_config = output_parser.parse(output)
+            output = model_instance.run(prompts)
+            rule_config = output_parser.parse(output.content)
+        except LLMError as e:
+            raise e
         except OutputParserException:
             raise ValueError('Please give a valid input for intended audience or hoping to solve problems.')
-        except Exception:
-            logging.exception("Error generating prompt")
+        except Exception as e:
+            logging.exception(e)
             rule_config = {
                 "prompt": "",
                 "variables": [],
@@ -185,3 +173,23 @@ class LLMGenerator:
             }
 
         return rule_config
+
+    @classmethod
+    def generate_qa_document(cls, tenant_id: str, query):
+        prompt = GENERATOR_QA_PROMPT
+
+        model_instance = ModelFactory.get_text_generation_model(
+            tenant_id=tenant_id,
+            model_kwargs=ModelKwargs(
+                max_tokens=2000
+            )
+        )
+
+        prompts = [
+            PromptMessage(content=prompt, type=MessageType.SYSTEM),
+            PromptMessage(content=query)
+        ]
+
+        response = model_instance.run(prompts)
+        answer = response.content
+        return answer.strip()

api/core/helper/encrypter.py

@@ -0,0 +1,20 @@
+import base64
+
+from extensions.ext_database import db
+from libs import rsa
+
+from models.account import Tenant
+
+
+def obfuscated_token(token: str):
+    return token[:6] + '*' * (len(token) - 8) + token[-2:]
+
+
+def encrypt_token(tenant_id: str, token: str):
+    tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).first()
+    encrypted_token = rsa.encrypt(token, tenant.encrypt_public_key)
+    return base64.b64encode(encrypted_token).decode()
+
+
+def decrypt_token(tenant_id: str, token: str):
+    return rsa.decrypt(base64.b64decode(token), tenant_id)

api/core/index/index.py

@@ -1,10 +1,9 @@
 from flask import current_app
-from langchain.embeddings import OpenAIEmbeddings
 
 from core.embedding.cached_embedding import CacheEmbedding
 from core.index.keyword_table_index.keyword_table_index import KeywordTableIndex, KeywordTableConfig
 from core.index.vector_index.vector_index import VectorIndex
-from core.llm.llm_builder import LLMBuilder
+from core.model_providers.model_factory import ModelFactory
 from models.dataset import Dataset
 
 
@@ -15,15 +14,11 @@ class IndexBuilder:
             if not ignore_high_quality_check and dataset.indexing_technique != 'high_quality':
                 return None
 
-            model_credentials = LLMBuilder.get_model_credentials(
-                tenant_id=dataset.tenant_id,
-                model_provider=LLMBuilder.get_default_provider(dataset.tenant_id, 'text-embedding-ada-002'),
-                model_name='text-embedding-ada-002'
+            embedding_model = ModelFactory.get_embedding_model(
+                tenant_id=dataset.tenant_id
             )
 
-            embeddings = CacheEmbedding(OpenAIEmbeddings(
-                **model_credentials
-            ))
+            embeddings = CacheEmbedding(embedding_model)
 
             return VectorIndex(
                 dataset=dataset,