feat: When voice is turned on, the page will not display an empty reply message when the answer is empty #1877 (#2447)

### What problem does this PR solve?

feat: When voice is turned on, the page will not display an empty reply
message when the answer is empty #1877

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Dockerfile.arm

@@ -4,6 +4,10 @@ USER  root
 WORKDIR /ragflow
 
 COPY requirements_arm.txt /ragflow/requirements.txt
+
+
+RUN pip install nltk --default-timeout=10000
+
 RUN pip install -i https://mirrors.aliyun.com/pypi/simple/ --default-timeout=1000 -r requirements.txt &&\
 python -c "import nltk;nltk.download('punkt');nltk.download('wordnet')"
 
@@ -14,6 +18,11 @@ RUN apt-get update && \
 RUN curl -sL https://deb.nodesource.com/setup_20.x | bash - && \
     apt-get install -y --fix-missing nodejs nginx ffmpeg libsm6 libxext6 libgl1
 
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+RUN pip install graspologic
+
 ADD ./web ./web
 RUN cd ./web && npm i --force && npm run build
 

README.md

@@ -18,7 +18,7 @@
     <a href="https://demo.ragflow.io" target="_blank">
         <img alt="Static Badge" src="https://img.shields.io/badge/Online-Demo-4e6b99"></a>
     <a href="https://hub.docker.com/r/infiniflow/ragflow" target="_blank">
-        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.10.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.10.0"></a>
+        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.11.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.11.0"></a>
     <a href="https://github.com/infiniflow/ragflow/blob/main/LICENSE">
     <img height="21" src="https://img.shields.io/badge/License-Apache--2.0-ffffff?labelColor=d4eaf7&color=2e6cc4" alt="license">
   </a>
@@ -66,24 +66,15 @@ Try our demo at [https://demo.ragflow.io](https://demo.ragflow.io).
 
 ## 🔥 Latest Updates
 
+- 2024-09-13 Adds search mode for knowledge base Q&A.
+- 2024-09-09 Adds a medical consultant agent template.
 - 2024-08-22 Support text to SQL statements through RAG.
-  
 - 2024-08-02 Supports GraphRAG inspired by [graphrag](https://github.com/microsoft/graphrag) and mind map.
-  
 - 2024-07-23 Supports audio file parsing.
-
+- 2024-07-08 Supports workflow based on [Graph](./agent/README.md).
-- 2024-07-21 Supports more LLMs (LocalAI, OpenRouter, StepFun, and Nvidia).
+- 2024-06-27 Supports Markdown and Docx in the Q&A parsing method, extracting images from Docx files, extracting tables from Markdown files.
-
-- 2024-07-18 Adds more components (Wikipedia, PubMed, Baidu, and Duckduckgo) to the graph.
-
-- 2024-07-08 Supports workflow based on [Graph](./graph/README.md).
-- 2024-06-27 Supports Markdown and Docx in the Q&A parsing method. 
-- 2024-06-27 Supports extracting images from Docx files. 
-- 2024-06-27 Supports extracting tables from Markdown files.
-- 2024-06-06 Supports [Self-RAG](https://huggingface.co/papers/2310.11511), which is enabled by default in dialog settings.
-- 2024-05-30 Integrates [BCE](https://github.com/netease-youdao/BCEmbedding) and [BGE](https://github.com/FlagOpen/FlagEmbedding) reranker models.
 - 2024-05-23 Supports [RAPTOR](https://arxiv.org/html/2401.18059v1) for better text retrieval.
-- 2024-05-15 Integrates OpenAI GPT-4o.
+
 
 ## 🌟 Key Features
 
@@ -160,7 +151,7 @@ Try our demo at [https://demo.ragflow.io](https://demo.ragflow.io).
 
 3. Build the pre-built Docker images and start up the server:
 
-   > Running the following commands automatically downloads the *dev* version RAGFlow Docker image. To download and run a specified Docker version, update `RAGFLOW_VERSION` in **docker/.env** to the intended version, for example `RAGFLOW_VERSION=v0.10.0`, before running the following commands.
+   > Running the following commands automatically downloads the *dev* version RAGFlow Docker image. To download and run a specified Docker version, update `RAGFLOW_VERSION` in **docker/.env** to the intended version, for example `RAGFLOW_VERSION=v0.11.0`, before running the following commands.
 
    ```bash
    $ cd ragflow/docker
@@ -192,7 +183,7 @@ Try our demo at [https://demo.ragflow.io](https://demo.ragflow.io).
     * Running on http://x.x.x.x:9380
     INFO:werkzeug:Press CTRL+C to quit
    ```
-   > If you skip this confirmation step and directly log in to RAGFlow, your browser may prompt a `network anomaly` error because, at that moment, your RAGFlow may not be fully initialized.  
+   > If you skip this confirmation step and directly log in to RAGFlow, your browser may prompt a `network abnormal` error because, at that moment, your RAGFlow may not be fully initialized.  
 
 5. In your web browser, enter the IP address of your server and log in to RAGFlow.
    > With the default settings, you only need to enter `http://IP_OF_YOUR_MACHINE` (**sans** port number) as the default HTTP serving port `80` can be omitted when using the default configurations.

README_ja.md

@@ -18,8 +18,8 @@
     <a href="https://demo.ragflow.io" target="_blank">
         <img alt="Static Badge" src="https://img.shields.io/badge/Online-Demo-4e6b99"></a>
     <a href="https://hub.docker.com/r/infiniflow/ragflow" target="_blank">
-        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.10.0-brightgreen"
+        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.11.0-brightgreen"
-            alt="docker pull infiniflow/ragflow:v0.10.0"></a>
+            alt="docker pull infiniflow/ragflow:v0.11.0"></a>
       <a href="https://github.com/infiniflow/ragflow/blob/main/LICENSE">
     <img height="21" src="https://img.shields.io/badge/License-Apache--2.0-ffffff?labelColor=d4eaf7&color=2e6cc4" alt="license">
   </a>
@@ -48,19 +48,15 @@
 
 ## 🔥 最新情報
 
+- 2024-09-13 ナレッジベース Q&A の検索モードを追加しました。
+- 2024-09-09 エージェントに医療相談テンプレートを追加しました。
 - 2024-08-22 RAG を介して SQL ステートメントへのテキストをサポートします。
 - 2024-08-02 [graphrag](https://github.com/microsoft/graphrag) からインスピレーションを得た GraphRAG とマインド マップをサポートします。
 - 2024-07-23 音声ファイルの解析をサポートしました。
-- 2024-07-21 より多くの LLM サプライヤー (LocalAI/OpenRouter/StepFun/Nvidia) をサポートします。
+- 2024-07-08 [Graph](./agent/README.md) ベースのワークフローをサポート
-- 2024-07-18 グラフにコンポーネント(Wikipedia/PubMed/Baidu/Duckduckgo)を追加しました。
+- 2024-06-27 Q&A 解析メソッドで Markdown と Docx をサポートし、Docx ファイルから画像を抽出し、Markdown ファイルからテーブルを抽出します。
-- 2024-07-08 [Graph](./graph/README.md) ベースのワークフローをサポート
-- 2024-06-27 Q&A解析方式はMarkdownファイルとDocxファイルをサポートしています。
-- 2024-06-27 Docxファイルからの画像の抽出をサポートします。
-- 2024-06-27 Markdownファイルからテーブルを抽出することをサポートします。
-- 2024-06-06 会話設定でデフォルトでチェックされている [Self-RAG](https://huggingface.co/papers/2310.11511) をサポートします。
-- 2024-05-30 [BCE](https://github.com/netease-youdao/BCEmbedding) 、[BGE](https://github.com/FlagOpen/FlagEmbedding) reranker を統合。
 - 2024-05-23 より良いテキスト検索のために [RAPTOR](https://arxiv.org/html/2401.18059v1) をサポート。
-- 2024-05-15 OpenAI GPT-4oを統合しました。
+
 
 ## 🌟 主な特徴
 
@@ -143,7 +139,7 @@
    $ docker compose up -d
    ```
 
-   > 上記のコマンドを実行すると、RAGFlowの開発版dockerイメージが自動的にダウンロードされます。 特定のバージョンのDockerイメージをダウンロードして実行したい場合は、docker/.envファイルのRAGFLOW_VERSION変数を見つけて、対応するバージョンに変更してください。 例えば、RAGFLOW_VERSION=v0.10.0として、上記のコマンドを実行してください。
+   > 上記のコマンドを実行すると、RAGFlowの開発版dockerイメージが自動的にダウンロードされます。 特定のバージョンのDockerイメージをダウンロードして実行したい場合は、docker/.envファイルのRAGFLOW_VERSION変数を見つけて、対応するバージョンに変更してください。 例えば、RAGFLOW_VERSION=v0.11.0として、上記のコマンドを実行してください。
 
    > コアイメージのサイズは約 9 GB で、ロードに時間がかかる場合があります。
 
@@ -205,7 +201,7 @@
 ```bash
 $ git clone https://github.com/infiniflow/ragflow.git
 $ cd ragflow/
-$ docker build -t infiniflow/ragflow:v0.10.0 .
+$ docker build -t infiniflow/ragflow:v0.11.0 .
 $ cd ragflow/docker
 $ chmod +x ./entrypoint.sh
 $ docker compose up -d

README_ko.md

@@ -18,7 +18,7 @@
     <a href="https://demo.ragflow.io" target="_blank">
         <img alt="Static Badge" src="https://img.shields.io/badge/Online-Demo-4e6b99"></a>
     <a href="https://hub.docker.com/r/infiniflow/ragflow" target="_blank">
-        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.10.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.10.0"></a>
+        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.11.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.11.0"></a>
     <a href="https://github.com/infiniflow/ragflow/blob/main/LICENSE">
     <img height="21" src="https://img.shields.io/badge/License-Apache--2.0-ffffff?labelColor=d4eaf7&color=2e6cc4" alt="license">
   </a>
@@ -49,31 +49,22 @@
 
 ## 🔥 업데이트
 
+- 2024-09-13 지식베이스 Q&A 검색 모드를 추가합니다.
+  
+- 2024-09-09 Agent에 의료상담 템플릿을 추가하였습니다.
+  
 - 2024-08-22 RAG를 통해 SQL 문에 텍스트를 지원합니다.
   
 - 2024-08-02: [graphrag](https://github.com/microsoft/graphrag)와 마인드맵에서 영감을 받은 GraphRAG를 지원합니다.
 
 - 2024-07-23: 오디오 파일 분석을 지원합니다.
 
-- 2024-07-21: 더 많은 LLMs(LocalAI, OpenRouter, StepFun, Nvidia)를 지원합니다.
+- 2024-07-08: [Graph](./agent/README.md)를 기반으로 한 워크플로우를 지원합니다.
 
-- 2024-07-18: 그래프에 더 많은 구성요소(Wikipedia, PubMed, Baidu, Duckduckgo)를 추가합니다.
+- 2024-06-27 Q&A 구문 분석 방식에서 Markdown 및 Docx를 지원하고, Docx 파일에서 이미지 추출, Markdown 파일에서 테이블 추출을 지원합니다.
-
-- 2024-07-08: [Graph](./graph/README.md)를 기반으로 한 워크플로우를 지원합니다.
-
-- 2024-06-27: Q&A 분석 방법에서 Markdown과 Docx를 지원합니다.
-
-- 2024-06-27: Docx 파일에서 이미지 추출을 지원합니다.
-
-- 2024-06-27: Markdown 파일에서 표 추출을 지원합니다.
-
-- 2024-06-06: 대화 설정에서 기본으로 [Self-RAG](https://huggingface.co/papers/2310.11511)를 지원합니다.
-
-- 2024-05-30: [BCE](https://github.com/netease-youdao/BCEmbedding) 및 [BGE](https://github.com/FlagOpen/FlagEmbedding) reranker 모델을 통합합니다.
 
 - 2024-05-23: 더 나은 텍스트 검색을 위해 [RAPTOR](https://arxiv.org/html/2401.18059v1)를 지원합니다.
 
-- 2024-05-15: OpenAI GPT-4o를 통합합니다.
 
 
 ## 🌟 주요 기능
@@ -147,7 +138,7 @@
 
 3. 미리 빌드된 Docker 이미지를 생성하고 서버를 시작하세요:
 
-   > 다음 명령어를 실행하면 *dev* 버전의 RAGFlow Docker 이미지가 자동으로 다운로드됩니다. 특정 Docker 버전을 다운로드하고 실행하려면, **docker/.env** 파일에서 `RAGFLOW_VERSION`을 원하는 버전으로 업데이트한 후, 예를 들어 `RAGFLOW_VERSION=v0.10.0`로 업데이트 한 뒤, 다음 명령어를 실행하세요.
+   > 다음 명령어를 실행하면 *dev* 버전의 RAGFlow Docker 이미지가 자동으로 다운로드됩니다. 특정 Docker 버전을 다운로드하고 실행하려면, **docker/.env** 파일에서 `RAGFLOW_VERSION`을 원하는 버전으로 업데이트한 후, 예를 들어 `RAGFLOW_VERSION=v0.11.0`로 업데이트 한 뒤, 다음 명령어를 실행하세요.
    ```bash
    $ cd ragflow/docker
    $ chmod +x ./entrypoint.sh
@@ -178,7 +169,7 @@
     * Running on http://x.x.x.x:9380
     INFO:werkzeug:Press CTRL+C to quit
    ```
-   > 만약 확인 단계를 건너뛰고 바로 RAGFlow에 로그인하면, RAGFlow가 완전히 초기화되지 않았기 때문에 브라우저에서 `network anomaly` 오류가 발생할 수 있습니다.
+   > 만약 확인 단계를 건너뛰고 바로 RAGFlow에 로그인하면, RAGFlow가 완전히 초기화되지 않았기 때문에 브라우저에서 `network abnormal` 오류가 발생할 수 있습니다.
 
 5. 웹 브라우저에 서버의 IP 주소를 입력하고 RAGFlow에 로그인하세요.
    > 기본 설정을 사용할 경우, `http://IP_OF_YOUR_MACHINE`만 입력하면 됩니다 (포트 번호는 제외). 기본 HTTP 서비스 포트 `80`은 기본 구성으로 사용할 때 생략할 수 있습니다.

README_zh.md

@@ -18,7 +18,7 @@
     <a href="https://demo.ragflow.io" target="_blank">
         <img alt="Static Badge" src="https://img.shields.io/badge/Online-Demo-4e6b99"></a>
     <a href="https://hub.docker.com/r/infiniflow/ragflow" target="_blank">
-        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.10.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.10.0"></a>
+        <img src="https://img.shields.io/badge/docker_pull-ragflow:v0.11.0-brightgreen" alt="docker pull infiniflow/ragflow:v0.11.0"></a>
     <a href="https://github.com/infiniflow/ragflow/blob/main/LICENSE">
     <img height="21" src="https://img.shields.io/badge/License-Apache--2.0-ffffff?labelColor=d4eaf7&color=2e6cc4" alt="license">
   </a>
@@ -47,19 +47,14 @@
 
 ## 🔥 近期更新
 
+- 2024-09-13 增加知识库问答搜索模式。
+- 2024-09-09 在 Agent 中加入医疗问诊模板。
 - 2024-08-22 支持用RAG技术实现从自然语言到SQL语句的转换。
 - 2024-08-02 支持 GraphRAG 启发于 [graphrag](https://github.com/microsoft/graphrag) 和思维导图。
 - 2024-07-23 支持解析音频文件。
-- 2024-07-21 支持更多的大模型供应商(LocalAI/OpenRouter/StepFun/Nvidia)。
+- 2024-07-08 支持 Agentic RAG: 基于 [Graph](./agent/README.md) 的工作流。
-- 2024-07-18 在Graph中支持算子：Wikipedia、PubMed、Baidu和Duckduckgo。
+- 2024-06-27 Q&A 解析方式支持 Markdown 文件和 Docx 文件，支持提取出 Docx 文件中的图片和 Markdown 文件中的表格。
-- 2024-07-08 支持 Agentic RAG: 基于 [Graph](./graph/README.md) 的工作流。
-- 2024-06-27 Q&A 解析方式支持 Markdown 文件和 Docx 文件。
-- 2024-06-27 支持提取出 Docx 文件中的图片。
-- 2024-06-27 支持提取出 Markdown 文件中的表格。
-- 2024-06-06 支持 [Self-RAG](https://huggingface.co/papers/2310.11511) ，在对话设置里面默认勾选。
-- 2024-05-30 集成 [BCE](https://github.com/netease-youdao/BCEmbedding) 和 [BGE](https://github.com/FlagOpen/FlagEmbedding) 重排序模型。
 - 2024-05-23 实现 [RAPTOR](https://arxiv.org/html/2401.18059v1) 提供更好的文本检索。
-- 2024-05-15 集成大模型 OpenAI GPT-4o。
 
 ## 🌟 主要功能
 
@@ -142,7 +137,7 @@
    $ docker compose -f docker-compose-CN.yml up -d
    ```
 
-   > 请注意，运行上述命令会自动下载 RAGFlow 的开发版本 docker 镜像。如果你想下载并运行特定版本的 docker 镜像，请在 docker/.env 文件中找到 RAGFLOW_VERSION 变量，将其改为对应版本。例如 RAGFLOW_VERSION=v0.10.0，然后运行上述命令。
+   > 请注意，运行上述命令会自动下载 RAGFlow 的开发版本 docker 镜像。如果你想下载并运行特定版本的 docker 镜像，请在 docker/.env 文件中找到 RAGFLOW_VERSION 变量，将其改为对应版本。例如 RAGFLOW_VERSION=v0.11.0，然后运行上述命令。
 
    > 核心镜像文件大约 9 GB，可能需要一定时间拉取。请耐心等待。
 
@@ -167,7 +162,7 @@
     * Running on http://x.x.x.x:9380
     INFO:werkzeug:Press CTRL+C to quit
    ```
-   > 如果您跳过这一步系统确认步骤就登录 RAGFlow，你的浏览器有可能会提示 `network anomaly` 或 `网络异常`，因为 RAGFlow 可能并未完全启动成功。
+   > 如果您跳过这一步系统确认步骤就登录 RAGFlow，你的浏览器有可能会提示 `network abnormal` 或 `网络异常`，因为 RAGFlow 可能并未完全启动成功。
 
 5. 在你的浏览器中输入你的服务器对应的 IP 地址并登录 RAGFlow。
    > 上面这个例子中，您只需输入 http://IP_OF_YOUR_MACHINE 即可：未改动过配置则无需输入端口（默认的 HTTP 服务端口 80）。
@@ -204,7 +199,7 @@
 ```bash
 $ git clone https://github.com/infiniflow/ragflow.git
 $ cd ragflow/
-$ docker build -t infiniflow/ragflow:v0.10.0 .
+$ docker build -t infiniflow/ragflow:v0.11.0 .
 $ cd ragflow/docker
 $ chmod +x ./entrypoint.sh
 $ docker compose up -d

SECURITY.md

@@ -18,7 +18,7 @@ main
 ### Actual behavior
 
 The restricted_loads function at [api/utils/__init__.py#L215](https://github.com/infiniflow/ragflow/blob/main/api/utils/__init__.py#L215) is still vulnerable leading via code execution.
-The main reson is that numpy module has a numpy.f2py.diagnose.run_command function directly execute commands, but the restricted_loads function allows users import functions in module numpy.
+The main reason is that numpy module has a numpy.f2py.diagnose.run_command function directly execute commands, but the restricted_loads function allows users import functions in module numpy.
 
 
 ### Steps to reproduce

agent/canvas.py

@@ -260,7 +260,7 @@ class Canvas(ABC):
 
     def get_history(self, window_size):
         convs = []
-        for role, obj in self.history[window_size * -2:]:
+        for role, obj in self.history[(window_size + 1) * -1:]:
             convs.append({"role": role, "content": (obj if role == "user" else
                                                     '\n'.join(pd.DataFrame(obj)['content']))})
         return convs
@@ -300,3 +300,6 @@ class Canvas(ABC):
                 return pat + " => " + pat
 
         return False
+
+    def get_prologue(self):
+        return self.components["begin"]["obj"]._param.prologue

agent/component/__init__.py

@@ -22,6 +22,12 @@ from .github import GitHub, GitHubParam
 from .baidufanyi import BaiduFanyi, BaiduFanyiParam
 from .qweather import QWeather, QWeatherParam
 from .exesql import ExeSQL, ExeSQLParam
+from .yahoofinance import YahooFinance, YahooFinanceParam
+from .wencai import WenCai, WenCaiParam
+from .jin10 import Jin10, Jin10Param
+from .tushare import TuShare, TuShareParam
+from .akshare import AkShare, AkShareParam
+
 
 def component_class(class_name):
     m = importlib.import_module("agent.component")

agent/component/akshare.py

@@ -0,0 +1,56 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+from abc import ABC
+import pandas as pd
+import akshare as ak
+from agent.component.base import ComponentBase, ComponentParamBase
+
+
+class AkShareParam(ComponentParamBase):
+    """
+    Define the AkShare component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.top_n = 10
+
+    def check(self):
+        self.check_positive_integer(self.top_n, "Top N")
+
+
+class AkShare(ComponentBase, ABC):
+    component_name = "AkShare"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = ",".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return AkShare.be_output("")
+
+        try:
+            ak_res = []
+            stock_news_em_df = ak.stock_news_em(symbol=ans)
+            stock_news_em_df = stock_news_em_df.head(self._param.top_n)
+            ak_res = [{"content": '<a href="' + i["新闻链接"] + '">' + i["新闻标题"] + '</a>\n 新闻内容: ' + i[
+                "新闻内容"] + " \n发布时间:" + i["发布时间"] + " \n文章来源: " + i["文章来源"]} for index, i in stock_news_em_df.iterrows()]
+        except Exception as e:
+            return AkShare.be_output("**ERROR**: " + str(e))
+
+        if not ak_res:
+            return AkShare.be_output("")
+
+        return pd.DataFrame(ak_res)

agent/component/categorize.py

@@ -82,6 +82,6 @@ class Categorize(Generate, ABC):
             if ans.lower().find(c.lower()) >= 0:
                 return Categorize.be_output(self._param.category_description[c]["to"])
 
-        return Categorize.be_output(self._param.category_description.items()[-1][1]["to"])
+        return Categorize.be_output(list(self._param.category_description.items())[-1][1]["to"])
 
 

agent/component/exesql.py

@@ -54,7 +54,7 @@ class ExeSQL(ComponentBase, ABC):
             setattr(self, "_loop", 0)
         if self._loop >= self._param.loop:
             self._loop = 0
-            raise Exception("Maximum loop time exceeds. Can't query the correct data via sql statement.")
+            raise Exception("Maximum loop time exceeds. Can't query the correct data via SQL statement.")
         self._loop += 1
 
         ans = self.get_input()
@@ -63,7 +63,7 @@ class ExeSQL(ComponentBase, ABC):
         ans = re.sub(r';.*?SELECT ', '; SELECT ', ans, flags=re.IGNORECASE)
         ans = re.sub(r';[^;]*$', r';', ans)
         if not ans:
-            return ExeSQL.be_output("SQL statement not found!")
+            raise Exception("SQL statement not found!")
 
         if self._param.db_type in ["mysql", "mariadb"]:
             db = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
@@ -75,13 +75,16 @@ class ExeSQL(ComponentBase, ABC):
         try:
             db.connect()
         except Exception as e:
-            return ExeSQL.be_output("**Error**: \nDatabase Connection Failed! \n" + str(e))
+            raise Exception("Database Connection Failed! \n" + str(e))
         sql_res = []
-        for single_sql in re.split(r';', ans):
+        for single_sql in re.split(r';', ans.replace(r"\n", " ")):
             if not single_sql:
                 continue
             try:
                 query = db.execute_sql(single_sql)
+                if query.rowcount == 0:
+                    sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n No record in the database!"})
+                    continue
                 single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
                 single_res.columns = [i[0] for i in query.description]
                 sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
@@ -91,6 +94,6 @@ class ExeSQL(ComponentBase, ABC):
         db.close()
 
         if not sql_res:
-            return ExeSQL.be_output("No record in the database!")
+            return ExeSQL.be_output("")
 
         return pd.DataFrame(sql_res)

agent/component/generate.py

@@ -112,8 +112,7 @@ class Generate(ComponentBase):
 
         kwargs["input"] = input
         for n, v in kwargs.items():
-            # prompt = re.sub(r"\{%s\}"%n, re.escape(str(v)), prompt)
+            prompt = re.sub(r"\{%s\}" % n, re.escape(str(v)), prompt)
-            prompt = re.sub(r"\{%s\}" % n, str(v), prompt)
 
         downstreams = self._canvas.get_component(self._id)["downstream"]
         if kwargs.get("stream") and len(downstreams) == 1 and self._canvas.get_component(downstreams[0])[

agent/component/jin10.py

@@ -0,0 +1,130 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+import json
+from abc import ABC
+import pandas as pd
+import requests
+from agent.component.base import ComponentBase, ComponentParamBase
+
+
+class Jin10Param(ComponentParamBase):
+    """
+    Define the Jin10 component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.type = "flash"
+        self.secret_key = "xxx"
+        self.flash_type = '1'
+        self.calendar_type = 'cj'
+        self.calendar_datatype = 'data'
+        self.symbols_type = 'GOODS'
+        self.symbols_datatype = 'symbols'
+        self.contain = ""
+        self.filter = ""
+
+    def check(self):
+        self.check_valid_value(self.type, "Type", ['flash', 'calendar', 'symbols', 'news'])
+        self.check_valid_value(self.flash_type, "Flash Type", ['1', '2', '3', '4', '5'])
+        self.check_valid_value(self.calendar_type, "Calendar Type", ['cj', 'qh', 'hk', 'us'])
+        self.check_valid_value(self.calendar_datatype, "Calendar DataType", ['data', 'event', 'holiday'])
+        self.check_valid_value(self.symbols_type, "Symbols Type", ['GOODS', 'FOREX', 'FUTURE', 'CRYPTO'])
+        self.check_valid_value(self.symbols_datatype, 'Symbols DataType', ['symbols', 'quotes'])
+
+
+class Jin10(ComponentBase, ABC):
+    component_name = "Jin10"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = " - ".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return Jin10.be_output("")
+
+        jin10_res = []
+        headers = {'secret-key': self._param.secret_key}
+        try:
+            if self._param.type == "flash":
+                params = {
+                    'category': self._param.flash_type,
+                    'contain': self._param.contain,
+                    'filter': self._param.filter
+                }
+                response = requests.get(
+                    url='https://open-data-api.jin10.com/data-api/flash?category=' + self._param.flash_type,
+                    headers=headers, data=json.dumps(params))
+                response = response.json()
+                for i in response['data']:
+                    jin10_res.append({"content": i['data']['content']})
+            if self._param.type == "calendar":
+                params = {
+                    'category': self._param.calendar_type
+                }
+                response = requests.get(
+                    url='https://open-data-api.jin10.com/data-api/calendar/' + self._param.calendar_datatype + '?category=' + self._param.calendar_type,
+                    headers=headers, data=json.dumps(params))
+
+                response = response.json()
+                jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
+            if self._param.type == "symbols":
+                params = {
+                    'type': self._param.symbols_type
+                }
+                if self._param.symbols_datatype == "quotes":
+                    params['codes'] = 'BTCUSD'
+                response = requests.get(
+                    url='https://open-data-api.jin10.com/data-api/' + self._param.symbols_datatype + '?type=' + self._param.symbols_type,
+                    headers=headers, data=json.dumps(params))
+                response = response.json()
+                if self._param.symbols_datatype == "symbols":
+                    for i in response['data']:
+                        i['Commodity Code'] = i['c']
+                        i['Stock Exchange'] = i['e']
+                        i['Commodity Name'] = i['n']
+                        i['Commodity Type'] = i['t']
+                        del i['c'], i['e'], i['n'], i['t']
+                if self._param.symbols_datatype == "quotes":
+                    for i in response['data']:
+                        i['Selling Price'] = i['a']
+                        i['Buying Price'] = i['b']
+                        i['Commodity Code'] = i['c']
+                        i['Stock Exchange'] = i['e']
+                        i['Highest Price'] = i['h']
+                        i['Yesterday’s Closing Price'] = i['hc']
+                        i['Lowest Price'] = i['l']
+                        i['Opening Price'] = i['o']
+                        i['Latest Price'] = i['p']
+                        i['Market Quote Time'] = i['t']
+                        del i['a'], i['b'], i['c'], i['e'], i['h'], i['hc'], i['l'], i['o'], i['p'], i['t']
+                jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
+            if self._param.type == "news":
+                params = {
+                    'contain': self._param.contain,
+                    'filter': self._param.filter
+                }
+                response = requests.get(
+                    url='https://open-data-api.jin10.com/data-api/news',
+                    headers=headers, data=json.dumps(params))
+                response = response.json()
+                jin10_res.append({"content": pd.DataFrame(response['data']).to_markdown()})
+        except Exception as e:
+            return Jin10.be_output("**ERROR**: " + str(e))
+
+        if not jin10_res:
+            return Jin10.be_output("")
+
+        return pd.DataFrame(jin10_res)

agent/component/pubmed.py

@@ -15,6 +15,7 @@
 #
 from abc import ABC
 from Bio import Entrez
+import re
 import pandas as pd
 import xml.etree.ElementTree as ET
 from agent.settings import DEBUG
@@ -47,12 +48,15 @@ class PubMed(ComponentBase, ABC):
         try:
             Entrez.email = self._param.email
             pubmedids = Entrez.read(Entrez.esearch(db='pubmed', retmax=self._param.top_n, term=ans))['IdList']
-            pubmedcnt = ET.fromstring(
+            pubmedcnt = ET.fromstring(re.sub(r'<(/?)b>|<(/?)i>', '', Entrez.efetch(db='pubmed', id=",".join(pubmedids),
-                Entrez.efetch(db='pubmed', id=",".join(pubmedids), retmode="xml").read().decode("utf-8"))
+                                                                                   retmode="xml").read().decode(
+                "utf-8")))
             pubmed_res = [{"content": 'Title:' + child.find("MedlineCitation").find("Article").find(
                 "ArticleTitle").text + '\nUrl:<a href=" https://pubmed.ncbi.nlm.nih.gov/' + child.find(
-                "MedlineCitation").find("PMID").text + '">' + '</a>\n' + 'Abstract:' + child.find(
+                "MedlineCitation").find("PMID").text + '">' + '</a>\n' + 'Abstract:' + (
-                "MedlineCitation").find("Article").find("Abstract").find("AbstractText").text} for child in
+                                          child.find("MedlineCitation").find("Article").find("Abstract").find(
+                                              "AbstractText").text if child.find("MedlineCitation").find(
+                                              "Article").find("Abstract") else "No abstract available")} for child in
                           pubmedcnt.findall("PubmedArticle")]
         except Exception as e:
             return PubMed.be_output("**ERROR**: " + str(e))

agent/component/qweather.py

@@ -51,7 +51,7 @@ class QWeatherParam(ComponentParamBase):
                                ['zh', 'zh-hant', 'en', 'de', 'es', 'fr', 'it', 'ja', 'ko', 'ru', 'hi', 'th', 'ar', 'pt',
                                 'bn', 'ms', 'nl', 'el', 'la', 'sv', 'id', 'pl', 'tr', 'cs', 'et', 'vi', 'fil', 'fi',
                                 'he', 'is', 'nb'])
-        self.check_vaild_value(self.time_period, "Time period", ['now', '3d', '7d', '10d', '15d', '30d'])
+        self.check_valid_value(self.time_period, "Time period", ['now', '3d', '7d', '10d', '15d', '30d'])
 
 
 class QWeather(ComponentBase, ABC):

agent/component/retrieval.py

@@ -54,8 +54,8 @@ class Retrieval(ComponentBase, ABC):
         for role, cnt in history[::-1][:self._param.message_history_window_size]:
             if role != "user":continue
             query.append(cnt)
-        query = "\n".join(query)
+        # query = "\n".join(query)
-
+        query = query[0]
         kbs = KnowledgebaseService.get_by_ids(self._param.kb_ids)
         if not kbs:
             raise ValueError("Can't find knowledgebases by {}".format(self._param.kb_ids))
@@ -76,7 +76,8 @@ class Retrieval(ComponentBase, ABC):
 
         if not kbinfos["chunks"]:
             df = Retrieval.be_output("")
-            df["empty_response"] = self._param.empty_response
+            if self._param.empty_response and self._param.empty_response.strip():
+                df["empty_response"] = self._param.empty_response
             return df
 
         df = pd.DataFrame(kbinfos["chunks"])

agent/component/rewrite.py

@@ -54,7 +54,7 @@ class RewriteQuestion(Generate, ABC):
             setattr(self, "_loop", 0)
         if self._loop >= self._param.loop:
             self._loop = 0
-            raise Exception("Maximum loop time exceeds. Can't find relevant information.")
+            raise Exception("Sorry! Nothing relevant found.")
         self._loop += 1
         q = "Question: "
         for r, c in self._canvas.history[::-1]:

agent/component/switch.py

@@ -42,8 +42,6 @@ class SwitchParam(ComponentParamBase):
         self.check_empty(self.conditions, "[Switch] conditions")
         for cond in self.conditions:
             if not cond["to"]: raise ValueError(f"[Switch] 'To' can not be empty!")
-            if cond["logical_operator"] not in ["and", "or"] and len(cond["items"]) > 1:
-                raise ValueError(f"[Switch] Please set logical_operator correctly!")
 
 
 class Switch(ComponentBase, ABC):

agent/component/tushare.py

@@ -0,0 +1,72 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+import json
+from abc import ABC
+import pandas as pd
+import time
+import requests
+from agent.component.base import ComponentBase, ComponentParamBase
+
+
+class TuShareParam(ComponentParamBase):
+    """
+    Define the TuShare component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.token = "xxx"
+        self.src = "eastmoney"
+        self.start_date = "2024-01-01 09:00:00"
+        self.end_date = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+        self.keyword = ""
+
+    def check(self):
+        self.check_valid_value(self.src, "Quick News Source",
+                               ["sina", "wallstreetcn", "10jqka", "eastmoney", "yuncaijing", "fenghuang", "jinrongjie"])
+
+
+class TuShare(ComponentBase, ABC):
+    component_name = "TuShare"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = ",".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return TuShare.be_output("")
+
+        try:
+            tus_res = []
+            params = {
+                "api_name": "news",
+                "token": self._param.token,
+                "params": {"src": self._param.src, "start_date": self._param.start_date,
+                           "end_date": self._param.end_date}
+            }
+            response = requests.post(url="http://api.tushare.pro", data=json.dumps(params).encode('utf-8'))
+            response = response.json()
+            if response['code'] != 0:
+                return TuShare.be_output(response['msg'])
+            df = pd.DataFrame(response['data']['items'])
+            df.columns = response['data']['fields']
+            tus_res.append({"content": (df[df['content'].str.contains(self._param.keyword, case=False)]).to_markdown()})
+        except Exception as e:
+            return TuShare.be_output("**ERROR**: " + str(e))
+
+        if not tus_res:
+            return TuShare.be_output("")
+
+        return pd.DataFrame(tus_res)

agent/component/wencai.py

@@ -0,0 +1,74 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+from abc import ABC
+import pandas as pd
+import pywencai
+from agent.component.base import ComponentBase, ComponentParamBase
+
+
+class WenCaiParam(ComponentParamBase):
+    """
+    Define the WenCai component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.top_n = 10
+        self.query_type = "stock"
+
+    def check(self):
+        self.check_positive_integer(self.top_n, "Top N")
+        self.check_valid_value(self.query_type, "Query type",
+                               ['stock', 'zhishu', 'fund', 'hkstock', 'usstock', 'threeboard', 'conbond', 'insurance',
+                                'futures', 'lccp',
+                                'foreign_exchange'])
+
+
+class WenCai(ComponentBase, ABC):
+    component_name = "WenCai"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = ",".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return WenCai.be_output("")
+
+        try:
+            wencai_res = []
+            res = pywencai.get(query=ans, query_type=self._param.query_type, perpage=self._param.top_n)
+            if isinstance(res, pd.DataFrame):
+                wencai_res.append({"content": res.to_markdown()})
+            if isinstance(res, dict):
+                for item in res.items():
+                    if isinstance(item[1], list):
+                        wencai_res.append({"content": item[0] + "\n" + pd.DataFrame(item[1]).to_markdown()})
+                        continue
+                    if isinstance(item[1], str):
+                        wencai_res.append({"content": item[0] + "\n" + item[1]})
+                        continue
+                    if isinstance(item[1], dict):
+                        if "meta" in item[1].keys():
+                            continue
+                        wencai_res.append({"content": pd.DataFrame.from_dict(item[1], orient='index').to_markdown()})
+                        continue
+                    wencai_res.append({"content": item[0] + "\n" + str(item[1])})
+        except Exception as e:
+            return WenCai.be_output("**ERROR**: " + str(e))
+
+        if not wencai_res:
+            return WenCai.be_output("")
+
+        return pd.DataFrame(wencai_res)

agent/component/yahoofinance.py

@@ -0,0 +1,83 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+from abc import ABC
+import pandas as pd
+from agent.component.base import ComponentBase, ComponentParamBase
+import yfinance as yf
+
+
+class YahooFinanceParam(ComponentParamBase):
+    """
+    Define the YahooFinance component parameters.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.info = True
+        self.history = False
+        self.count = False
+        self.financials = False
+        self.income_stmt = False
+        self.balance_sheet = False
+        self.cash_flow_statement = False
+        self.news = True
+
+    def check(self):
+        self.check_boolean(self.info, "get all stock info")
+        self.check_boolean(self.history, "get historical market data")
+        self.check_boolean(self.count, "show share count")
+        self.check_boolean(self.financials, "show financials")
+        self.check_boolean(self.income_stmt, "income statement")
+        self.check_boolean(self.balance_sheet, "balance sheet")
+        self.check_boolean(self.cash_flow_statement, "cash flow statement")
+        self.check_boolean(self.news, "show news")
+
+
+class YahooFinance(ComponentBase, ABC):
+    component_name = "YahooFinance"
+
+    def _run(self, history, **kwargs):
+        ans = self.get_input()
+        ans = "".join(ans["content"]) if "content" in ans else ""
+        if not ans:
+            return YahooFinance.be_output("")
+
+        yohoo_res = []
+        try:
+            msft = yf.Ticker(ans)
+            if self._param.info:
+                yohoo_res.append({"content": "info:\n" + pd.Series(msft.info).to_markdown() + "\n"})
+            if self._param.history:
+                yohoo_res.append({"content": "history:\n" + msft.history().to_markdown() + "\n"})
+            if self._param.financials:
+                yohoo_res.append({"content": "calendar:\n" + pd.DataFrame(msft.calendar).to_markdown() + "\n"})
+            if self._param.balance_sheet:
+                yohoo_res.append({"content": "balance sheet:\n" + msft.balance_sheet.to_markdown() + "\n"})
+                yohoo_res.append(
+                    {"content": "quarterly balance sheet:\n" + msft.quarterly_balance_sheet.to_markdown() + "\n"})
+            if self._param.cash_flow_statement:
+                yohoo_res.append({"content": "cash flow statement:\n" + msft.cashflow.to_markdown() + "\n"})
+                yohoo_res.append(
+                    {"content": "quarterly cash flow statement:\n" + msft.quarterly_cashflow.to_markdown() + "\n"})
+            if self._param.news:
+                yohoo_res.append({"content": "news:\n" + pd.DataFrame(msft.news).to_markdown() + "\n"})
+        except Exception as e:
+            print("**ERROR** " + str(e))
+
+        if not yohoo_res:
+            return YahooFinance.be_output("")
+
+        return pd.DataFrame(yohoo_res)

agent/templates/DB Assistant.json

@@ -1,7 +1,7 @@
 {
   "id": 6,
   "title": "DB Assistant",
-  "description": "Database query assistant. It converts questions into SQL statements and queries them in the database. You need to provide 3 kinds of knowledge base: 1. DDL data in the database. 2. Sample text questions converted to SQL statements. 3. A description of the database contents, including but not limited to: tables, records, and so on. You will also need to set up database configuration information: like IP Port ...",
+  "description": "An advanced agent that converts user queries into SQL statements, executes the queries, and assesses and returns the results. You must prepare three knowledge bases: 1: DDL for your database; 2: Examples of user queries converted to SQL statements; 3: A comprehensive description of your database, including but not limited to tables and records. You are also required to configure the corresponding database.",
   "canvas_type": "chatbot",
   "dsl": {
             "answer": [],
@@ -63,7 +63,7 @@
                                 }
                             ],
                             "presence_penalty": 0.4,
-                            "prompt": "## You are the Repair SQL Statement Helper, please modify the original SQL statement based on the SQL query error report.\n\n## The contents of the SQL query error report and the original SQL statement are as follows:\n{exesql_input}\n\n## Answer only the modified SQL statement. Please do not give any explanation, just answer the code.",
+                            "prompt": "## You are the Repair SQL Statement Helper, please modify the original SQL statement based on the SQL query error report.\n\n## The contents of the SQL query error report and the original SQL statement are as follows:\n{exesql_input}\n\n## Answer only the modified SQL statement. Each SQL statement ends with semicolon and do not give any explanation, just answer the code.",
                             "temperature": 0.1,
                             "top_p": 0.3
                         }
@@ -102,7 +102,7 @@
                                 }
                             ],
                             "presence_penalty": 0.4,
-                            "prompt": "\n##The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.\n\n##You may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you: {ddl_input}.\n\n##You may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you: {db_input}.\n\n##You may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you: {sql_input}.\n\n##Respond with only SQL code. Do not answer with any explanations -- just the code.",
+                            "prompt": "\n##The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.\n\n##You may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you: {ddl_input}.\n\n##You may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you: {db_input}.\n\n##You may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you: {sql_input}.\n\n##Respond with only SQL code. Each SQL code ends with semicolon and do not give any explanation -- just the code.",
                             "temperature": 0.1,
                             "top_p": 0.3
                         }
@@ -120,6 +120,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in DB-Description!",
                             "kb_ids": [
                                 "b510f8f45f6011ef904f0242ac160006"
                             ],
@@ -139,6 +140,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in DDL!",
                             "kb_ids": [
                                 "9870268e5f6011efb8570242ac160006"
                             ],
@@ -158,6 +160,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in Q->SQL!",
                             "kb_ids": [
                                 "dd401bcc5b9e11efae770242ac160006"
                             ],
@@ -408,6 +411,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in Q->SQL!",
                                 "kb_ids": [
                                     "dd401bcc5b9e11efae770242ac160006"
                                 ],
@@ -493,6 +497,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in DB-Description!",
                                 "kb_ids": [
                                     "b510f8f45f6011ef904f0242ac160006"
                                 ],
@@ -523,6 +528,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in DDL!",
                                 "kb_ids": [
                                     "9870268e5f6011efb8570242ac160006"
                                 ],

agent/templates/HR_callout_zh.json

@@ -1,7 +1,7 @@
 {
   "id": 2,
-  "title": "HR call-out assistant(Chinese)",
+  "title": "HR recruitment pitch assistant (Chinese)",
-  "description": "A HR call-out assistant. It will introduce the given job, answer the candidates' question about this job. And the most important thing is that it will try to obtain the contact information of the candidates. What you need to do is to link a knowledgebase which contains job description in 'Retrieval' component.",
+  "description": "A recruitment pitch assistant capable of pitching a candidate, presenting a job opportunity, addressing queries, and requesting the candidate's contact details. Let's begin by linking a knowledge base containing the job description in 'Retrieval'!",
   "canvas_type": "chatbot",
   "dsl": {
     "answer": [],

agent/templates/customer_service.json

@@ -1,7 +1,7 @@
 {
   "id": 3,
   "title": "Customer service",
-  "description": "A call-in customer service chat bot. It will provide useful information about the products, answer customers' questions and soothe the customers' bad emotions.",
+  "description": "A customer service chatbot that explains product specifications, addresses customer queries, and alleviates negative emotions.",
   "canvas_type": "chatbot",
   "dsl": {
     "answer": [],

agent/templates/general_chat_bot.json

@@ -1,7 +1,7 @@
 {
   "id": 1,
-  "title": "Chat bot",
+  "title": "General-purpose chatbot",
-  "description": "A general chat bot. It is based on Self-RAG mechanism. What you need to do is setting up knowleage base in 'Retrieval'",
+  "description": "A general-purpose chatbot based on Self-RAG. Let's begin by setting up your knowledge base in 'Retrieval'!",
   "canvas_type": "chatbot",
   "dsl": {
     "answer": [],

agent/templates/interpreter.json

@@ -1,7 +1,7 @@
 {
   "id": 4,
   "title": "Interpreter",
-  "description": "An interpreter. Type the content you want to translate and the object language like: Hi there => Spanish. Hava a try!",
+  "description": "A simple interpreter that translates user input into a target language. Try 'Hi there => Spanish' to see the translation!",
   "canvas_type": "chatbot",
   "dsl": {
     "answer": [],

agent/templates/text2sql.json

@@ -1,7 +1,7 @@
 {
   "id": 5,
   "title": "Text To SQL",
-  "description": "An agent tool provides the ability to convert text questions into SQL statements, you will need to provide 3 kinds of knowledge bases. 1. DDL data for the database. 2. Examples of text questions converted into SQL statements. 3. A description of the database contents, including but not limited to: tables, records, etc...  ",
+  "description": "An agent that converts user queries into SQL statements. You must prepare three knowledge bases: 1: DDL for your database; 2: Examples of user queries converted to SQL statements; 3: A comprehensive description of your database, including but not limited to tables and records.",
   "canvas_type": "chatbot",
   "dsl": {
             "answer": [],
@@ -69,6 +69,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in DB-Description!",
                             "kb_ids": [
                                 "0ab5de985ba911efad9942010a8a0006"
                             ],
@@ -88,6 +89,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in DDL!",
                             "kb_ids": [
                                 "b1a6a45e5ba811ef80dc42010a8a0006"
                             ],
@@ -107,6 +109,7 @@
                     "obj": {
                         "component_name": "Retrieval",
                         "params": {
+                            "empty_response": "Nothing found in Q-SQL!",
                             "kb_ids": [
                                 "31257b925b9f11ef9f0142010a8a0004"
                             ],
@@ -286,6 +289,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in Q-SQL!",
                                 "kb_ids": [
                                     "31257b925b9f11ef9f0142010a8a0004"
                                 ],
@@ -371,6 +375,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in DB-Description!",
                                 "kb_ids": [
                                     "0ab5de985ba911efad9942010a8a0006"
                                 ],
@@ -401,6 +406,7 @@
                     {
                         "data": {
                             "form": {
+                                "empty_response": "Nothing found in DDL!",
                                 "kb_ids": [
                                     "b1a6a45e5ba811ef80dc42010a8a0006"
                                 ],

agent/templates/websearch_assistant.json

@@ -1,7 +1,7 @@
 {
   "id": 0,
   "title": "WebSearch Assistant",
-  "description": "A chat assistant that combines information both from knowledge base and web search engines. It integrates information from the knowledge base and relevant search engines to answer a given question. What you need to do is setting up knowleage base in 'Retrieval'.",
+  "description": "A chat assistant template that integrates information extracted from a knowledge base and web searches to respond to queries. Let's begin by setting up your knowledge base in 'Retrieval'!",
   "canvas_type": "chatbot",
   "dsl": {
     "answer": [],

api/apps/__init__.py

@@ -65,7 +65,7 @@ commands.register_commands(app)
 
 def search_pages_path(pages_dir):
     app_path_list = [path for path in pages_dir.glob('*_app.py') if not path.name.startswith('.')]
-    api_path_list = [path for path in pages_dir.glob('*_api.py') if not path.name.startswith('.')]
+    api_path_list = [path for path in pages_dir.glob('*sdk/*.py') if not path.name.startswith('.')]
     app_path_list.extend(api_path_list)
     return app_path_list
 
@@ -73,7 +73,7 @@ def search_pages_path(pages_dir):
 def register_page(page_path):
     path = f'{page_path}'
 
-    page_name = page_path.stem.rstrip('_api') if "_api" in path else page_path.stem.rstrip('_app')
+    page_name = page_path.stem.rstrip('_app')
     module_name = '.'.join(page_path.parts[page_path.parts.index('api'):-1] + (page_name,))
 
     spec = spec_from_file_location(module_name, page_path)
@@ -83,7 +83,7 @@ def register_page(page_path):
     sys.modules[module_name] = page
     spec.loader.exec_module(page)
     page_name = getattr(page, 'page_name', page_name)
-    url_prefix = f'/api/{API_VERSION}/{page_name}' if "_api" in path else f'/{API_VERSION}/{page_name}'
+    url_prefix = f'/api/{API_VERSION}/{page_name}' if "/sdk/" in path else f'/{API_VERSION}/{page_name}'
 
     app.register_blueprint(page.manager, url_prefix=url_prefix)
     return url_prefix
@@ -91,7 +91,8 @@ def register_page(page_path):
 
 pages_dir = [
     Path(__file__).parent,
-    Path(__file__).parent.parent / 'api' / 'apps', # FIXME: ragflow/api/api/apps, can be remove?
+    Path(__file__).parent.parent / 'api' / 'apps',
+    Path(__file__).parent.parent / 'api' / 'apps' / 'sdk',
 ]
 
 client_urls_prefix = [

api/apps/api_app.py

@@ -39,7 +39,7 @@ from itsdangerous import URLSafeTimedSerializer
 
 from api.utils.file_utils import filename_type, thumbnail
 from rag.nlp import keyword_extraction
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 
 from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
 from agent.canvas import Canvas
@@ -151,14 +151,17 @@ def set_conversation():
     req = request.json
     try:
         if objs[0].source == "agent":
-            e, c = UserCanvasService.get_by_id(objs[0].dialog_id)
+            e, cvs = UserCanvasService.get_by_id(objs[0].dialog_id)
             if not e:
                 return server_error_response("canvas not found.")
+            if not isinstance(cvs.dsl, str):
+                cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
+            canvas = Canvas(cvs.dsl, objs[0].tenant_id)
             conv = {
                 "id": get_uuid(),
-                "dialog_id": c.id,
+                "dialog_id": cvs.id,
                 "user_id": request.args.get("user_id", ""),
-                "message": [{"role": "assistant", "content": "Hi there!"}],
+                "message": [{"role": "assistant", "content": canvas.get_prologue()}],
                 "source": "agent"
             }
             API4ConversationService.save(**conv)
@@ -199,15 +202,18 @@ def completion():
             continue
         if m["role"] == "assistant" and not msg:
             continue
-        msg.append({"role": m["role"], "content": m["content"]})
+        msg.append(m)
+    if not msg[-1].get("id"): msg[-1]["id"] = get_uuid()
+    message_id = msg[-1]["id"]
 
     def fillin_conv(ans):
-        nonlocal conv
+        nonlocal conv, message_id
         if not conv.reference:
             conv.reference.append(ans["reference"])
         else:
             conv.reference[-1] = ans["reference"]
-        conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
+        conv.message[-1] = {"role": "assistant", "content": ans["answer"], "id": message_id}
+        ans["id"] = message_id
 
     def rename_field(ans):
         reference = ans['reference']
@@ -233,7 +239,7 @@ def completion():
 
             if not conv.reference:
                 conv.reference = []
-            conv.message.append({"role": "assistant", "content": ""})
+            conv.message.append({"role": "assistant", "content": "", "id": message_id})
             conv.reference.append({"chunks": [], "doc_aggs": []})
 
             final_ans = {"reference": [], "content": ""}
@@ -260,7 +266,7 @@ def completion():
                             yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
                                                        ensure_ascii=False) + "\n\n"
 
-                        canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
+                        canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
                         if final_ans.get("reference"):
                             canvas.reference.append(final_ans["reference"])
                         cvs.dsl = json.loads(str(canvas))
@@ -279,7 +285,7 @@ def completion():
                 return resp
 
             final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
-            canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
+            canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
             if final_ans.get("reference"):
                 canvas.reference.append(final_ans["reference"])
             cvs.dsl = json.loads(str(canvas))
@@ -300,7 +306,7 @@ def completion():
 
         if not conv.reference:
             conv.reference = []
-        conv.message.append({"role": "assistant", "content": ""})
+        conv.message.append({"role": "assistant", "content": "", "id": message_id})
         conv.reference.append({"chunks": [], "doc_aggs": []})
 
         def stream():
@@ -342,12 +348,22 @@ def completion():
 @manager.route('/conversation/<conversation_id>', methods=['GET'])
 # @login_required
 def get(conversation_id):
+    token = request.headers.get('Authorization').split()[1]
+    objs = APIToken.query(token=token)
+    if not objs:
+        return get_json_result(
+            data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
+    
     try:
         e, conv = API4ConversationService.get_by_id(conversation_id)
         if not e:
             return get_data_error_result(retmsg="Conversation not found!")
 
         conv = conv.to_dict()
+        if token != APIToken.query(dialog_id=conv['dialog_id'])[0].token:
+            return get_json_result(data=False, retmsg='Token is not valid for this conversation_id!"',
+                                   retcode=RetCode.AUTHENTICATION_ERROR)
+            
         for referenct_i in conv['reference']:
             if referenct_i is None or len(referenct_i) == 0:
                 continue
@@ -411,10 +427,10 @@ def upload():
                 retmsg="This type of file has not been supported yet!")
 
         location = filename
-        while MINIO.obj_exist(kb_id, location):
+        while STORAGE_IMPL.obj_exist(kb_id, location):
             location += "_"
         blob = request.files['file'].read()
-        MINIO.put(kb_id, location, blob)
+        STORAGE_IMPL.put(kb_id, location, blob)
         doc = {
             "id": get_uuid(),
             "kb_id": kb.id,
@@ -634,7 +650,7 @@ def document_rm():
             FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
             File2DocumentService.delete_by_document_id(doc_id)
 
-            MINIO.rm(b, n)
+            STORAGE_IMPL.rm(b, n)
         except Exception as e:
             errors += str(e)
 
@@ -707,7 +723,7 @@ def completion_faq():
             if ans["reference"]["chunks"][chunk_idx]["img_id"]:
                 try:
                     bkt, nm = ans["reference"]["chunks"][chunk_idx]["img_id"].split("-")
-                    response = MINIO.get(bkt, nm)
+                    response = STORAGE_IMPL.get(bkt, nm)
                     data_type_picture["url"] = base64.b64encode(response).decode('utf-8')
                     data.append(data_type_picture)
                     break

api/apps/canvas_app.py

@@ -18,8 +18,9 @@ from functools import partial
 from flask import request, Response
 from flask_login import login_required, current_user
 from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
+from api.settings import RetCode
 from api.utils import get_uuid
-from api.utils.api_utils import get_json_result, server_error_response, validate_request
+from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result
 from agent.canvas import Canvas
 from peewee import MySQLDatabase, PostgresqlDatabase
 
@@ -43,6 +44,10 @@ def canvas_list():
 @login_required
 def rm():
     for i in request.json["canvas_ids"]:
+        if not UserCanvasService.query(user_id=current_user.id,id=i):
+            return get_json_result(
+                data=False, retmsg=f'Only owner of canvas authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
         UserCanvasService.delete_by_id(i)
     return get_json_result(data=True)
 
@@ -61,10 +66,13 @@ def save():
             return server_error_response(ValueError("Duplicated title."))
         req["id"] = get_uuid()
         if not UserCanvasService.save(**req):
-            return server_error_response("Fail to save canvas.")
+            return get_data_error_result(retmsg="Fail to save canvas.")
     else:
+        if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
+            return get_json_result(
+                data=False, retmsg=f'Only owner of canvas authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
         UserCanvasService.update_by_id(req["id"], req)
-
     return get_json_result(data=req)
 
 
@@ -73,7 +81,7 @@ def save():
 def get(canvas_id):
     e, c = UserCanvasService.get_by_id(canvas_id)
     if not e:
-        return server_error_response("canvas not found.")
+        return get_data_error_result(retmsg="canvas not found.")
     return get_json_result(data=c.to_dict())
 
 
@@ -85,16 +93,21 @@ def run():
     stream = req.get("stream", True)
     e, cvs = UserCanvasService.get_by_id(req["id"])
     if not e:
-        return server_error_response("canvas not found.")
+        return get_data_error_result(retmsg="canvas not found.")
+    if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
+        return get_json_result(
+            data=False, retmsg=f'Only owner of canvas authorized for this operation.',
+            retcode=RetCode.OPERATING_ERROR)
 
     if not isinstance(cvs.dsl, str):
         cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
 
     final_ans = {"reference": [], "content": ""}
+    message_id = req.get("message_id", get_uuid())
     try:
         canvas = Canvas(cvs.dsl, current_user.id)
         if "message" in req:
-            canvas.messages.append({"role": "user", "content": req["message"]})
+            canvas.messages.append({"role": "user", "content": req["message"], "id": message_id})
             canvas.add_user_input(req["message"])
         answer = canvas.run(stream=stream)
         print(canvas)
@@ -115,7 +128,7 @@ def run():
                     ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
                     yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
 
-                canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
+                canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
                 if final_ans.get("reference"):
                     canvas.reference.append(final_ans["reference"])
                 cvs.dsl = json.loads(str(canvas))
@@ -134,7 +147,7 @@ def run():
         return resp
 
     final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
-    canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
+    canvas.messages.append({"role": "assistant", "content": final_ans["content"], "id": message_id})
     if final_ans.get("reference"):
         canvas.reference.append(final_ans["reference"])
     cvs.dsl = json.loads(str(canvas))
@@ -150,7 +163,11 @@ def reset():
     try:
         e, user_canvas = UserCanvasService.get_by_id(req["id"])
         if not e:
-            return server_error_response("canvas not found.")
+            return get_data_error_result(retmsg="canvas not found.")
+        if not UserCanvasService.query(user_id=current_user.id, id=req["id"]):
+            return get_json_result(
+                data=False, retmsg=f'Only owner of canvas authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
 
         canvas = Canvas(json.dumps(user_canvas.dsl), current_user.id)
         canvas.reset()

api/apps/chunk_app.py

@@ -58,7 +58,7 @@ def list_chunk():
         }
         if "available_int" in req:
             query["available_int"] = int(req["available_int"])
-        sres = retrievaler.search(query, search.index_name(tenant_id))
+        sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
         res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
         for id in sres.ids:
             d = {
@@ -259,12 +259,25 @@ def retrieval_test():
     size = int(req.get("size", 30))
     question = req["question"]
     kb_id = req["kb_id"]
+    if isinstance(kb_id, str): kb_id = [kb_id]
     doc_ids = req.get("doc_ids", [])
-    similarity_threshold = float(req.get("similarity_threshold", 0.2))
+    similarity_threshold = float(req.get("similarity_threshold", 0.0))
     vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
     top = int(req.get("top_k", 1024))
+
     try:
-        e, kb = KnowledgebaseService.get_by_id(kb_id)
+        tenants = UserTenantService.query(user_id=current_user.id)
+        for kid in kb_id:
+            for tenant in tenants:
+                if KnowledgebaseService.query(
+                        tenant_id=tenant.tenant_id, id=kid):
+                    break
+            else:
+                return get_json_result(
+                    data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+                    retcode=RetCode.OPERATING_ERROR)
+
+        e, kb = KnowledgebaseService.get_by_id(kb_id[0])
         if not e:
             return get_data_error_result(retmsg="Knowledgebase not found!")
 
@@ -281,9 +294,9 @@ def retrieval_test():
             question += keyword_extraction(chat_mdl, question)
 
         retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
-        ranks = retr.retrieval(question, embd_mdl, kb.tenant_id, [kb_id], page, size,
+        ranks = retr.retrieval(question, embd_mdl, kb.tenant_id, kb_id, page, size,
                                similarity_threshold, vector_similarity_weight, top,
-                               doc_ids, rerank_mdl=rerank_mdl)
+                               doc_ids, rerank_mdl=rerank_mdl, highlight=req.get("highlight"))
         for c in ranks["chunks"]:
             if "vector" in c:
                 del c["vector"]

api/apps/conversation_app.py

@@ -13,14 +13,23 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import json
+import re
+import traceback
 from copy import deepcopy
+from api.db.services.user_service import UserTenantService
 from flask import request, Response
-from flask_login import login_required
+from flask_login import login_required, current_user
-from api.db.services.dialog_service import DialogService, ConversationService, chat
+
-from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
+from api.db import LLMType
+from api.db.services.dialog_service import DialogService, ConversationService, chat, ask
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.llm_service import LLMBundle, TenantService, TenantLLMService
+from api.settings import RetCode, retrievaler
 from api.utils import get_uuid
 from api.utils.api_utils import get_json_result
-import json
+from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
+from graphrag.mind_map_extractor import MindMapExtractor
 
 
 @manager.route('/set', methods=['POST'])
@@ -70,6 +79,14 @@ def get():
         e, conv = ConversationService.get_by_id(conv_id)
         if not e:
             return get_data_error_result(retmsg="Conversation not found!")
+        tenants = UserTenantService.query(user_id=current_user.id)
+        for tenant in tenants:
+            if DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id):
+                break
+        else:
+            return get_json_result(
+                data=False, retmsg=f'Only owner of conversation authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
         conv = conv.to_dict()
         return get_json_result(data=conv)
     except Exception as e:
@@ -82,6 +99,17 @@ def rm():
     conv_ids = request.json["conversation_ids"]
     try:
         for cid in conv_ids:
+            exist, conv = ConversationService.get_by_id(cid)
+            if not exist:
+                return get_data_error_result(retmsg="Conversation not found!")
+            tenants = UserTenantService.query(user_id=current_user.id)
+            for tenant in tenants:
+                if DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id):
+                    break
+            else:
+                return get_json_result(
+                    data=False, retmsg=f'Only owner of conversation authorized for this operation.',
+                    retcode=RetCode.OPERATING_ERROR)
             ConversationService.delete_by_id(cid)
         return get_json_result(data=True)
     except Exception as e:
@@ -93,6 +121,10 @@ def rm():
 def list_convsersation():
     dialog_id = request.args["dialog_id"]
     try:
+        if not DialogService.query(tenant_id=current_user.id, id=dialog_id):
+            return get_json_result(
+                data=False, retmsg=f'Only owner of dialog authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
         convs = ConversationService.query(
             dialog_id=dialog_id,
             order_by=ConversationService.model.create_time,
@@ -105,26 +137,25 @@ def list_convsersation():
 
 @manager.route('/completion', methods=['POST'])
 @login_required
-#@validate_request("conversation_id", "messages")
+@validate_request("conversation_id", "messages")
 def completion():
     req = request.json
-    #req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
+    # req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
     #    {"role": "user", "content": "上海有吗？"}
-    #]}
+    # ]}
     msg = []
     for m in req["messages"]:
         if m["role"] == "system":
             continue
         if m["role"] == "assistant" and not msg:
             continue
-        msg.append({"role": m["role"], "content": m["content"]})
+        msg.append(m)
-        if "doc_ids" in m:
+    message_id = msg[-1].get("id")
-            msg[-1]["doc_ids"] = m["doc_ids"]
     try:
         e, conv = ConversationService.get_by_id(req["conversation_id"])
         if not e:
             return get_data_error_result(retmsg="Conversation not found!")
-        conv.message.append(deepcopy(msg[-1]))
+        conv.message = deepcopy(req["messages"])
         e, dia = DialogService.get_by_id(conv.dialog_id)
         if not e:
             return get_data_error_result(retmsg="Dialog not found!")
@@ -133,28 +164,31 @@ def completion():
 
         if not conv.reference:
             conv.reference = []
-        conv.message.append({"role": "assistant", "content": ""})
+        conv.message.append({"role": "assistant", "content": "", "id": message_id})
         conv.reference.append({"chunks": [], "doc_aggs": []})
 
         def fillin_conv(ans):
-            nonlocal conv
+            nonlocal conv, message_id
             if not conv.reference:
                 conv.reference.append(ans["reference"])
-            else: conv.reference[-1] = ans["reference"]
+            else:
-            conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
+                conv.reference[-1] = ans["reference"]
+            conv.message[-1] = {"role": "assistant", "content": ans["answer"],
+                                "id": message_id, "prompt": ans.get("prompt", "")}
+            ans["id"] = message_id
 
         def stream():
             nonlocal dia, msg, req, conv
             try:
                 for ans in chat(dia, msg, True, **req):
                     fillin_conv(ans)
-                    yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+                    yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
                 ConversationService.update_by_id(conv.id, conv.to_dict())
             except Exception as e:
                 yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
-                                            "data": {"answer": "**ERROR**: "+str(e), "reference": []}},
+                                            "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
                                            ensure_ascii=False) + "\n\n"
-            yield "data:"+json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
+            yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
 
         if req.get("stream", True):
             resp = Response(stream(), mimetype="text/event-stream")
@@ -175,3 +209,168 @@ def completion():
     except Exception as e:
         return server_error_response(e)
 
+
+@manager.route('/tts', methods=['POST'])
+@login_required
+def tts():
+    req = request.json
+    text = req["text"]
+
+    tenants = TenantService.get_by_user_id(current_user.id)
+    if not tenants:
+        return get_data_error_result(retmsg="Tenant not found!")
+
+    tts_id = tenants[0]["tts_id"]
+    if not tts_id:
+        return get_data_error_result(retmsg="No default TTS model is set")
+
+    tts_mdl = LLMBundle(tenants[0]["tenant_id"], LLMType.TTS, tts_id)
+
+    def stream_audio():
+        try:
+            for chunk in tts_mdl.tts(text):
+                yield chunk
+        except Exception as e:
+            yield ("data:" + json.dumps({"retcode": 500, "retmsg": str(e),
+                                         "data": {"answer": "**ERROR**: " + str(e)}},
+                                        ensure_ascii=False)).encode('utf-8')
+
+    resp = Response(stream_audio(), mimetype="audio/mpeg")
+    resp.headers.add_header("Cache-Control", "no-cache")
+    resp.headers.add_header("Connection", "keep-alive")
+    resp.headers.add_header("X-Accel-Buffering", "no")
+
+    return resp
+
+
+@manager.route('/delete_msg', methods=['POST'])
+@login_required
+@validate_request("conversation_id", "message_id")
+def delete_msg():
+    req = request.json
+    e, conv = ConversationService.get_by_id(req["conversation_id"])
+    if not e:
+        return get_data_error_result(retmsg="Conversation not found!")
+
+    conv = conv.to_dict()
+    for i, msg in enumerate(conv["message"]):
+        if req["message_id"] != msg.get("id", ""):
+            continue
+        assert conv["message"][i + 1]["id"] == req["message_id"]
+        conv["message"].pop(i)
+        conv["message"].pop(i)
+        conv["reference"].pop(max(0, i // 2 - 1))
+        break
+
+    ConversationService.update_by_id(conv["id"], conv)
+    return get_json_result(data=conv)
+
+
+@manager.route('/thumbup', methods=['POST'])
+@login_required
+@validate_request("conversation_id", "message_id")
+def thumbup():
+    req = request.json
+    e, conv = ConversationService.get_by_id(req["conversation_id"])
+    if not e:
+        return get_data_error_result(retmsg="Conversation not found!")
+    up_down = req.get("set")
+    feedback = req.get("feedback", "")
+    conv = conv.to_dict()
+    for i, msg in enumerate(conv["message"]):
+        if req["message_id"] == msg.get("id", "") and msg.get("role", "") == "assistant":
+            if up_down:
+                msg["thumbup"] = True
+                if "feedback" in msg: del msg["feedback"]
+            else:
+                msg["thumbup"] = False
+                if feedback: msg["feedback"] = feedback
+            break
+
+    ConversationService.update_by_id(conv["id"], conv)
+    return get_json_result(data=conv)
+
+
+@manager.route('/ask', methods=['POST'])
+@login_required
+@validate_request("question", "kb_ids")
+def ask_about():
+    req = request.json
+    uid = current_user.id
+    def stream():
+        nonlocal req, uid
+        try:
+            for ans in ask(req["question"], req["kb_ids"], uid):
+                yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+        except Exception as e:
+            yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
+                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                       ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+    resp = Response(stream(), mimetype="text/event-stream")
+    resp.headers.add_header("Cache-control", "no-cache")
+    resp.headers.add_header("Connection", "keep-alive")
+    resp.headers.add_header("X-Accel-Buffering", "no")
+    resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
+    return resp
+
+
+@manager.route('/mindmap', methods=['POST'])
+@login_required
+@validate_request("question", "kb_ids")
+def mindmap():
+    req = request.json
+    kb_ids = req["kb_ids"]
+    e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
+    if not e:
+        return get_data_error_result(retmsg="Knowledgebase not found!")
+
+    embd_mdl = TenantLLMService.model_instance(
+        kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
+    chat_mdl = LLMBundle(current_user.id, LLMType.CHAT)
+    ranks = retrievaler.retrieval(req["question"], embd_mdl, kb.tenant_id, kb_ids, 1, 12,
+                           0.3, 0.3, aggs=False)
+    mindmap = MindMapExtractor(chat_mdl)
+    mind_map = mindmap([c["content_with_weight"] for c in ranks["chunks"]]).output
+    if "error" in mind_map:
+        return server_error_response(Exception(mind_map["error"]))
+    return get_json_result(data=mind_map)
+
+
+@manager.route('/related_questions', methods=['POST'])
+@login_required
+@validate_request("question")
+def related_questions():
+    req = request.json
+    question = req["question"]
+    chat_mdl = LLMBundle(current_user.id, LLMType.CHAT)
+    prompt = """
+Objective: To generate search terms related to the user's search keywords, helping users find more valuable information.
+Instructions:
+ - Based on the keywords provided by the user, generate 5-10 related search terms.
+ - Each search term should be directly or indirectly related to the keyword, guiding the user to find more valuable information.
+ - Use common, general terms as much as possible, avoiding obscure words or technical jargon.
+ - Keep the term length between 2-4 words, concise and clear.
+ - DO NOT translate, use the language of the original keywords.
+
+### Example:
+Keywords: Chinese football
+Related search terms:
+1. Current status of Chinese football
+2. Reform of Chinese football
+3. Youth training of Chinese football
+4. Chinese football in the Asian Cup
+5. Chinese football in the World Cup
+
+Reason:
+ - When searching, users often only use one or two keywords, making it difficult to fully express their information needs.
+ - Generating related search terms can help users dig deeper into relevant information and improve search efficiency. 
+ - At the same time, related terms can also help search engines better understand user needs and return more accurate search results.
+ 
+"""
+    ans = chat_mdl.chat(prompt, [{"role": "user", "content": f"""
+Keywords: {question}
+Related search terms:
+    """}], {"temperature": 0.9})
+    return get_json_result(data=[re.sub(r"^[0-9]\. ", "", a) for a in ans.split("\n") if re.match(r"^[0-9]\. ", a)])

api/apps/dataset_api.py

@@ -42,7 +42,7 @@ from api.utils.file_utils import filename_type, thumbnail
 from rag.app import book, laws, manual, naive, one, paper, presentation, qa, resume, table, picture, audio, email
 from rag.nlp import search
 from rag.utils.es_conn import ELASTICSEARCH
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 
 MAXIMUM_OF_UPLOADING_FILES = 256
 
@@ -352,7 +352,7 @@ def upload_documents(dataset_id):
 
             # upload to the minio
             location = filename
-            while MINIO.obj_exist(dataset_id, location):
+            while STORAGE_IMPL.obj_exist(dataset_id, location):
                 location += "_"
 
             blob = file.read()
@@ -361,7 +361,7 @@ def upload_documents(dataset_id):
             if blob == b'':
                 warnings.warn(f"[WARNING]: The content of the file {filename} is empty.")
 
-            MINIO.put(dataset_id, location, blob)
+            STORAGE_IMPL.put(dataset_id, location, blob)
 
             doc = {
                 "id": get_uuid(),
@@ -441,7 +441,7 @@ def delete_document(document_id, dataset_id):  # string
         File2DocumentService.delete_by_document_id(document_id)
 
         # delete it from minio
-        MINIO.rm(dataset_id, location)
+        STORAGE_IMPL.rm(dataset_id, location)
     except Exception as e:
         errors += str(e)
     if errors:
@@ -596,7 +596,7 @@ def download_document(dataset_id, document_id):
 
         # The process of downloading
         doc_id, doc_location = File2DocumentService.get_minio_address(doc_id=document_id)  # minio address
-        file_stream = MINIO.get(doc_id, doc_location)
+        file_stream = STORAGE_IMPL.get(doc_id, doc_location)
         if not file_stream:
             return construct_json_result(message="This file is empty.", code=RetCode.DATA_ERROR)
 
@@ -737,7 +737,7 @@ def parsing_document_internal(id):
         doc_id = doc_attributes["id"]
 
         bucket, doc_name = File2DocumentService.get_minio_address(doc_id=doc_id)
-        binary = MINIO.get(bucket, doc_name)
+        binary = STORAGE_IMPL.get(bucket, doc_name)
         parser_name = doc_attributes["parser_id"]
         if binary:
             res = doc_parse(binary, doc_name, parser_name, tenant_id, doc_id)

api/apps/dialog_app.py

@@ -19,7 +19,8 @@ from flask_login import login_required, current_user
 from api.db.services.dialog_service import DialogService
 from api.db import StatusEnum
 from api.db.services.knowledgebase_service import KnowledgebaseService
-from api.db.services.user_service import TenantService
+from api.db.services.user_service import TenantService, UserTenantService
+from api.settings import RetCode
 from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 from api.utils import get_uuid
 from api.utils.api_utils import get_json_result
@@ -164,9 +165,19 @@ def list_dialogs():
 @validate_request("dialog_ids")
 def rm():
     req = request.json
+    dialog_list=[]
+    tenants = UserTenantService.query(user_id=current_user.id)
     try:
-        DialogService.update_many_by_id(
+        for id in req["dialog_ids"]:
-            [{"id": id, "status": StatusEnum.INVALID.value} for id in req["dialog_ids"]])
+            for tenant in tenants:
+                if DialogService.query(tenant_id=tenant.tenant_id, id=id):
+                    break
+            else:
+                return get_json_result(
+                    data=False, retmsg=f'Only owner of dialog authorized for this operation.',
+                    retcode=RetCode.OPERATING_ERROR)
+            dialog_list.append({"id": id,"status":StatusEnum.INVALID.value})
+        DialogService.update_many_by_id(dialog_list)
         return get_json_result(data=True)
     except Exception as e:
         return server_error_response(e)

api/apps/document_app.py

@@ -35,7 +35,7 @@ from api.db.services.file2document_service import File2DocumentService
 from api.db.services.file_service import FileService
 from api.db.services.llm_service import LLMBundle
 from api.db.services.task_service import TaskService, queue_tasks
-from api.db.services.user_service import TenantService
+from api.db.services.user_service import TenantService, UserTenantService
 from graphrag.mind_map_extractor import MindMapExtractor
 from rag.app import naive
 from rag.nlp import search
@@ -48,7 +48,7 @@ from api.db import FileType, TaskStatus, ParserType, FileSource, LLMType
 from api.db.services.document_service import DocumentService, doc_upload_and_parse
 from api.settings import RetCode, stat_logger
 from api.utils.api_utils import get_json_result
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 from api.utils.file_utils import filename_type, thumbnail, get_project_base_directory
 from api.utils.web_utils import html2pdf, is_valid_url
 
@@ -118,9 +118,9 @@ def web_crawl():
             raise RuntimeError("This type of file has not been supported yet!")
 
         location = filename
-        while MINIO.obj_exist(kb_id, location):
+        while STORAGE_IMPL.obj_exist(kb_id, location):
             location += "_"
-        MINIO.put(kb_id, location, blob)
+        STORAGE_IMPL.put(kb_id, location, blob)
         doc = {
             "id": get_uuid(),
             "kb_id": kb.id,
@@ -189,6 +189,15 @@ def list_docs():
     if not kb_id:
         return get_json_result(
             data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
+    tenants = UserTenantService.query(user_id=current_user.id)
+    for tenant in tenants:
+        if KnowledgebaseService.query(
+                tenant_id=tenant.tenant_id, id=kb_id):
+            break
+    else:
+        return get_json_result(
+            data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+            retcode=RetCode.OPERATING_ERROR)
     keywords = request.args.get("keywords", "")
 
     page_number = int(request.args.get("page", 1))
@@ -298,7 +307,7 @@ def rm():
             FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
             File2DocumentService.delete_by_document_id(doc_id)
 
-            MINIO.rm(b, n)
+            STORAGE_IMPL.rm(b, n)
         except Exception as e:
             errors += str(e)
 
@@ -385,7 +394,7 @@ def get(doc_id):
             return get_data_error_result(retmsg="Document not found!")
 
         b, n = File2DocumentService.get_minio_address(doc_id=doc_id)
-        response = flask.make_response(MINIO.get(b, n))
+        response = flask.make_response(STORAGE_IMPL.get(b, n))
 
         ext = re.search(r"\.([^.]+)$", doc.name)
         if ext:
@@ -449,7 +458,7 @@ def change_parser():
 def get_image(image_id):
     try:
         bkt, nm = image_id.split("-")
-        response = flask.make_response(MINIO.get(bkt, nm))
+        response = flask.make_response(STORAGE_IMPL.get(bkt, nm))
         response.headers.set('Content-Type', 'image/JPEG')
         return response
     except Exception as e:

api/apps/file_app.py

@@ -34,7 +34,7 @@ from api.utils.api_utils import get_json_result
 from api.utils.file_utils import filename_type
 from rag.nlp import search
 from rag.utils.es_conn import ELASTICSEARCH
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 
 
 @manager.route('/upload', methods=['POST'])
@@ -98,7 +98,7 @@ def upload():
             # file type
             filetype = filename_type(file_obj_names[file_len - 1])
             location = file_obj_names[file_len - 1]
-            while MINIO.obj_exist(last_folder.id, location):
+            while STORAGE_IMPL.obj_exist(last_folder.id, location):
                 location += "_"
             blob = file_obj.read()
             filename = duplicate_name(
@@ -116,7 +116,7 @@ def upload():
                 "size": len(blob),
             }
             file = FileService.insert(file)
-            MINIO.put(last_folder.id, location, blob)
+            STORAGE_IMPL.put(last_folder.id, location, blob)
             file_res.append(file.to_json())
         return get_json_result(data=file_res)
     except Exception as e:
@@ -260,7 +260,7 @@ def rm():
                     e, file = FileService.get_by_id(inner_file_id)
                     if not e:
                         return get_data_error_result(retmsg="File not found!")
-                    MINIO.rm(file.parent_id, file.location)
+                    STORAGE_IMPL.rm(file.parent_id, file.location)
                 FileService.delete_folder_by_pf_id(current_user.id, file_id)
             else:
                 if not FileService.delete(file):
@@ -296,7 +296,8 @@ def rename():
         e, file = FileService.get_by_id(req["file_id"])
         if not e:
             return get_data_error_result(retmsg="File not found!")
-        if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
+        if file.type != FileType.FOLDER.value \
+            and pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
                 file.name.lower()).suffix:
             return get_json_result(
                 data=False,
@@ -332,7 +333,7 @@ def get(file_id):
         if not e:
             return get_data_error_result(retmsg="Document not found!")
         b, n = File2DocumentService.get_minio_address(file_id=file_id)
-        response = flask.make_response(MINIO.get(b, n))
+        response = flask.make_response(STORAGE_IMPL.get(b, n))
         ext = re.search(r"\.([^.]+)$", file.name)
         if ext:
             if file.type == FileType.VISUAL.value:

api/apps/kb_app.py

@@ -100,6 +100,15 @@ def update():
 def detail():
     kb_id = request.args["kb_id"]
     try:
+        tenants = UserTenantService.query(user_id=current_user.id)
+        for tenant in tenants:
+            if KnowledgebaseService.query(
+                    tenant_id=tenant.tenant_id, id=kb_id):
+                break
+        else:
+            return get_json_result(
+                data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
         kb = KnowledgebaseService.get_detail(kb_id)
         if not kb:
             return get_data_error_result(

api/apps/llm_app.py

@@ -20,16 +20,28 @@ from api.utils.api_utils import server_error_response, get_data_error_result, va
 from api.db import StatusEnum, LLMType
 from api.db.db_models import TenantLLM
 from api.utils.api_utils import get_json_result
-from rag.llm import EmbeddingModel, ChatModel, RerankModel,CvModel
+from rag.llm import EmbeddingModel, ChatModel, RerankModel, CvModel, TTSModel
 import requests
-import ast
+
 
 @manager.route('/factories', methods=['GET'])
 @login_required
 def factories():
     try:
         fac = LLMFactoriesService.get_all()
-        return get_json_result(data=[f.to_dict() for f in fac if f.name not in ["Youdao", "FastEmbed", "BAAI"]])
+        fac = [f.to_dict() for f in fac if f.name not in ["Youdao", "FastEmbed", "BAAI"]]
+        llms = LLMService.get_all()
+        mdl_types = {}
+        for m in llms:
+            if m.status != StatusEnum.VALID.value:
+                continue
+            if m.fid not in mdl_types:
+                mdl_types[m.fid] = set([])
+            mdl_types[m.fid].add(m.model_type)
+        for f in fac:
+            f["model_types"] = list(mdl_types.get(f["name"], [LLMType.CHAT, LLMType.EMBEDDING, LLMType.RERANK,
+                                                              LLMType.IMAGE2TEXT, LLMType.SPEECH2TEXT, LLMType.TTS]))
+        return get_json_result(data=fac)
     except Exception as e:
         return server_error_response(e)
 
@@ -43,7 +55,7 @@ def set_api_key():
     chat_passed, embd_passed, rerank_passed = False, False, False
     factory = req["llm_factory"]
     msg = ""
-    for llm in LLMService.query(fid=factory):
+    for llm in LLMService.query(fid=factory)[:3]:
         if not embd_passed and llm.model_type == LLMType.EMBEDDING.value:
             mdl = EmbeddingModel[factory](
                 req["api_key"], llm.llm_name, base_url=req.get("base_url"))
@@ -113,18 +125,19 @@ def add_llm():
 
     if factory == "VolcEngine":
         # For VolcEngine, due to its special authentication method
-        # Assemble volc_ak, volc_sk, endpoint_id into api_key
+        # Assemble ark_api_key endpoint_id into api_key
-        temp = list(ast.literal_eval(req["llm_name"]).items())[0]
+        llm_name = req["llm_name"]
-        llm_name = temp[0]
+        api_key = '{' + f'"ark_api_key": "{req.get("ark_api_key", "")}", ' \
-        endpoint_id = temp[1]
+                        f'"ep_id": "{req.get("endpoint_id", "")}", ' + '}'
-        api_key = '{' + f'"volc_ak": "{req.get("volc_ak", "")}", ' \
-                        f'"volc_sk": "{req.get("volc_sk", "")}", ' \
-                        f'"ep_id": "{endpoint_id}", ' + '}'
     elif factory == "Tencent Hunyuan":
         api_key = '{' + f'"hunyuan_sid": "{req.get("hunyuan_sid", "")}", ' \
                         f'"hunyuan_sk": "{req.get("hunyuan_sk", "")}"' + '}'
         req["api_key"] = api_key
         return set_api_key()
+    elif factory == "Tencent Cloud":
+        api_key = '{' + f'"tencent_cloud_sid": "{req.get("tencent_cloud_sid", "")}", ' \
+                f'"tencent_cloud_sk": "{req.get("tencent_cloud_sk", "")}"' + '}'
+        req["api_key"] = api_key
     elif factory == "Bedrock":
         # For Bedrock, due to its special authentication method
         # Assemble bedrock_ak, bedrock_sk, bedrock_region
@@ -145,6 +158,18 @@ def add_llm():
         llm_name = req["llm_name"]
         api_key = '{' + f'"yiyan_ak": "{req.get("yiyan_ak", "")}", ' \
                 f'"yiyan_sk": "{req.get("yiyan_sk", "")}"' + '}'
+    elif factory == "Fish Audio":
+        llm_name = req["llm_name"]
+        api_key = '{' + f'"fish_audio_ak": "{req.get("fish_audio_ak", "")}", ' \
+                f'"fish_audio_refid": "{req.get("fish_audio_refid", "59cb5986671546eaa6ca8ae6f29f6d22")}"' + '}'
+    elif factory == "Google Cloud":
+        llm_name = req["llm_name"]
+        api_key = (
+            "{" + f'"google_project_id": "{req.get("google_project_id", "")}", '
+            f'"google_region": "{req.get("google_region", "")}", '
+            f'"google_service_account_key": "{req.get("google_service_account_key", "")}"'
+            + "}"
+        )
     else:
         llm_name = req["llm_name"]
         api_key = req.get("api_key","xxxxxxxxxxxxxxx") 
@@ -218,6 +243,15 @@ def add_llm():
                 pass
         except Exception as e:
             msg += f"\nFail to access model({llm['llm_name']})." + str(e)
+    elif llm["model_type"] == LLMType.TTS:
+        mdl = TTSModel[factory](
+            key=llm["api_key"], model_name=llm["llm_name"], base_url=llm["api_base"]
+        )
+        try:
+            for resp in mdl.tts("Hello~ Ragflower!"):
+                pass
+        except RuntimeError as e:
+            msg += f"\nFail to access model({llm['llm_name']})." + str(e)
     else:
         # TODO: check other type of models
         pass

api/apps/sdk/assistant.py

@@ -0,0 +1,304 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+from flask import request
+
+from api.db import StatusEnum
+from api.db.db_models import TenantLLM
+from api.db.services.dialog_service import DialogService
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.llm_service import LLMService, TenantLLMService
+from api.db.services.user_service import TenantService
+from api.settings import RetCode
+from api.utils import get_uuid
+from api.utils.api_utils import get_data_error_result, token_required
+from api.utils.api_utils import get_json_result
+
+
+@manager.route('/save', methods=['POST'])
+@token_required
+def save(tenant_id):
+    req = request.json
+    # dataset
+    if req.get("knowledgebases") == []:
+        return get_data_error_result(retmsg="knowledgebases can not be empty list")
+    kb_list = []
+    if req.get("knowledgebases"):
+        for kb in req.get("knowledgebases"):
+            if not kb["id"]:
+                return get_data_error_result(retmsg="knowledgebase needs id")
+            if not KnowledgebaseService.query(id=kb["id"], tenant_id=tenant_id):
+                return get_data_error_result(retmsg="you do not own the knowledgebase")
+            # if not DocumentService.query(kb_id=kb["id"]):
+            #  return get_data_error_result(retmsg="There is a invalid knowledgebase")
+            kb_list.append(kb["id"])
+    req["kb_ids"] = kb_list
+    # llm
+    llm = req.get("llm")
+    if llm:
+        if "model_name" in llm:
+            req["llm_id"] = llm.pop("model_name")
+        req["llm_setting"] = req.pop("llm")
+    e, tenant = TenantService.get_by_id(tenant_id)
+    if not e:
+        return get_data_error_result(retmsg="Tenant not found!")
+    # prompt
+    prompt = req.get("prompt")
+    key_mapping = {"parameters": "variables",
+                   "prologue": "opener",
+                   "quote": "show_quote",
+                   "system": "prompt",
+                   "rerank_id": "rerank_model",
+                   "vector_similarity_weight": "keywords_similarity_weight"}
+    key_list = ["similarity_threshold", "vector_similarity_weight", "top_n", "rerank_id"]
+    if prompt:
+        for new_key, old_key in key_mapping.items():
+            if old_key in prompt:
+                prompt[new_key] = prompt.pop(old_key)
+        for key in key_list:
+            if key in prompt:
+                req[key] = prompt.pop(key)
+        req["prompt_config"] = req.pop("prompt")
+    # create
+    if "id" not in req:
+        # dataset
+        if not kb_list:
+            return get_data_error_result(retmsg="knowledgebases are required!")
+        # init
+        req["id"] = get_uuid()
+        req["description"] = req.get("description", "A helpful Assistant")
+        req["icon"] = req.get("avatar", "")
+        req["top_n"] = req.get("top_n", 6)
+        req["top_k"] = req.get("top_k", 1024)
+        req["rerank_id"] = req.get("rerank_id", "")
+        if req.get("llm_id"):
+            if not TenantLLMService.query(llm_name=req["llm_id"]):
+                return get_data_error_result(retmsg="the model_name does not exist.")
+        else:
+            req["llm_id"] = tenant.llm_id
+        if not req.get("name"):
+            return get_data_error_result(retmsg="name is required.")
+        if DialogService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value):
+            return get_data_error_result(retmsg="Duplicated assistant name in creating dataset.")
+        # tenant_id
+        if req.get("tenant_id"):
+            return get_data_error_result(retmsg="tenant_id must not be provided.")
+        req["tenant_id"] = tenant_id
+        # prompt more parameter
+        default_prompt = {
+            "system": """你是一个智能助手，请总结知识库的内容来回答问题，请列举知识库中的数据详细回答。当所有知识库内容都与问题无关时，你的回答必须包括“知识库中未找到您要的答案！”这句话。回答需要考虑聊天历史。
+                以下是知识库：
+                {knowledge}
+                以上是知识库。""",
+            "prologue": "您好，我是您的助手小樱，长得可爱又善良，can I help you?",
+            "parameters": [
+                {"key": "knowledge", "optional": False}
+            ],
+            "empty_response": "Sorry! 知识库中未找到相关内容！"
+        }
+        key_list_2 = ["system", "prologue", "parameters", "empty_response"]
+        if "prompt_config" not in req:
+            req['prompt_config'] = {}
+        for key in key_list_2:
+            temp = req['prompt_config'].get(key)
+            if not temp:
+                req['prompt_config'][key] = default_prompt[key]
+        for p in req['prompt_config']["parameters"]:
+            if p["optional"]:
+                continue
+            if req['prompt_config']["system"].find("{%s}" % p["key"]) < 0:
+                return get_data_error_result(
+                    retmsg="Parameter '{}' is not used".format(p["key"]))
+        # save
+        if not DialogService.save(**req):
+            return get_data_error_result(retmsg="Fail to new an assistant!")
+        # response
+        e, res = DialogService.get_by_id(req["id"])
+        if not e:
+            return get_data_error_result(retmsg="Fail to new an assistant!")
+        res = res.to_json()
+        renamed_dict = {}
+        for key, value in res["prompt_config"].items():
+            new_key = key_mapping.get(key, key)
+            renamed_dict[new_key] = value
+        res["prompt"] = renamed_dict
+        del res["prompt_config"]
+        new_dict = {"similarity_threshold": res["similarity_threshold"],
+                    "keywords_similarity_weight": res["vector_similarity_weight"],
+                    "top_n": res["top_n"],
+                    "rerank_model": res['rerank_id']}
+        res["prompt"].update(new_dict)
+        for key in key_list:
+            del res[key]
+        res["llm"] = res.pop("llm_setting")
+        res["llm"]["model_name"] = res.pop("llm_id")
+        del res["kb_ids"]
+        res["knowledgebases"] = req["knowledgebases"]
+        res["avatar"] = res.pop("icon")
+        return get_json_result(data=res)
+    else:
+        # authorization
+        if not DialogService.query(tenant_id=tenant_id, id=req["id"], status=StatusEnum.VALID.value):
+            return get_json_result(data=False, retmsg='You do not own the assistant', retcode=RetCode.OPERATING_ERROR)
+        # prompt
+        if not req["id"]:
+            return get_data_error_result(retmsg="id can not be empty")
+        e, res = DialogService.get_by_id(req["id"])
+        res = res.to_json()
+        if "llm_id" in req:
+            if not TenantLLMService.query(llm_name=req["llm_id"]):
+                return get_data_error_result(retmsg="the model_name does not exist.")
+        if "name" in req:
+            if not req.get("name"):
+                return get_data_error_result(retmsg="name is not empty.")
+            if req["name"].lower() != res["name"].lower() \
+                    and len(
+                DialogService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value)) > 0:
+                return get_data_error_result(retmsg="Duplicated assistant name in updating dataset.")
+        if "prompt_config" in req:
+            res["prompt_config"].update(req["prompt_config"])
+            for p in res["prompt_config"]["parameters"]:
+                if p["optional"]:
+                    continue
+                if res["prompt_config"]["system"].find("{%s}" % p["key"]) < 0:
+                    return get_data_error_result(retmsg="Parameter '{}' is not used".format(p["key"]))
+        if "llm_setting" in req:
+            res["llm_setting"].update(req["llm_setting"])
+        req["prompt_config"] = res["prompt_config"]
+        req["llm_setting"] = res["llm_setting"]
+        # avatar
+        if "avatar" in req:
+            req["icon"] = req.pop("avatar")
+        assistant_id = req.pop("id")
+        if "knowledgebases" in req:
+            req.pop("knowledgebases")
+        if not DialogService.update_by_id(assistant_id, req):
+            return get_data_error_result(retmsg="Assistant not found!")
+        return get_json_result(data=True)
+
+
+@manager.route('/delete', methods=['DELETE'])
+@token_required
+def delete(tenant_id):
+    req = request.args
+    if "id" not in req:
+        return get_data_error_result(retmsg="id is required")
+    id = req['id']
+    if not DialogService.query(tenant_id=tenant_id, id=id, status=StatusEnum.VALID.value):
+        return get_json_result(data=False, retmsg='you do not own the assistant.', retcode=RetCode.OPERATING_ERROR)
+
+    temp_dict = {"status": StatusEnum.INVALID.value}
+    DialogService.update_by_id(req["id"], temp_dict)
+    return get_json_result(data=True)
+
+
+@manager.route('/get', methods=['GET'])
+@token_required
+def get(tenant_id):
+    req = request.args
+    if "id" in req:
+        id = req["id"]
+        ass = DialogService.query(tenant_id=tenant_id, id=id, status=StatusEnum.VALID.value)
+        if not ass:
+            return get_json_result(data=False, retmsg='You do not own the assistant.', retcode=RetCode.OPERATING_ERROR)
+        if "name" in req:
+            name = req["name"]
+            if ass[0].name != name:
+                return get_json_result(data=False, retmsg='name does not match id.', retcode=RetCode.OPERATING_ERROR)
+        res = ass[0].to_json()
+    else:
+        if "name" in req:
+            name = req["name"]
+            ass = DialogService.query(name=name, tenant_id=tenant_id, status=StatusEnum.VALID.value)
+            if not ass:
+                return get_json_result(data=False, retmsg='You do not own the assistant.',
+                                       retcode=RetCode.OPERATING_ERROR)
+            res = ass[0].to_json()
+        else:
+            return get_data_error_result(retmsg="At least one of `id` or `name` must be provided.")
+    renamed_dict = {}
+    key_mapping = {"parameters": "variables",
+                   "prologue": "opener",
+                   "quote": "show_quote",
+                   "system": "prompt",
+                   "rerank_id": "rerank_model",
+                   "vector_similarity_weight": "keywords_similarity_weight"}
+    key_list = ["similarity_threshold", "vector_similarity_weight", "top_n", "rerank_id"]
+    for key, value in res["prompt_config"].items():
+        new_key = key_mapping.get(key, key)
+        renamed_dict[new_key] = value
+    res["prompt"] = renamed_dict
+    del res["prompt_config"]
+    new_dict = {"similarity_threshold": res["similarity_threshold"],
+                "keywords_similarity_weight": res["vector_similarity_weight"],
+                "top_n": res["top_n"],
+                "rerank_model": res['rerank_id']}
+    res["prompt"].update(new_dict)
+    for key in key_list:
+        del res[key]
+    res["llm"] = res.pop("llm_setting")
+    res["llm"]["model_name"] = res.pop("llm_id")
+    kb_list = []
+    for kb_id in res["kb_ids"]:
+        kb = KnowledgebaseService.query(id=kb_id)
+        kb_list.append(kb[0].to_json())
+    del res["kb_ids"]
+    res["knowledgebases"] = kb_list
+    res["avatar"] = res.pop("icon")
+    return get_json_result(data=res)
+
+
+@manager.route('/list', methods=['GET'])
+@token_required
+def list_assistants(tenant_id):
+    assts = DialogService.query(
+        tenant_id=tenant_id,
+        status=StatusEnum.VALID.value,
+        reverse=True,
+        order_by=DialogService.model.create_time)
+    assts = [d.to_dict() for d in assts]
+    list_assts = []
+    renamed_dict = {}
+    key_mapping = {"parameters": "variables",
+                   "prologue": "opener",
+                   "quote": "show_quote",
+                   "system": "prompt",
+                   "rerank_id": "rerank_model",
+                   "vector_similarity_weight": "keywords_similarity_weight"}
+    key_list = ["similarity_threshold", "vector_similarity_weight", "top_n", "rerank_id"]
+    for res in assts:
+        for key, value in res["prompt_config"].items():
+            new_key = key_mapping.get(key, key)
+            renamed_dict[new_key] = value
+        res["prompt"] = renamed_dict
+        del res["prompt_config"]
+        new_dict = {"similarity_threshold": res["similarity_threshold"],
+                    "keywords_similarity_weight": res["vector_similarity_weight"],
+                    "top_n": res["top_n"],
+                    "rerank_model": res['rerank_id']}
+        res["prompt"].update(new_dict)
+        for key in key_list:
+            del res[key]
+        res["llm"] = res.pop("llm_setting")
+        res["llm"]["model_name"] = res.pop("llm_id")
+        kb_list = []
+        for kb_id in res["kb_ids"]:
+            kb = KnowledgebaseService.query(id=kb_id)
+            kb_list.append(kb[0].to_json())
+        del res["kb_ids"]
+        res["knowledgebases"] = kb_list
+        res["avatar"] = res.pop("icon")
+        list_assts.append(res)
+    return get_json_result(data=list_assts)

api/apps/sdk/dataset.py

@@ -0,0 +1,224 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+
+from flask import request
+
+from api.db import StatusEnum, FileSource
+from api.db.db_models import File
+from api.db.services.document_service import DocumentService
+from api.db.services.file2document_service import File2DocumentService
+from api.db.services.file_service import FileService
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.user_service import TenantService
+from api.settings import RetCode
+from api.utils import get_uuid
+from api.utils.api_utils import get_json_result, token_required, get_data_error_result
+
+
+@manager.route('/save', methods=['POST'])
+@token_required
+def save(tenant_id):
+    req = request.json
+    e, t = TenantService.get_by_id(tenant_id)
+    if "id" not in req:
+        if "tenant_id" in req or "embedding_model" in req:
+            return get_data_error_result(
+                retmsg="Tenant_id or embedding_model must not be provided")
+        if "name" not in req:
+            return get_data_error_result(
+                retmsg="Name is not empty!")
+        req['id'] = get_uuid()
+        req["name"] = req["name"].strip()
+        if req["name"] == "":
+            return get_data_error_result(
+                retmsg="Name is not empty string!")
+        if KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value):
+            return get_data_error_result(
+                retmsg="Duplicated knowledgebase name in creating dataset.")
+        req["tenant_id"] = req['created_by'] = tenant_id
+        req['embedding_model'] = t.embd_id
+        key_mapping = {
+            "chunk_num": "chunk_count",
+            "doc_num": "document_count",
+            "parser_id": "parse_method",
+            "embd_id": "embedding_model"
+        }
+        mapped_keys = {new_key: req[old_key] for new_key, old_key in key_mapping.items() if old_key in req}
+        req.update(mapped_keys)
+        if not KnowledgebaseService.save(**req):
+            return get_data_error_result(retmsg="Create dataset error.(Database error)")
+        renamed_data = {}
+        e, k = KnowledgebaseService.get_by_id(req["id"])
+        for key, value in k.to_dict().items():
+            new_key = key_mapping.get(key, key)
+            renamed_data[new_key] = value
+        return get_json_result(data=renamed_data)
+    else:
+        invalid_keys = {"embd_id", "chunk_num", "doc_num", "parser_id"}
+        if any(key in req for key in invalid_keys):
+            return get_data_error_result(retmsg="The input parameters are invalid.")
+
+        if "tenant_id" in req:
+            if req["tenant_id"] != tenant_id:
+                return get_data_error_result(
+                    retmsg="Can't change tenant_id.")
+
+        if "embedding_model" in req:
+            if req["embedding_model"] != t.embd_id:
+                return get_data_error_result(
+                    retmsg="Can't change embedding_model.")
+            req.pop("embedding_model")
+
+        if not KnowledgebaseService.query(
+                created_by=tenant_id, id=req["id"]):
+            return get_json_result(
+                data=False, retmsg='You do not own the dataset.',
+                retcode=RetCode.OPERATING_ERROR)
+
+        if not req["id"]:
+            return get_data_error_result(
+                retmsg="id can not be empty.")
+        e, kb = KnowledgebaseService.get_by_id(req["id"])
+
+        if "chunk_count" in req:
+            if req["chunk_count"] != kb.chunk_num:
+                return get_data_error_result(
+                    retmsg="Can't change chunk_count.")
+            req.pop("chunk_count")
+
+        if "document_count" in req:
+            if req['document_count'] != kb.doc_num:
+                return get_data_error_result(
+                    retmsg="Can't change document_count.")
+            req.pop("document_count")
+
+        if "parse_method" in req:
+            if kb.chunk_num != 0 and req['parse_method'] != kb.parser_id:
+                return get_data_error_result(
+                    retmsg="If chunk count is not 0, parse method is not changable.")
+            req['parser_id'] = req.pop('parse_method')
+        if "name" in req:
+            req["name"] = req["name"].strip()
+            if req["name"].lower() != kb.name.lower() \
+                    and len(KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id,
+                                                       status=StatusEnum.VALID.value)) > 0:
+                return get_data_error_result(
+                    retmsg="Duplicated knowledgebase name in updating dataset.")
+
+        del req["id"]
+        if not KnowledgebaseService.update_by_id(kb.id, req):
+            return get_data_error_result(retmsg="Update dataset error.(Database error)")
+        return get_json_result(data=True)
+
+
+@manager.route('/delete', methods=['DELETE'])
+@token_required
+def delete(tenant_id):
+    req = request.args
+    if "id" not in req:
+        return get_data_error_result(
+            retmsg="id is required")
+    kbs = KnowledgebaseService.query(
+        created_by=tenant_id, id=req["id"])
+    if not kbs:
+        return get_json_result(
+            data=False, retmsg='You do not own the dataset',
+            retcode=RetCode.OPERATING_ERROR)
+
+    for doc in DocumentService.query(kb_id=req["id"]):
+        if not DocumentService.remove_document(doc, kbs[0].tenant_id):
+            return get_data_error_result(
+                retmsg="Remove document error.(Database error)")
+        f2d = File2DocumentService.get_by_document_id(doc.id)
+        FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
+        File2DocumentService.delete_by_document_id(doc.id)
+
+    if not KnowledgebaseService.delete_by_id(req["id"]):
+        return get_data_error_result(
+            retmsg="Delete dataset error.(Database serror)")
+    return get_json_result(data=True)
+
+
+@manager.route('/list', methods=['GET'])
+@token_required
+def list_datasets(tenant_id):
+    page_number = int(request.args.get("page", 1))
+    items_per_page = int(request.args.get("page_size", 1024))
+    orderby = request.args.get("orderby", "create_time")
+    desc = bool(request.args.get("desc", True))
+    tenants = TenantService.get_joined_tenants_by_user_id(tenant_id)
+    kbs = KnowledgebaseService.get_by_tenant_ids(
+        [m["tenant_id"] for m in tenants], tenant_id, page_number, items_per_page, orderby, desc)
+    renamed_list = []
+    for kb in kbs:
+        key_mapping = {
+            "chunk_num": "chunk_count",
+            "doc_num": "document_count",
+            "parser_id": "parse_method",
+            "embd_id": "embedding_model"
+        }
+        renamed_data = {}
+        for key, value in kb.items():
+            new_key = key_mapping.get(key, key)
+            renamed_data[new_key] = value
+        renamed_list.append(renamed_data)
+    return get_json_result(data=renamed_list)
+
+
+@manager.route('/detail', methods=['GET'])
+@token_required
+def detail(tenant_id):
+    req = request.args
+    key_mapping = {
+        "chunk_num": "chunk_count",
+        "doc_num": "document_count",
+        "parser_id": "parse_method",
+        "embd_id": "embedding_model"
+    }
+    renamed_data = {}
+    if "id" in req:
+        id = req["id"]
+        kb = KnowledgebaseService.query(created_by=tenant_id, id=req["id"])
+        if not kb:
+            return get_json_result(
+                data=False, retmsg='You do not own the dataset.',
+                retcode=RetCode.OPERATING_ERROR)
+        if "name" in req:
+            name = req["name"]
+            if kb[0].name != name:
+                return get_json_result(
+                    data=False, retmsg='You do not own the dataset.',
+                    retcode=RetCode.OPERATING_ERROR)
+        e, k = KnowledgebaseService.get_by_id(id)
+        for key, value in k.to_dict().items():
+            new_key = key_mapping.get(key, key)
+            renamed_data[new_key] = value
+        return get_json_result(data=renamed_data)
+    else:
+        if "name" in req:
+            name = req["name"]
+            e, k = KnowledgebaseService.get_by_name(kb_name=name, tenant_id=tenant_id)
+            if not e:
+                return get_json_result(
+                    data=False, retmsg='You do not own the dataset.',
+                    retcode=RetCode.OPERATING_ERROR)
+            for key, value in k.to_dict().items():
+                new_key = key_mapping.get(key, key)
+                renamed_data[new_key] = value
+            return get_json_result(data=renamed_data)
+        else:
+            return get_data_error_result(
+                retmsg="At least one of `id` or `name` must be provided.")

api/apps/sdk/doc.py

@@ -0,0 +1,529 @@
+import pathlib
+import re
+import datetime
+import json
+import traceback
+
+from flask import request
+from flask_login import login_required, current_user
+from elasticsearch_dsl import Q
+
+from rag.app.qa import rmPrefix, beAdoc
+from rag.nlp import search, rag_tokenizer, keyword_extraction
+from rag.utils.es_conn import ELASTICSEARCH
+from rag.utils import rmSpace
+from api.db import LLMType, ParserType
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.llm_service import TenantLLMService
+from api.db.services.user_service import UserTenantService
+from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
+from api.db.services.document_service import DocumentService
+from api.settings import RetCode, retrievaler, kg_retrievaler
+from api.utils.api_utils import get_json_result
+import hashlib
+import re
+from api.utils.api_utils import get_json_result, token_required, get_data_error_result
+
+from api.db.db_models import Task, File
+
+from api.db.services.task_service import TaskService, queue_tasks
+from api.db.services.user_service import TenantService, UserTenantService
+
+from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
+
+from api.utils.api_utils import get_json_result
+
+from functools import partial
+from io import BytesIO
+
+from elasticsearch_dsl import Q
+from flask import request, send_file
+from flask_login import login_required
+
+from api.db import FileSource, TaskStatus, FileType
+from api.db.db_models import File
+from api.db.services.document_service import DocumentService
+from api.db.services.file2document_service import File2DocumentService
+from api.db.services.file_service import FileService
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.settings import RetCode, retrievaler
+from api.utils.api_utils import construct_json_result, construct_error_response
+from rag.app import book, laws, manual, naive, one, paper, presentation, qa, resume, table, picture, audio, email
+from rag.nlp import search
+from rag.utils import rmSpace
+from rag.utils.es_conn import ELASTICSEARCH
+from rag.utils.storage_factory import STORAGE_IMPL
+
+MAXIMUM_OF_UPLOADING_FILES = 256
+
+MAXIMUM_OF_UPLOADING_FILES = 256
+
+
+@manager.route('/dataset/<dataset_id>/documents/upload', methods=['POST'])
+@token_required
+def upload(dataset_id, tenant_id):
+    if 'file' not in request.files:
+        return get_json_result(
+            data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
+    file_objs = request.files.getlist('file')
+    for file_obj in file_objs:
+        if file_obj.filename == '':
+            return get_json_result(
+                data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
+    e, kb = KnowledgebaseService.get_by_id(dataset_id)
+    if not e:
+        raise LookupError(f"Can't find the knowledgebase with ID {dataset_id}!")
+    err, _ = FileService.upload_document(kb, file_objs, tenant_id)
+    if err:
+        return get_json_result(
+            data=False, retmsg="\n".join(err), retcode=RetCode.SERVER_ERROR)
+    return get_json_result(data=True)
+
+
+@manager.route('/infos', methods=['GET'])
+@token_required
+def docinfos(tenant_id):
+    req = request.args
+    if "id" in req:
+        doc_id = req["id"]
+        e, doc = DocumentService.get_by_id(doc_id)
+        return get_json_result(data=doc.to_json())
+    if "name" in req:
+        doc_name = req["name"]
+        doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
+        e, doc = DocumentService.get_by_id(doc_id)
+        return get_json_result(data=doc.to_json())
+
+
+@manager.route('/save', methods=['POST'])
+@token_required
+def save_doc(tenant_id):
+    req = request.json
+    #get doc by id or name
+    doc_id = None
+    if "id" in req:
+        doc_id = req["id"]
+    elif "name" in req:
+        doc_name = req["name"]
+        doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
+    if not doc_id:
+        return get_json_result(retcode=400, retmsg="Document ID or name is required")
+    e, doc = DocumentService.get_by_id(doc_id)
+    if not e:
+        return get_data_error_result(retmsg="Document not found!")
+    #other value can't be changed
+    if "chunk_num" in req:
+        if req["chunk_num"] != doc.chunk_num:
+            return get_data_error_result(
+                retmsg="Can't change chunk_count.")
+    if "progress" in req:
+        if req['progress'] != doc.progress:
+            return get_data_error_result(
+                retmsg="Can't change progress.")
+    #change name or parse_method
+    if "name" in req and req["name"] != doc.name:
+        try:
+            if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
+                    doc.name.lower()).suffix:
+                return get_json_result(
+                    data=False,
+                    retmsg="The extension of file can't be changed",
+                    retcode=RetCode.ARGUMENT_ERROR)
+            for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
+                if d.name == req["name"]:
+                    return get_data_error_result(
+                        retmsg="Duplicated document name in the same knowledgebase.")
+
+            if not DocumentService.update_by_id(
+                    doc_id, {"name": req["name"]}):
+                return get_data_error_result(
+                    retmsg="Database error (Document rename)!")
+
+            informs = File2DocumentService.get_by_document_id(doc_id)
+            if informs:
+                e, file = FileService.get_by_id(informs[0].file_id)
+                FileService.update_by_id(file.id, {"name": req["name"]})
+        except Exception as e:
+            return server_error_response(e)
+    if "parser_id" in req:
+        try:
+            if doc.parser_id.lower() == req["parser_id"].lower():
+                if "parser_config" in req:
+                    if req["parser_config"] == doc.parser_config:
+                        return get_json_result(data=True)
+                else:
+                    return get_json_result(data=True)
+
+            if doc.type == FileType.VISUAL or re.search(
+                    r"\.(ppt|pptx|pages)$", doc.name):
+                return get_data_error_result(retmsg="Not supported yet!")
+
+            e = DocumentService.update_by_id(doc.id,
+                                             {"parser_id": req["parser_id"], "progress": 0, "progress_msg": "",
+                                              "run": TaskStatus.UNSTART.value})
+            if not e:
+                return get_data_error_result(retmsg="Document not found!")
+            if "parser_config" in req:
+                DocumentService.update_parser_config(doc.id, req["parser_config"])
+            if doc.token_num > 0:
+                e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
+                                                        doc.process_duation * -1)
+                if not e:
+                    return get_data_error_result(retmsg="Document not found!")
+                tenant_id = DocumentService.get_tenant_id(req["doc_id"])
+                if not tenant_id:
+                    return get_data_error_result(retmsg="Tenant not found!")
+                ELASTICSEARCH.deleteByQuery(
+                    Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
+        except Exception as e:
+            return server_error_response(e)
+    return get_json_result(data=True)
+
+
+
+@manager.route('/change_parser', methods=['POST'])
+@token_required
+def change_parser(tenant_id):
+    req = request.json
+    try:
+        e, doc = DocumentService.get_by_id(req["doc_id"])
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        if doc.parser_id.lower() == req["parser_id"].lower():
+            if "parser_config" in req:
+                if req["parser_config"] == doc.parser_config:
+                    return get_json_result(data=True)
+            else:
+                return get_json_result(data=True)
+
+        if doc.type == FileType.VISUAL or re.search(
+                r"\.(ppt|pptx|pages)$", doc.name):
+            return get_data_error_result(retmsg="Not supported yet!")
+
+        e = DocumentService.update_by_id(doc.id,
+                                         {"parser_id": req["parser_id"], "progress": 0, "progress_msg": "",
+                                          "run": TaskStatus.UNSTART.value})
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        if "parser_config" in req:
+            DocumentService.update_parser_config(doc.id, req["parser_config"])
+        if doc.token_num > 0:
+            e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
+                                                    doc.process_duation * -1)
+            if not e:
+                return get_data_error_result(retmsg="Document not found!")
+            tenant_id = DocumentService.get_tenant_id(req["doc_id"])
+            if not tenant_id:
+                return get_data_error_result(retmsg="Tenant not found!")
+            ELASTICSEARCH.deleteByQuery(
+                Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
+
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+
+@manager.route('/rename', methods=['POST'])
+@login_required
+@validate_request("doc_id", "name")
+def rename():
+    req = request.json
+    try:
+        e, doc = DocumentService.get_by_id(req["doc_id"])
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
+                doc.name.lower()).suffix:
+            return get_json_result(
+                data=False,
+                retmsg="The extension of file can't be changed",
+                retcode=RetCode.ARGUMENT_ERROR)
+        for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
+            if d.name == req["name"]:
+                return get_data_error_result(
+                    retmsg="Duplicated document name in the same knowledgebase.")
+
+        if not DocumentService.update_by_id(
+                req["doc_id"], {"name": req["name"]}):
+            return get_data_error_result(
+                retmsg="Database error (Document rename)!")
+
+        informs = File2DocumentService.get_by_document_id(req["doc_id"])
+        if informs:
+            e, file = FileService.get_by_id(informs[0].file_id)
+            FileService.update_by_id(file.id, {"name": req["name"]})
+
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route("/<document_id>", methods=["GET"])
+@token_required
+def download_document(dataset_id, document_id):
+    try:
+        # Check whether there is this document
+        exist, document = DocumentService.get_by_id(document_id)
+        if not exist:
+            return construct_json_result(message=f"This document '{document_id}' cannot be found!",
+                                         code=RetCode.ARGUMENT_ERROR)
+
+        # The process of downloading
+        doc_id, doc_location = File2DocumentService.get_minio_address(doc_id=document_id)  # minio address
+        file_stream = STORAGE_IMPL.get(doc_id, doc_location)
+        if not file_stream:
+            return construct_json_result(message="This file is empty.", code=RetCode.DATA_ERROR)
+
+        file = BytesIO(file_stream)
+
+        # Use send_file with a proper filename and MIME type
+        return send_file(
+            file,
+            as_attachment=True,
+            download_name=document.name,
+            mimetype='application/octet-stream'  # Set a default MIME type
+        )
+
+    # Error
+    except Exception as e:
+        return construct_error_response(e)
+
+
+@manager.route('/dataset/<dataset_id>/documents', methods=['GET'])
+@token_required
+def list_docs(dataset_id, tenant_id):
+    kb_id = request.args.get("kb_id")
+    if not kb_id:
+        return get_json_result(
+            data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
+    tenants = UserTenantService.query(user_id=tenant_id)
+    for tenant in tenants:
+        if KnowledgebaseService.query(
+                tenant_id=tenant.tenant_id, id=kb_id):
+            break
+    else:
+        return get_json_result(
+            data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+            retcode=RetCode.OPERATING_ERROR)
+    keywords = request.args.get("keywords", "")
+
+    page_number = int(request.args.get("page", 1))
+    items_per_page = int(request.args.get("page_size", 15))
+    orderby = request.args.get("orderby", "create_time")
+    desc = request.args.get("desc", True)
+    try:
+        docs, tol = DocumentService.get_by_kb_id(
+            kb_id, page_number, items_per_page, orderby, desc, keywords)
+        return get_json_result(data={"total": tol, "docs": docs})
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/delete', methods=['DELETE'])
+@token_required
+def rm(tenant_id):
+    req = request.args
+    if "doc_id" not in req:
+        return get_data_error_result(
+            retmsg="doc_id is required")
+    doc_ids = req["doc_id"]
+    if isinstance(doc_ids, str): doc_ids = [doc_ids]
+    root_folder = FileService.get_root_folder(tenant_id)
+    pf_id = root_folder["id"]
+    FileService.init_knowledgebase_docs(pf_id, tenant_id)
+    errors = ""
+    for doc_id in doc_ids:
+        try:
+            e, doc = DocumentService.get_by_id(doc_id)
+            if not e:
+                return get_data_error_result(retmsg="Document not found!")
+            tenant_id = DocumentService.get_tenant_id(doc_id)
+            if not tenant_id:
+                return get_data_error_result(retmsg="Tenant not found!")
+
+            b, n = File2DocumentService.get_minio_address(doc_id=doc_id)
+
+            if not DocumentService.remove_document(doc, tenant_id):
+                return get_data_error_result(
+                    retmsg="Database error (Document removal)!")
+
+            f2d = File2DocumentService.get_by_document_id(doc_id)
+            FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
+            File2DocumentService.delete_by_document_id(doc_id)
+
+            STORAGE_IMPL.rm(b, n)
+        except Exception as e:
+            errors += str(e)
+
+    if errors:
+        return get_json_result(data=False, retmsg=errors, retcode=RetCode.SERVER_ERROR)
+
+    return get_json_result(data=True, retmsg="success")
+
+@manager.route("/<document_id>/status", methods=["GET"])
+@token_required
+def show_parsing_status(tenant_id, document_id):
+    try:
+        # valid document
+        exist, _ = DocumentService.get_by_id(document_id)
+        if not exist:
+            return construct_json_result(code=RetCode.DATA_ERROR,
+                                         message=f"This document: '{document_id}' is not a valid document.")
+
+        _, doc = DocumentService.get_by_id(document_id)  # get doc object
+        doc_attributes = doc.to_dict()
+
+        return construct_json_result(
+            data={"progress": doc_attributes["progress"], "status": TaskStatus(doc_attributes["status"]).name},
+            code=RetCode.SUCCESS
+        )
+    except Exception as e:
+        return construct_error_response(e)
+
+
+
+@manager.route('/run', methods=['POST'])
+@token_required
+def run(tenant_id):
+    req = request.json
+    try:
+        for id in req["doc_ids"]:
+            info = {"run": str(req["run"]), "progress": 0}
+            if str(req["run"]) == TaskStatus.RUNNING.value:
+                info["progress_msg"] = ""
+                info["chunk_num"] = 0
+                info["token_num"] = 0
+            DocumentService.update_by_id(id, info)
+            # if str(req["run"]) == TaskStatus.CANCEL.value:
+            tenant_id = DocumentService.get_tenant_id(id)
+            if not tenant_id:
+                return get_data_error_result(retmsg="Tenant not found!")
+            ELASTICSEARCH.deleteByQuery(
+                Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
+
+            if str(req["run"]) == TaskStatus.RUNNING.value:
+                TaskService.filter_delete([Task.doc_id == id])
+                e, doc = DocumentService.get_by_id(id)
+                doc = doc.to_dict()
+                doc["tenant_id"] = tenant_id
+                bucket, name = File2DocumentService.get_minio_address(doc_id=doc["id"])
+                queue_tasks(doc, bucket, name)
+
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/chunk/list', methods=['POST'])
+@token_required
+@validate_request("doc_id")
+def list_chunk(tenant_id):
+    req = request.json
+    doc_id = req["doc_id"]
+    page = int(req.get("page", 1))
+    size = int(req.get("size", 30))
+    question = req.get("keywords", "")
+    try:
+        tenant_id = DocumentService.get_tenant_id(req["doc_id"])
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
+        e, doc = DocumentService.get_by_id(doc_id)
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        query = {
+            "doc_ids": [doc_id], "page": page, "size": size, "question": question, "sort": True
+        }
+        if "available_int" in req:
+            query["available_int"] = int(req["available_int"])
+        sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
+        res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
+        for id in sres.ids:
+            d = {
+                "chunk_id": id,
+                "content_with_weight": rmSpace(sres.highlight[id]) if question and id in sres.highlight else sres.field[
+                    id].get(
+                    "content_with_weight", ""),
+                "doc_id": sres.field[id]["doc_id"],
+                "docnm_kwd": sres.field[id]["docnm_kwd"],
+                "important_kwd": sres.field[id].get("important_kwd", []),
+                "img_id": sres.field[id].get("img_id", ""),
+                "available_int": sres.field[id].get("available_int", 1),
+                "positions": sres.field[id].get("position_int", "").split("\t")
+            }
+            if len(d["positions"]) % 5 == 0:
+                poss = []
+                for i in range(0, len(d["positions"]), 5):
+                    poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]),
+                                 float(d["positions"][i + 3]), float(d["positions"][i + 4])])
+                d["positions"] = poss
+            res["chunks"].append(d)
+        return get_json_result(data=res)
+    except Exception as e:
+        if str(e).find("not_found") > 0:
+            return get_json_result(data=False, retmsg=f'No chunk found!',
+                                   retcode=RetCode.DATA_ERROR)
+        return server_error_response(e)
+
+
+@manager.route('/chunk/create', methods=['POST'])
+@token_required
+@validate_request("doc_id", "content_with_weight")
+def create(tenant_id):
+    req = request.json
+    md5 = hashlib.md5()
+    md5.update((req["content_with_weight"] + req["doc_id"]).encode("utf-8"))
+    chunck_id = md5.hexdigest()
+    d = {"id": chunck_id, "content_ltks": rag_tokenizer.tokenize(req["content_with_weight"]),
+         "content_with_weight": req["content_with_weight"]}
+    d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
+    d["important_kwd"] = req.get("important_kwd", [])
+    d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_kwd", [])))
+    d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
+    d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
+
+    try:
+        e, doc = DocumentService.get_by_id(req["doc_id"])
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        d["kb_id"] = [doc.kb_id]
+        d["docnm_kwd"] = doc.name
+        d["doc_id"] = doc.id
+
+        tenant_id = DocumentService.get_tenant_id(req["doc_id"])
+        if not tenant_id:
+            return get_data_error_result(retmsg="Tenant not found!")
+
+        embd_id = DocumentService.get_embd_id(req["doc_id"])
+        embd_mdl = TenantLLMService.model_instance(
+            tenant_id, LLMType.EMBEDDING.value, embd_id)
+
+        v, c = embd_mdl.encode([doc.name, req["content_with_weight"]])
+        v = 0.1 * v[0] + 0.9 * v[1]
+        d["q_%d_vec" % len(v)] = v.tolist()
+        ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
+
+        DocumentService.increment_chunk_num(
+            doc.id, doc.kb_id, c, 1, 0)
+        return get_json_result(data={"chunk": d})
+        # return get_json_result(data={"chunk_id": chunck_id})
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/chunk/rm', methods=['POST'])
+@token_required
+@validate_request("chunk_ids", "doc_id")
+def rm_chunk():
+    req = request.json
+    try:
+        if not ELASTICSEARCH.deleteByQuery(
+                Q("ids", values=req["chunk_ids"]), search.index_name(current_user.id)):
+            return get_data_error_result(retmsg="Index updating failure")
+        e, doc = DocumentService.get_by_id(req["doc_id"])
+        if not e:
+            return get_data_error_result(retmsg="Document not found!")
+        deleted_chunk_ids = req["chunk_ids"]
+        chunk_number = len(deleted_chunk_ids)
+        DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)

api/apps/sdk/session.py

@@ -0,0 +1,263 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+import json
+from uuid import uuid4
+
+from flask import request, Response
+
+from api.db import StatusEnum
+from api.db.services.dialog_service import DialogService, ConversationService, chat
+from api.settings import RetCode
+from api.utils import get_uuid
+from api.utils.api_utils import get_data_error_result
+from api.utils.api_utils import get_json_result, token_required
+
+
+@manager.route('/save', methods=['POST'])
+@token_required
+def set_conversation(tenant_id):
+    req = request.json
+    conv_id = req.get("id")
+    if "assistant_id" in req:
+        req["dialog_id"] = req.pop("assistant_id")
+    if "id" in req:
+        del req["id"]
+        conv = ConversationService.query(id=conv_id)
+        if not conv:
+            return get_data_error_result(retmsg="Session does not exist")
+        if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+            return get_data_error_result(retmsg="You do not own the session")
+        if req.get("dialog_id"):
+            dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
+            if not dia:
+                return get_data_error_result(retmsg="You do not own the assistant")
+        if "dialog_id" in req and not req.get("dialog_id"):
+            return get_data_error_result(retmsg="assistant_id can not be empty.")
+        if "message" in req:
+            return get_data_error_result(retmsg="message can not be change")
+        if "reference" in req:
+            return get_data_error_result(retmsg="reference can not be change")
+        if "name" in req and not req.get("name"):
+            return get_data_error_result(retmsg="name can not be empty.")
+        if not ConversationService.update_by_id(conv_id, req):
+            return get_data_error_result(retmsg="Session updates error")
+        return get_json_result(data=True)
+
+    if not req.get("dialog_id"):
+        return get_data_error_result(retmsg="assistant_id is required.")
+    dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
+    if not dia:
+        return get_data_error_result(retmsg="You do not own the assistant")
+    conv = {
+        "id": get_uuid(),
+        "dialog_id": req["dialog_id"],
+        "name": req.get("name", "New session"),
+        "message": [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
+    }
+    if not conv.get("name"):
+        return get_data_error_result(retmsg="name can not be empty.")
+    ConversationService.save(**conv)
+    e, conv = ConversationService.get_by_id(conv["id"])
+    if not e:
+        return get_data_error_result(retmsg="Fail to new session!")
+    conv = conv.to_dict()
+    conv['messages'] = conv.pop("message")
+    conv["assistant_id"] = conv.pop("dialog_id")
+    del conv["reference"]
+    return get_json_result(data=conv)
+
+
+@manager.route('/completion', methods=['POST'])
+@token_required
+def completion(tenant_id):
+    req = request.json
+    # req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
+    #    {"role": "user", "content": "上海有吗？"}
+    # ]}
+    if "id" not in req:
+        return get_data_error_result(retmsg="id is required")
+    conv = ConversationService.query(id=req["id"])
+    if not conv:
+        return get_data_error_result(retmsg="Session does not exist")
+    conv = conv[0]
+    if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_data_error_result(retmsg="You do not own the session")
+    msg = []
+    question = {
+        "content": req.get("question"),
+        "role": "user",
+        "id": str(uuid4())
+    }
+    conv.message.append(question)
+    for m in conv.message:
+        if m["role"] == "system": continue
+        if m["role"] == "assistant" and not msg: continue
+        msg.append(m)
+    message_id = msg[-1].get("id")
+    e, dia = DialogService.get_by_id(conv.dialog_id)
+    del req["id"]
+
+    if not conv.reference:
+        conv.reference = []
+    conv.message.append({"role": "assistant", "content": "", "id": message_id})
+    conv.reference.append({"chunks": [], "doc_aggs": []})
+
+    def fillin_conv(ans):
+        nonlocal conv, message_id
+        if not conv.reference:
+            conv.reference.append(ans["reference"])
+        else:
+            conv.reference[-1] = ans["reference"]
+        conv.message[-1] = {"role": "assistant", "content": ans["answer"],
+                            "id": message_id, "prompt": ans.get("prompt", "")}
+        ans["id"] = message_id
+
+    def stream():
+        nonlocal dia, msg, req, conv
+        try:
+            for ans in chat(dia, msg, **req):
+                fillin_conv(ans)
+                yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans}, ensure_ascii=False) + "\n\n"
+            ConversationService.update_by_id(conv.id, conv.to_dict())
+        except Exception as e:
+            yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
+                                        "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
+                                       ensure_ascii=False) + "\n\n"
+        yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
+
+    if req.get("stream", True):
+        resp = Response(stream(), mimetype="text/event-stream")
+        resp.headers.add_header("Cache-control", "no-cache")
+        resp.headers.add_header("Connection", "keep-alive")
+        resp.headers.add_header("X-Accel-Buffering", "no")
+        resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
+        return resp
+
+    else:
+        answer = None
+        for ans in chat(dia, msg, **req):
+            answer = ans
+            fillin_conv(ans)
+            ConversationService.update_by_id(conv.id, conv.to_dict())
+            break
+        return get_json_result(data=answer)
+
+
+@manager.route('/get', methods=['GET'])
+@token_required
+def get(tenant_id):
+    req = request.args
+    if "id" not in req:
+        return get_data_error_result(retmsg="id is required")
+    conv_id = req["id"]
+    conv = ConversationService.query(id=conv_id)
+    if not conv:
+        return get_data_error_result(retmsg="Session does not exist")
+    if not DialogService.query(id=conv[0].dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_data_error_result(retmsg="You do not own the session")
+    conv = conv[0].to_dict()
+    conv['messages'] = conv.pop("message")
+    conv["assistant_id"] = conv.pop("dialog_id")
+    if conv["reference"]:
+        messages = conv["messages"]
+        message_num = 0
+        chunk_num = 0
+        while message_num < len(messages):
+            if message_num != 0 and messages[message_num]["role"] != "user":
+                chunk_list = []
+                if "chunks" in conv["reference"][chunk_num]:
+                    chunks = conv["reference"][chunk_num]["chunks"]
+                    for chunk in chunks:
+                        new_chunk = {
+                            "id": chunk["chunk_id"],
+                            "content": chunk["content_with_weight"],
+                            "document_id": chunk["doc_id"],
+                            "document_name": chunk["docnm_kwd"],
+                            "knowledgebase_id": chunk["kb_id"],
+                            "image_id": chunk["img_id"],
+                            "similarity": chunk["similarity"],
+                            "vector_similarity": chunk["vector_similarity"],
+                            "term_similarity": chunk["term_similarity"],
+                            "positions": chunk["positions"],
+                        }
+                        chunk_list.append(new_chunk)
+                chunk_num += 1
+                messages[message_num]["reference"] = chunk_list
+            message_num += 1
+    del conv["reference"]
+    return get_json_result(data=conv)
+
+
+@manager.route('/list', methods=["GET"])
+@token_required
+def list(tenant_id):
+    assistant_id = request.args["assistant_id"]
+    if not DialogService.query(tenant_id=tenant_id, id=assistant_id, status=StatusEnum.VALID.value):
+        return get_json_result(
+            data=False, retmsg=f'Only owner of the assistant is authorized for this operation.',
+            retcode=RetCode.OPERATING_ERROR)
+    convs = ConversationService.query(
+        dialog_id=assistant_id,
+        order_by=ConversationService.model.create_time,
+        reverse=True)
+    convs = [d.to_dict() for d in convs]
+    for conv in convs:
+        conv['messages'] = conv.pop("message")
+        conv["assistant_id"] = conv.pop("dialog_id")
+        if conv["reference"]:
+            messages = conv["messages"]
+            message_num = 0
+            chunk_num = 0
+            while message_num < len(messages):
+                if message_num != 0 and messages[message_num]["role"] != "user":
+                    chunk_list = []
+                    if "chunks" in conv["reference"][chunk_num]:
+                        chunks = conv["reference"][chunk_num]["chunks"]
+                        for chunk in chunks:
+                            new_chunk = {
+                                "id": chunk["chunk_id"],
+                                "content": chunk["content_with_weight"],
+                                "document_id": chunk["doc_id"],
+                                "document_name": chunk["docnm_kwd"],
+                                "knowledgebase_id": chunk["kb_id"],
+                                "image_id": chunk["img_id"],
+                                "similarity": chunk["similarity"],
+                                "vector_similarity": chunk["vector_similarity"],
+                                "term_similarity": chunk["term_similarity"],
+                                "positions": chunk["positions"],
+                            }
+                            chunk_list.append(new_chunk)
+                    chunk_num += 1
+                    messages[message_num]["reference"] = chunk_list
+                message_num += 1
+        del conv["reference"]
+    return get_json_result(data=convs)
+
+
+@manager.route('/delete', methods=["DELETE"])
+@token_required
+def delete(tenant_id):
+    id = request.args.get("id")
+    if not id:
+        return get_data_error_result(retmsg="`id` is required in deleting operation")
+    conv = ConversationService.query(id=id)
+    if not conv:
+        return get_data_error_result(retmsg="Session doesn't exist")
+    conv = conv[0]
+    if not DialogService.query(id=conv.dialog_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
+        return get_data_error_result(retmsg="You don't own the session")
+    ConversationService.delete_by_id(id)
+    return get_json_result(data=True)

api/apps/system_app.py

@@ -22,7 +22,7 @@ from api.utils.api_utils import get_json_result
 from api.versions import get_rag_version
 from rag.settings import SVR_QUEUE_NAME
 from rag.utils.es_conn import ELASTICSEARCH
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 from timeit import default_timer as timer
 
 from rag.utils.redis_conn import REDIS_CONN
@@ -47,7 +47,7 @@ def status():
 
     st = timer()
     try:
-        MINIO.health()
+        STORAGE_IMPL.health()
         res["minio"] = {"status": "green", "elapsed": "{:.1f}".format((timer() - st)*1000.)}
     except Exception as e:
         res["minio"] = {"status": "red", "elapsed": "{:.1f}".format((timer() - st)*1000.), "error": str(e)}

api/apps/tenant_app.py

@@ -0,0 +1,85 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+
+from flask import request
+from flask_login import current_user, login_required
+
+from api.db import UserTenantRole, StatusEnum
+from api.db.db_models import UserTenant
+from api.db.services.user_service import TenantService, UserTenantService
+from api.settings import RetCode
+
+from api.utils import get_uuid
+from api.utils.api_utils import get_json_result, validate_request, server_error_response
+
+
+@manager.route("/list", methods=["GET"])
+@login_required
+def tenant_list():
+    try:
+        tenants = TenantService.get_by_user_id(current_user.id)
+        return get_json_result(data=tenants)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route("/<tenant_id>/user/list", methods=["GET"])
+@login_required
+def user_list(tenant_id):
+    try:
+        users = UserTenantService.get_by_tenant_id(tenant_id)
+        return get_json_result(data=users)
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/<tenant_id>/user', methods=['POST'])
+@login_required
+@validate_request("user_id")
+def create(tenant_id):
+    user_id = request.json.get("user_id")
+    if not user_id:
+        return get_json_result(
+            data=False, retmsg='Lack of "USER ID"', retcode=RetCode.ARGUMENT_ERROR)
+
+    try:
+        user_tenants = UserTenantService.query(user_id=user_id, tenant_id=tenant_id)
+        if user_tenants:
+            uuid = user_tenants[0].id
+            return get_json_result(data={"id": uuid})
+
+        uuid = get_uuid()
+        UserTenantService.save(
+            id = uuid,
+            user_id = user_id,
+            tenant_id = tenant_id,
+            role = UserTenantRole.NORMAL.value,
+            status = StatusEnum.VALID.value)
+
+        return get_json_result(data={"id": uuid})
+    except Exception as e:
+        return server_error_response(e)
+
+
+@manager.route('/<tenant_id>/user/<user_id>', methods=['DELETE'])
+@login_required
+def rm(tenant_id, user_id):
+    try:
+        UserTenantService.filter_delete([UserTenant.tenant_id == tenant_id, UserTenant.user_id == user_id])
+        return get_json_result(data=True)
+    except Exception as e:
+        return server_error_response(e)
+    

api/db/__init__.py

@@ -55,6 +55,7 @@ class LLMType(StrEnum):
     SPEECH2TEXT = 'speech2text'
     IMAGE2TEXT = 'image2text'
     RERANK = 'rerank'
+    TTS    = 'tts'
 
 
 class ChatStyle(StrEnum):

api/db/db_models.py

@@ -18,18 +18,19 @@ import os
 import sys
 import typing
 import operator
+from enum import Enum
 from functools import wraps
 from itsdangerous.url_safe import URLSafeTimedSerializer as Serializer
 from flask_login import UserMixin
-from playhouse.migrate import MySQLMigrator, migrate
+from playhouse.migrate import MySQLMigrator, PostgresqlMigrator, migrate
 from peewee import (
     BigIntegerField, BooleanField, CharField,
     CompositeKey, IntegerField, TextField, FloatField, DateTimeField,
     Field, Model, Metadata
 )
-from playhouse.pool import PooledMySQLDatabase
+from playhouse.pool import PooledMySQLDatabase, PooledPostgresqlDatabase
 from api.db import SerializedType, ParserType
-from api.settings import DATABASE, stat_logger, SECRET_KEY
+from api.settings import DATABASE, stat_logger, SECRET_KEY, DATABASE_TYPE
 from api.utils.log_utils import getLogger
 from api import utils
 
@@ -58,8 +59,13 @@ AUTO_DATE_TIMESTAMP_FIELD_PREFIX = {
     "write_access"}
 
 
+class TextFieldType(Enum):
+    MYSQL = 'LONGTEXT'
+    POSTGRES = 'TEXT'
+
+
 class LongTextField(TextField):
-    field_type = 'LONGTEXT'
+    field_type = TextFieldType[DATABASE_TYPE.upper()].value
 
 
 class JSONField(LongTextField):
@@ -266,18 +272,69 @@ class JsonSerializedField(SerializedField):
         super(JsonSerializedField, self).__init__(serialized_type=SerializedType.JSON, object_hook=object_hook,
                                                   object_pairs_hook=object_pairs_hook, **kwargs)
 
+class PooledDatabase(Enum):
+    MYSQL = PooledMySQLDatabase
+    POSTGRES = PooledPostgresqlDatabase
+
+
+class DatabaseMigrator(Enum):
+    MYSQL = MySQLMigrator
+    POSTGRES = PostgresqlMigrator
+
 
 @singleton
 class BaseDataBase:
     def __init__(self):
         database_config = DATABASE.copy()
         db_name = database_config.pop("name")
-        self.database_connection = PooledMySQLDatabase(
+        self.database_connection = PooledDatabase[DATABASE_TYPE.upper()].value(db_name, **database_config)
-            db_name, **database_config)
+        stat_logger.info('init database on cluster mode successfully')
-        stat_logger.info('init mysql database on cluster mode successfully')
 
+class PostgresDatabaseLock:
+    def __init__(self, lock_name, timeout=10, db=None):
+        self.lock_name = lock_name
+        self.timeout = int(timeout)
+        self.db = db if db else DB
 
-class DatabaseLock:
+    def lock(self):
+        cursor = self.db.execute_sql("SELECT pg_try_advisory_lock(%s)", self.timeout)
+        ret = cursor.fetchone()
+        if ret[0] == 0:
+            raise Exception(f'acquire postgres lock {self.lock_name} timeout')
+        elif ret[0] == 1:
+            return True
+        else:
+            raise Exception(f'failed to acquire lock {self.lock_name}')
+
+    def unlock(self):
+        cursor = self.db.execute_sql("SELECT pg_advisory_unlock(%s)", self.timeout)
+        ret = cursor.fetchone()
+        if ret[0] == 0:
+            raise Exception(
+                f'postgres lock {self.lock_name} was not established by this thread')
+        elif ret[0] == 1:
+            return True
+        else:
+            raise Exception(f'postgres lock {self.lock_name} does not exist')
+
+    def __enter__(self):
+        if isinstance(self.db, PostgresDatabaseLock):
+            self.lock()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if isinstance(self.db, PostgresDatabaseLock):
+            self.unlock()
+
+    def __call__(self, func):
+        @wraps(func)
+        def magic(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+
+        return magic
+
+class MysqlDatabaseLock:
     def __init__(self, lock_name, timeout=10, db=None):
         self.lock_name = lock_name
         self.timeout = int(timeout)
@@ -325,8 +382,13 @@ class DatabaseLock:
         return magic
 
 
+class DatabaseLock(Enum):
+    MYSQL = MysqlDatabaseLock
+    POSTGRES = PostgresDatabaseLock
+
+
 DB = BaseDataBase().database_connection
-DB.lock = DatabaseLock
+DB.lock = DatabaseLock[DATABASE_TYPE.upper()].value
 
 
 def close_connection():
@@ -449,6 +511,11 @@ class Tenant(DataBaseModel):
         null=False,
         help_text="default rerank model ID",
         index=True)
+    tts_id = CharField(
+        max_length=256,
+        null=True,
+        help_text="default tts model ID",
+        index=True)
     parser_ids = CharField(
         max_length=256,
         null=False,
@@ -783,6 +850,7 @@ class Task(DataBaseModel):
         null=True,
         help_text="process message",
         default="")
+    retry_count = IntegerField(default=0)
 
 
 class Dialog(DataBaseModel):
@@ -824,6 +892,7 @@ class Dialog(DataBaseModel):
     do_refer = CharField(
         max_length=1,
         null=False,
+        default="1",
         help_text="it needs to insert reference index into answer or not")
     
     rerank_id = CharField(
@@ -911,7 +980,7 @@ class CanvasTemplate(DataBaseModel):
 
 def migrate_db():
     with DB.transaction():
-        migrator = MySQLMigrator(DB)
+        migrator = DatabaseMigrator[DATABASE_TYPE.upper()].value(DB)
         try:
             migrate(
                 migrator.add_column('file', 'source_type', CharField(max_length=128, null=False, default="",
@@ -958,6 +1027,13 @@ def migrate_db():
             )
         except Exception as e:
             pass
+        try:
+            migrate(
+                migrator.add_column("tenant","tts_id",
+                    CharField(max_length=256,null=True,help_text="default tts model ID",index=True))
+            )
+        except Exception as e:
+            pass
         try:
             migrate(
                 migrator.add_column('api_4_conversation', 'source',
@@ -970,3 +1046,10 @@ def migrate_db():
             DB.execute_sql('ALTER TABLE llm ADD PRIMARY KEY (llm_name,fid);')
         except Exception as e:
             pass
+        try:
+            migrate(
+                migrator.add_column('task', 'retry_count', IntegerField(default=0))
+            )
+        except Exception as e:
+            pass
+

api/db/db_utils.py

@@ -17,6 +17,8 @@ import operator
 from functools import reduce
 from typing import Dict, Type, Union
 
+from playhouse.pool import PooledMySQLDatabase
+
 from api.utils import current_timestamp, timestamp_to_date
 
 from api.db.db_models import DB, DataBaseModel
@@ -49,7 +51,10 @@ def bulk_insert_into_db(model, data_source, replace_on_conflict=False):
         with DB.atomic():
             query = model.insert_many(data_source[i:i + batch_size])
             if replace_on_conflict:
-                query = query.on_conflict(preserve=preserve)
+                if isinstance(DB, PooledMySQLDatabase):
+                    query = query.on_conflict(preserve=preserve)
+                else:
+                    query = query.on_conflict(conflict_target="id", preserve=preserve)
             query.execute()
 
 

api/db/services/api_service.py

@@ -14,7 +14,9 @@
 #  limitations under the License.
 #
 from datetime import datetime
+
 import peewee
+
 from api.db.db_models import DB, API4Conversation, APIToken, Dialog
 from api.db.services.common_service import CommonService
 from api.utils import current_timestamp, datetime_format
@@ -41,7 +43,7 @@ class API4ConversationService(CommonService):
     @DB.connection_context()
     def append_message(cls, id, conversation):
         cls.update_by_id(id, conversation)
-        return cls.model.update(round=cls.model.round + 1).where(cls.model.id==id).execute()
+        return cls.model.update(round=cls.model.round + 1).where(cls.model.id == id).execute()
 
     @classmethod
     @DB.connection_context()
@@ -61,7 +63,7 @@ class API4ConversationService(CommonService):
                 cls.model.round).alias("round"),
             peewee.fn.SUM(
                 cls.model.thumb_up).alias("thumb_up")
-        ).join(Dialog, on=(cls.model.dialog_id == Dialog.id & Dialog.tenant_id == tenant_id)).where(
+        ).join(Dialog, on=((cls.model.dialog_id == Dialog.id) & (Dialog.tenant_id == tenant_id))).where(
             cls.model.create_date >= from_date,
             cls.model.create_date <= to_date,
             cls.model.source == source

api/db/services/dialog_service.py

@@ -13,11 +13,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import binascii
 import os
 import json
 import re
 from copy import deepcopy
-
+from timeit import default_timer as timer
 from api.db import LLMType, ParserType
 from api.db.db_models import Dialog, Conversation
 from api.db.services.common_service import CommonService
@@ -87,6 +88,7 @@ def llm_id2llm_type(llm_id):
 
 def chat(dialog, messages, stream=True, **kwargs):
     assert messages[-1]["role"] == "user", "The last content of this conversation is not from user."
+    st = timer()
     llm = LLMService.query(llm_name=dialog.llm_id)
     if not llm:
         llm = TenantLLMService.query(tenant_id=dialog.tenant_id, llm_name=dialog.llm_id)
@@ -120,6 +122,9 @@ def chat(dialog, messages, stream=True, **kwargs):
 
     prompt_config = dialog.prompt_config
     field_map = KnowledgebaseService.get_field_map(dialog.kb_ids)
+    tts_mdl = None
+    if prompt_config.get("tts"):
+        tts_mdl = LLMBundle(dialog.tenant_id, LLMType.TTS)
     # try to use sql if field mapping is good to go
     if field_map:
         chat_logger.info("Use SQL to retrieval:{}".format(questions[-1]))
@@ -154,24 +159,16 @@ def chat(dialog, messages, stream=True, **kwargs):
                                         doc_ids=attachments,
                                         top=dialog.top_k, aggs=False, rerank_mdl=rerank_mdl)
     knowledges = [ck["content_with_weight"] for ck in kbinfos["chunks"]]
-    #self-rag
-    if dialog.prompt_config.get("self_rag") and not relevant(dialog.tenant_id, dialog.llm_id, questions[-1], knowledges):
-        questions[-1] = rewrite(dialog.tenant_id, dialog.llm_id, questions[-1])
-        kbinfos = retr.retrieval(" ".join(questions), embd_mdl, dialog.tenant_id, dialog.kb_ids, 1, dialog.top_n,
-                                        dialog.similarity_threshold,
-                                        dialog.vector_similarity_weight,
-                                        doc_ids=attachments,
-                                        top=dialog.top_k, aggs=False, rerank_mdl=rerank_mdl)
-        knowledges = [ck["content_with_weight"] for ck in kbinfos["chunks"]]
-
     chat_logger.info(
         "{}->{}".format(" ".join(questions), "\n->".join(knowledges)))
+    retrieval_tm = timer()
 
     if not knowledges and prompt_config.get("empty_response"):
-        yield {"answer": prompt_config["empty_response"], "reference": kbinfos}
+        empty_res = prompt_config["empty_response"]
+        yield {"answer": empty_res, "reference": kbinfos, "audio_binary": tts(tts_mdl, empty_res)}
         return {"answer": prompt_config["empty_response"], "reference": kbinfos}
 
-    kwargs["knowledge"] = "\n".join(knowledges)
+    kwargs["knowledge"] = "\n------\n".join(knowledges)
     gen_conf = dialog.llm_setting
 
     msg = [{"role": "system", "content": prompt_config["system"].format(**kwargs)}]
@@ -179,6 +176,7 @@ def chat(dialog, messages, stream=True, **kwargs):
                 for m in messages if m["role"] != "system"])
     used_token_count, msg = message_fit_in(msg, int(max_tokens * 0.97))
     assert len(msg) >= 2, f"message_fit_in has bug: {msg}"
+    prompt = msg[0]["content"]
 
     if "max_tokens" in gen_conf:
         gen_conf["max_tokens"] = min(
@@ -186,7 +184,7 @@ def chat(dialog, messages, stream=True, **kwargs):
             max_tokens - used_token_count)
 
     def decorate_answer(answer):
-        nonlocal prompt_config, knowledges, kwargs, kbinfos
+        nonlocal prompt_config, knowledges, kwargs, kbinfos, prompt, retrieval_tm
         refs = []
         if knowledges and (prompt_config.get("quote", True) and kwargs.get("quote", True)):
             answer, idx = retr.insert_citations(answer,
@@ -210,20 +208,31 @@ def chat(dialog, messages, stream=True, **kwargs):
 
         if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0:
             answer += " Please set LLM API-Key in 'User Setting -> Model Providers -> API-Key'"
-        return {"answer": answer, "reference": refs}
+        done_tm = timer()
+        prompt += "\n### Elapsed\n  - Retrieval: %.1f ms\n  - LLM: %.1f ms"%((retrieval_tm-st)*1000, (done_tm-st)*1000)
+        return {"answer": answer, "reference": refs, "prompt": prompt}
 
     if stream:
+        last_ans = ""
         answer = ""
-        for ans in chat_mdl.chat_streamly(msg[0]["content"], msg[1:], gen_conf):
+        for ans in chat_mdl.chat_streamly(prompt, msg[1:], gen_conf):
             answer = ans
-            yield {"answer": answer, "reference": {}}
+            delta_ans = ans[len(last_ans):]
+            if num_tokens_from_string(delta_ans) < 16:
+                continue
+            last_ans = answer
+            yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
+        delta_ans = answer[len(last_ans):]
+        if delta_ans:
+            yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
         yield decorate_answer(answer)
     else:
-        answer = chat_mdl.chat(
+        answer = chat_mdl.chat(prompt, msg[1:], gen_conf)
-            msg[0]["content"], msg[1:], gen_conf)
         chat_logger.info("User: {}|Assistant: {}".format(
             msg[-1]["content"], answer))
-        yield decorate_answer(answer)
+        res = decorate_answer(answer)
+        res["audio_binary"] = tts(tts_mdl, answer)
+        yield res
 
 
 def use_sql(question, field_map, tenant_id, chat_mdl, quota=True):
@@ -334,7 +343,8 @@ def use_sql(question, field_map, tenant_id, chat_mdl, quota=True):
         chat_logger.warning("SQL missing field: " + sql)
         return {
             "answer": "\n".join([clmns, line, rows]),
-            "reference": {"chunks": [], "doc_aggs": []}
+            "reference": {"chunks": [], "doc_aggs": []},
+            "prompt": sys_prompt
         }
 
     docid_idx = list(docid_idx)[0]
@@ -348,7 +358,8 @@ def use_sql(question, field_map, tenant_id, chat_mdl, quota=True):
         "answer": "\n".join([clmns, line, rows]),
         "reference": {"chunks": [{"doc_id": r[docid_idx], "docnm_kwd": r[docnm_idx]} for r in tbl["rows"]],
                       "doc_aggs": [{"doc_id": did, "doc_name": d["doc_name"], "count": d["count"]} for did, d in
-                                   doc_aggs.items()]}
+                                   doc_aggs.items()]},
+        "prompt": sys_prompt
     }
 
 
@@ -390,3 +401,82 @@ def rewrite(tenant_id, llm_id, question):
     """
     ans = chat_mdl.chat(prompt, [{"role": "user", "content": question}], {"temperature": 0.8})
     return ans
+
+
+def tts(tts_mdl, text):
+    if not tts_mdl or not text: return
+    bin = b""
+    for chunk in tts_mdl.tts(text):
+        bin += chunk
+    return binascii.hexlify(bin).decode("utf-8")
+
+
+def ask(question, kb_ids, tenant_id):
+    kbs = KnowledgebaseService.get_by_ids(kb_ids)
+    embd_nms = list(set([kb.embd_id for kb in kbs]))
+
+    is_kg = all([kb.parser_id == ParserType.KG for kb in kbs])
+    retr = retrievaler if not is_kg else kg_retrievaler
+
+    embd_mdl = LLMBundle(tenant_id, LLMType.EMBEDDING, embd_nms[0])
+    chat_mdl = LLMBundle(tenant_id, LLMType.CHAT)
+    max_tokens = chat_mdl.max_length
+
+    kbinfos = retr.retrieval(question, embd_mdl, tenant_id, kb_ids, 1, 12, 0.1, 0.3, aggs=False)
+    knowledges = [ck["content_with_weight"] for ck in kbinfos["chunks"]]
+
+    used_token_count = 0
+    for i, c in enumerate(knowledges):
+        used_token_count += num_tokens_from_string(c)
+        if max_tokens * 0.97 < used_token_count:
+            knowledges = knowledges[:i]
+            break
+
+    prompt = """
+    Role: You're a smart assistant. Your name is Miss R.
+    Task: Summarize the information from knowledge bases and answer user's question.
+    Requirements and restriction:
+      - DO NOT make things up, especially for numbers.
+      - If the information from knowledge is irrelevant with user's question, JUST SAY: Sorry, no relevant information provided.
+      - Answer with markdown format text.
+      - Answer in language of user's question.
+      - DO NOT make things up, especially for numbers.
+      
+    ### Information from knowledge bases
+    %s
+    
+    The above is information from knowledge bases.
+     
+    """%"\n".join(knowledges)
+    msg = [{"role": "user", "content": question}]
+
+    def decorate_answer(answer):
+        nonlocal knowledges, kbinfos, prompt
+        answer, idx = retr.insert_citations(answer,
+                                           [ck["content_ltks"]
+                                            for ck in kbinfos["chunks"]],
+                                           [ck["vector"]
+                                            for ck in kbinfos["chunks"]],
+                                           embd_mdl,
+                                           tkweight=0.7,
+                                           vtweight=0.3)
+        idx = set([kbinfos["chunks"][int(i)]["doc_id"] for i in idx])
+        recall_docs = [
+            d for d in kbinfos["doc_aggs"] if d["doc_id"] in idx]
+        if not recall_docs: recall_docs = kbinfos["doc_aggs"]
+        kbinfos["doc_aggs"] = recall_docs
+        refs = deepcopy(kbinfos)
+        for c in refs["chunks"]:
+            if c.get("vector"):
+                del c["vector"]
+
+        if answer.lower().find("invalid key") >= 0 or answer.lower().find("invalid api") >= 0:
+            answer += " Please set LLM API-Key in 'User Setting -> Model Providers -> API-Key'"
+        return {"answer": answer, "reference": refs}
+
+    answer = ""
+    for ans in chat_mdl.chat_streamly(prompt, msg, {"temperature": 0.1}):
+        answer = ans
+        yield {"answer": answer, "reference": {}}
+    yield decorate_answer(answer)
+

api/db/services/document_service.py

@@ -34,7 +34,7 @@ from api.utils.file_utils import get_project_base_directory
 from graphrag.mind_map_extractor import MindMapExtractor
 from rag.settings import SVR_QUEUE_NAME
 from rag.utils.es_conn import ELASTICSEARCH
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 from rag.nlp import search, rag_tokenizer
 
 from api.db import FileType, TaskStatus, ParserType, LLMType
@@ -473,7 +473,7 @@ def doc_upload_and_parse(conversation_id, file_objs, user_id):
             else:
                 d["image"].save(output_buffer, format='JPEG')
 
-            MINIO.put(kb.id, d["_id"], output_buffer.getvalue())
+            STORAGE_IMPL.put(kb.id, d["_id"], output_buffer.getvalue())
             d["img_id"] = "{}-{}".format(kb.id, d["_id"])
             del d["image"]
             docs.append(d)

api/db/services/file2document_service.py

@@ -76,7 +76,7 @@ class File2DocumentService(CommonService):
             f2d = cls.get_by_file_id(file_id)
         if f2d:
             file = File.get_by_id(f2d[0].file_id)
-            if file.source_type == FileSource.LOCAL:
+            if not file.source_type or file.source_type == FileSource.LOCAL:
                 return file.parent_id, file.location
             doc_id = f2d[0].document_id
 

api/db/services/file_service.py

@@ -27,7 +27,7 @@ from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
 from api.utils import get_uuid
 from api.utils.file_utils import filename_type, thumbnail
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 
 
 class FileService(CommonService):
@@ -350,10 +350,10 @@ class FileService(CommonService):
                     raise RuntimeError("This type of file has not been supported yet!")
 
                 location = filename
-                while MINIO.obj_exist(kb.id, location):
+                while STORAGE_IMPL.obj_exist(kb.id, location):
                     location += "_"
                 blob = file.read()
-                MINIO.put(kb.id, location, blob)
+                STORAGE_IMPL.put(kb.id, location, blob)
                 doc = {
                     "id": get_uuid(),
                     "kb_id": kb.id,

api/db/services/llm_service.py

@@ -15,7 +15,7 @@
 #
 from api.db.services.user_service import TenantService
 from api.settings import database_logger
-from rag.llm import EmbeddingModel, CvModel, ChatModel, RerankModel, Seq2txtModel
+from rag.llm import EmbeddingModel, CvModel, ChatModel, RerankModel, Seq2txtModel, TTSModel
 from api.db import LLMType
 from api.db.db_models import DB, UserTenant
 from api.db.db_models import LLMFactories, LLM, TenantLLM
@@ -75,6 +75,8 @@ class TenantLLMService(CommonService):
             mdlnm = tenant.llm_id if not llm_name else llm_name
         elif llm_type == LLMType.RERANK:
             mdlnm = tenant.rerank_id if not llm_name else llm_name
+        elif llm_type == LLMType.TTS:
+            mdlnm = tenant.tts_id if not llm_name else llm_name
         else:
             assert False, "LLM type error"
 
@@ -127,6 +129,14 @@ class TenantLLMService(CommonService):
                 model_config["api_key"], model_config["llm_name"], lang,
                 base_url=model_config["api_base"]
             )
+        if llm_type == LLMType.TTS:
+            if model_config["llm_factory"] not in TTSModel:
+                return
+            return TTSModel[model_config["llm_factory"]](
+                model_config["api_key"],
+                model_config["llm_name"],
+                base_url=model_config["api_base"],
+            )
 
     @classmethod
     @DB.connection_context()
@@ -144,7 +154,9 @@ class TenantLLMService(CommonService):
         elif llm_type == LLMType.CHAT.value:
             mdlnm = tenant.llm_id if not llm_name else llm_name
         elif llm_type == LLMType.RERANK:
-            mdlnm = tenant.llm_id if not llm_name else llm_name
+            mdlnm = tenant.rerank_id if not llm_name else llm_name
+        elif llm_type == LLMType.TTS:
+            mdlnm = tenant.tts_id if not llm_name else llm_name
         else:
             assert False, "LLM type error"
 
@@ -178,7 +190,7 @@ class LLMBundle(object):
             tenant_id, llm_type, llm_name, lang=lang)
         assert self.mdl, "Can't find mole for {}/{}/{}".format(
             tenant_id, llm_type, llm_name)
-        self.max_length = 512
+        self.max_length = 8192
         for lm in LLMService.query(llm_name=llm_name):
             self.max_length = lm.max_tokens
             break
@@ -223,6 +235,17 @@ class LLMBundle(object):
                 "Can't update token usage for {}/SEQUENCE2TXT".format(self.tenant_id))
         return txt
 
+    def tts(self, text):
+        for chunk in self.mdl.tts(text):
+            if isinstance(chunk,int):
+                if not TenantLLMService.increase_usage(
+                    self.tenant_id, self.llm_type, chunk, self.llm_name):
+                        database_logger.error(
+                            "Can't update token usage for {}/TTS".format(self.tenant_id))
+                return
+            yield chunk     
+
+    
     def chat(self, system, history, gen_conf):
         txt, used_tokens = self.mdl.chat(system, history, gen_conf)
         if not TenantLLMService.increase_usage(

api/db/services/task_service.py

@@ -27,7 +27,7 @@ from api.db.services.document_service import DocumentService
 from api.utils import current_timestamp, get_uuid
 from deepdoc.parser.excel_parser import RAGFlowExcelParser
 from rag.settings import SVR_QUEUE_NAME
-from rag.utils.minio_conn import MINIO
+from rag.utils.storage_factory import STORAGE_IMPL
 from rag.utils.redis_conn import REDIS_CONN
 
 
@@ -42,6 +42,7 @@ class TaskService(CommonService):
             cls.model.doc_id,
             cls.model.from_page,
             cls.model.to_page,
+            cls.model.retry_count,
             Document.kb_id,
             Document.parser_id,
             Document.parser_config,
@@ -64,9 +65,20 @@ class TaskService(CommonService):
         docs = list(docs.dicts())
         if not docs: return []
 
-        cls.model.update(progress_msg=cls.model.progress_msg + "\n" + "Task has been received.",
+        msg = "\nTask has been received."
-                         progress=random.random() / 10.).where(
+        prog = random.random() / 10.
+        if docs[0]["retry_count"] >= 3:
+            msg = "\nERROR: Task is abandoned after 3 times attempts."
+            prog = -1
+
+        cls.model.update(progress_msg=cls.model.progress_msg + msg,
+                         progress=prog,
+                         retry_count=docs[0]["retry_count"]+1
+                         ).where(
             cls.model.id == docs[0]["id"]).execute()
+
+        if docs[0]["retry_count"] >= 3: return []
+
         return docs
 
     @classmethod
@@ -131,7 +143,7 @@ def queue_tasks(doc, bucket, name):
     tsks = []
 
     if doc["type"] == FileType.PDF.value:
-        file_bin = MINIO.get(bucket, name)
+        file_bin = STORAGE_IMPL.get(bucket, name)
         do_layout = doc["parser_config"].get("layout_recognize", True)
         pages = PdfParser.total_page_number(doc["name"], file_bin)
         page_size = doc["parser_config"].get("task_page_size", 12)
@@ -157,7 +169,7 @@ def queue_tasks(doc, bucket, name):
                 tsks.append(task)
 
     elif doc["parser_id"] == "table":
-        file_bin = MINIO.get(bucket, name)
+        file_bin = STORAGE_IMPL.get(bucket, name)
         rn = RAGFlowExcelParser.row_number(
             doc["name"], file_bin)
         for i in range(0, rn, 3000):

api/db/services/user_service.py

@@ -96,6 +96,7 @@ class TenantService(CommonService):
             cls.model.rerank_id,
             cls.model.asr_id,
             cls.model.img2txt_id,
+            cls.model.tts_id,
             cls.model.parser_ids,
             UserTenant.role]
         return list(cls.model.select(*fields)
@@ -136,3 +137,24 @@ class UserTenantService(CommonService):
             kwargs["id"] = get_uuid()
         obj = cls.model(**kwargs).save(force_insert=True)
         return obj
+
+    @classmethod
+    @DB.connection_context()
+    def get_by_tenant_id(cls, tenant_id):
+        fields = [
+            cls.model.user_id,
+            cls.model.tenant_id,
+            cls.model.role,
+            cls.model.status,
+            User.nickname,
+            User.email,
+            User.avatar,
+            User.is_authenticated,
+            User.is_active,
+            User.is_anonymous,
+            User.status,
+            User.is_superuser]
+        return list(cls.model.select(*fields)
+                    .join(User, on=((cls.model.user_id == User.id) & (cls.model.status == StatusEnum.VALID.value)))
+                    .where(cls.model.tenant_id == tenant_id)
+                    .dicts())

api/settings.py

@@ -164,7 +164,8 @@ RANDOM_INSTANCE_ID = get_base_config(
 PROXY = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("proxy")
 PROXY_PROTOCOL = get_base_config(RAG_FLOW_SERVICE_NAME, {}).get("protocol")
 
-DATABASE = decrypt_database_config(name="mysql")
+DATABASE_TYPE = os.getenv("DB_TYPE", 'mysql')
+DATABASE = decrypt_database_config(name=DATABASE_TYPE)
 
 # Switch
 # upload

api/utils/api_utils.py

@@ -13,30 +13,32 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import functools
 import json
 import random
 import time
+from base64 import b64encode
 from functools import wraps
+from hmac import HMAC
 from io import BytesIO
+from urllib.parse import quote, urlencode
+from uuid import uuid1
+
+import requests
 from flask import (
     Response, jsonify, send_file, make_response,
     request as flask_request,
 )
 from werkzeug.http import HTTP_STATUS_CODES
 
-from api.utils import json_dumps
+from api.db.db_models import APIToken
-from api.settings import RetCode
 from api.settings import (
     REQUEST_MAX_WAIT_SEC, REQUEST_WAIT_SEC,
     stat_logger, CLIENT_AUTHENTICATION, HTTP_APP_KEY, SECRET_KEY
 )
-import requests
+from api.settings import RetCode
-import functools
 from api.utils import CustomJSONEncoder
-from uuid import uuid1
+from api.utils import json_dumps
-from base64 import b64encode
-from hmac import HMAC
-from urllib.parse import quote, urlencode
 
 requests.models.complexjson.dumps = functools.partial(
     json.dumps, cls=CustomJSONEncoder)
@@ -96,7 +98,6 @@ def get_exponential_backoff_interval(retries, full_jitter=False):
 
 def get_json_result(retcode=RetCode.SUCCESS, retmsg='success',
                     data=None, job_id=None, meta=None):
-    import re
     result_dict = {
         "retcode": retcode,
         "retmsg": retmsg,
@@ -145,7 +146,8 @@ def server_error_response(e):
         return get_json_result(
             retcode=RetCode.EXCEPTION_ERROR, retmsg=repr(e.args[0]), data=e.args[1])
     if repr(e).find("index_not_found_exception") >= 0:
-        return get_json_result(retcode=RetCode.EXCEPTION_ERROR, retmsg="No chunk found, please upload file and parse it.")
+        return get_json_result(retcode=RetCode.EXCEPTION_ERROR,
+                               retmsg="No chunk found, please upload file and parse it.")
 
     return get_json_result(retcode=RetCode.EXCEPTION_ERROR, retmsg=repr(e))
 
@@ -190,7 +192,9 @@ def validate_request(*args, **kwargs):
                 return get_json_result(
                     retcode=RetCode.ARGUMENT_ERROR, retmsg=error_string)
             return func(*_args, **_kwargs)
+
         return decorated_function
+
     return wrapper
 
 
@@ -217,7 +221,7 @@ def get_json_result(retcode=RetCode.SUCCESS, retmsg='success', data=None):
 
 
 def construct_response(retcode=RetCode.SUCCESS,
-                 retmsg='success', data=None, auth=None):
+                       retmsg='success', data=None, auth=None):
     result_dict = {"retcode": retcode, "retmsg": retmsg, "data": data}
     response_dict = {}
     for key, value in result_dict.items():
@@ -235,6 +239,7 @@ def construct_response(retcode=RetCode.SUCCESS,
     response.headers["Access-Control-Expose-Headers"] = "Authorization"
     return response
 
+
 def construct_result(code=RetCode.DATA_ERROR, message='data is missing'):
     import re
     result_dict = {"code": code, "message": re.sub(r"rag", "seceum", message, flags=re.IGNORECASE)}
@@ -263,7 +268,23 @@ def construct_error_response(e):
         pass
     if len(e.args) > 1:
         return construct_json_result(code=RetCode.EXCEPTION_ERROR, message=repr(e.args[0]), data=e.args[1])
-    if repr(e).find("index_not_found_exception") >=0:
+    if repr(e).find("index_not_found_exception") >= 0:
-        return construct_json_result(code=RetCode.EXCEPTION_ERROR, message="No chunk found, please upload file and parse it.")
+        return construct_json_result(code=RetCode.EXCEPTION_ERROR,
+                                     message="No chunk found, please upload file and parse it.")
 
     return construct_json_result(code=RetCode.EXCEPTION_ERROR, message=repr(e))
+
+
+def token_required(func):
+    @wraps(func)
+    def decorated_function(*args, **kwargs):
+        token = flask_request.headers.get('Authorization').split()[1]
+        objs = APIToken.query(token=token)
+        if not objs:
+            return get_json_result(
+                data=False, retmsg='Token is not valid!', retcode=RetCode.AUTHENTICATION_ERROR
+            )
+        kwargs['tenant_id'] = objs[0].tenant_id
+        return func(*args, **kwargs)
+
+    return decorated_function

conf/service_conf.yaml

@@ -9,10 +9,34 @@ mysql:
   port: 3306
   max_connections: 100
   stale_timeout: 30
+postgres:
+  name: 'rag_flow'
+  user: 'rag_flow'
+  password: 'infini_rag_flow'
+  host: 'postgres'
+  port: 5432
+  max_connections: 100
+  stale_timeout: 30
 minio:
   user: 'rag_flow'
   password: 'infini_rag_flow'
   host: 'minio:9000'
+azure:
+  auth_type: 'sas'
+  container_url: 'container_url'
+  sas_token: 'sas_token'
+#azure:
+#  auth_type: 'spn'
+#  account_url: 'account_url'
+#  client_id: 'client_id'
+#  secret: 'secret'
+#  tenant_id: 'tenant_id'
+#  container_name: 'container_name'
+s3:
+  endpoint: 'endpoint'
+  access_key: 'access_key'
+  secret_key: 'secret_key'
+  region: 'region'
 es:
   hosts: 'http://es01:9200'
   username: 'elastic'

deepdoc/parser/pdf_parser.py

@@ -299,7 +299,7 @@ class RAGFlowPdfParser:
                 self.lefted_chars.append(c)
                 continue
             if c["text"] == " " and bxs[ii]["text"]:
-                if re.match(r"[0-9a-zA-Z,.?;:!%%]", bxs[ii]["text"][-1]):
+                if re.match(r"[0-9a-zA-Zа-яА-Я,.?;:!%%]", bxs[ii]["text"][-1]):
                     bxs[ii]["text"] += " "
             else:
                 bxs[ii]["text"] += c["text"]

deepdoc/parser/txt_parser.py

@@ -33,14 +33,30 @@ class RAGFlowTxtParser:
     def parser_txt(cls, txt, chunk_token_num=128, delimiter="\n!?;。；！？"):
         if type(txt) != str:
             raise TypeError("txt type should be str!")
-        sections = []
+        cks = [""]
-        for sec in re.split(r"[%s]+"%delimiter, txt):
+        tk_nums = [0]
-            if sections and sec in delimiter:
+
-                sections[-1][0] += sec
+        def add_chunk(t):
-                continue
+            nonlocal cks, tk_nums, delimiter
-            if num_tokens_from_string(sec) > 10 * int(chunk_token_num):
+            tnum = num_tokens_from_string(t)
-                sections.append([sec[: int(len(sec) / 2)], ""])
+            if tnum < 8:
-                sections.append([sec[int(len(sec) / 2) :], ""])
+                pos = ""
+            if tk_nums[-1] > chunk_token_num:
+                cks.append(t)
+                tk_nums.append(tnum)
             else:
-                sections.append([sec, ""])
+                cks[-1] += t
-        return sections
+                tk_nums[-1] += tnum
+
+        s, e = 0, 1
+        while e < len(txt):
+            if txt[e] in delimiter:
+                add_chunk(txt[s: e + 1])
+                s = e + 1
+                e = s + 1
+            else:
+                e += 1
+        if s < e:
+            add_chunk(txt[s: e + 1])
+
+        return [[c,""] for c in cks]

docker/docker-compose-admin-tool.yml

@@ -1,3 +1,7 @@
+include:
+  - path: ./docker-compose.yml
+    env_file: ./.env
+
 services:
   kibana:
     image: kibana:${STACK_VERSION}

docker/entrypoint.sh

@@ -1,5 +1,8 @@
 #!/bin/bash
 
+# unset http proxy which maybe set by docker daemon
+export http_proxy=""; export https_proxy=""; export no_proxy=""; export HTTP_PROXY=""; export HTTPS_PROXY=""; export NO_PROXY=""
+
 /usr/sbin/nginx
 
 export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/
@@ -11,13 +14,13 @@ fi
 
 function task_exe(){
     while [ 1 -eq 1 ];do
-      $PY rag/svr/task_executor.py ;
+      $PY rag/svr/task_executor.py $1;
     done
 }
 
 for ((i=0;i<WS;i++))
 do
-  task_exe  &
+  task_exe  $i &
 done
 
 while [ 1 -eq 1 ];do

docker/init-kibana.sh

@@ -1,30 +1,67 @@
 #!/bin/bash
 
-# 等待 Elasticsearch 啟動
+# unset http proxy which maybe set by docker daemon
-until curl -u "elastic:${ELASTIC_PASSWORD}" -s http://es01:9200 >/dev/null; do
+export http_proxy=""; export https_proxy=""; export no_proxy=""; export HTTP_PROXY=""; export HTTPS_PROXY=""; export NO_PROXY=""
-  echo "等待 Elasticsearch 啟動..."
+
-  sleep 5
+echo "Elasticsearch built-in user: elastic:${ELASTIC_PASSWORD}"
+
+# Wait Elasticsearch be healthy
+while true; do
+    response=$(curl -s -v -w "\n%{http_code}" -u "elastic:${ELASTIC_PASSWORD}" "http://es01:9200")
+    exit_code=$?
+    status=$(echo "$response" | tail -n1)
+    if [ $exit_code -eq 0 ] && [ "$status" = "200" ]; then
+        echo "Elasticsearch is healthy"
+        break
+    else
+        echo "Elasticsearch is unhealthy: $exit_code $status"
+        echo "$response"
+        sleep 5
+    fi
 done
 
+# Create new role with all privileges to all indices
+# https://www.elastic.co/guide/en/elasticsearch/reference/current/security-privileges.html#privileges-list-indices
+echo "Going to create Elasticsearch role own_indices with all privileges to all indices"
+while true; do
+    response=$(curl -s -v -w "\n%{http_code}" -u "elastic:${ELASTIC_PASSWORD}" -X POST http://es01:9200/_security/role/own_indices -H 'Content-Type: application/json' -d '{"indices": [{"names": ["*"], "privileges": ["all"]}]}')
+    exit_code=$?
+    status=$(echo "$response" | tail -n1)
+    if [ $exit_code -eq 0 ] && [ "$status" = "200" ]; then
+        echo "Elasticsearch role own_indices created"
+        break
+    else
+        echo "Elasticsearch role own_indices failure: $exit_code $status"
+        echo "$response"
+        sleep 5
+    fi
+done
 
-echo "使用者: elastic:${ELASTIC_PASSWORD}"
+echo "Elasticsearch role own_indices:"
+curl -u "elastic:${ELASTIC_PASSWORD}" -X GET "http://es01:9200/_security/role/own_indices"
+echo ""
 
+PAYLOAD="{\"password\": \"${KIBANA_PASSWORD}\", \"roles\": [\"kibana_admin\", \"kibana_system\", \"own_indices\"], \"full_name\": \"${KIBANA_USER}\", \"email\": \"${KIBANA_USER}@example.com\"}"
 
+echo "Going to create Elasticsearch user ${KIBANA_USER}: ${PAYLOAD}"
 
-PAYLOAD="{
+# Create new user
-  \"password\" : \"${KIBANA_PASSWORD}\",
+while true; do
-  \"roles\" : [ \"kibana_admin\",\"kibana_system\" ],
+    response=$(curl -s -v -w "\n%{http_code}" -u "elastic:${ELASTIC_PASSWORD}" -X POST http://es01:9200/_security/user/${KIBANA_USER} -H "Content-Type: application/json" -d "${PAYLOAD}")
-  \"full_name\" : \"${KIBANA_USER}\",
+    exit_code=$?
-  \"email\" : \"${KIBANA_USER}@example.com\"
+    status=$(echo "$response" | tail -n1)
-}"
+    if [ $exit_code -eq 0 ] && [ "$status" = "200" ]; then
-echo "新用戶帳戶: $PAYLOAD"
+        echo "Elasticsearch user ${KIBANA_USER} created"
+        break
+    else
+        echo "Elasticsearch user ${KIBANA_USER} failure: $exit_code $status"
+        echo "$response"
+        sleep 5
+    fi
+done
 
-# 創建新用戶帳戶
+echo "Elasticsearch user ${KIBANA_USER}:"
-curl -X POST "http://es01:9200/_security/user/${KIBANA_USER}" \
+curl -u "elastic:${ELASTIC_PASSWORD}" -X GET "http://es01:9200/_security/user/${KIBANA_USER}"
--u "elastic:${ELASTIC_PASSWORD}" \
+echo ""
--H "Content-Type: application/json" \
--d "$PAYLOAD"s
-
-echo "新用戶帳戶已創建"
 
 exit 0

docker/nginx/ragflow.conf

@@ -11,10 +11,11 @@ server {
     gzip_disable "MSIE [1-6]\.";
 
     location /v1 {
-      proxy_pass http://ragflow:9380;
+        proxy_pass http://ragflow:9380;
-      include proxy.conf;
+        include proxy.conf;
     }
 
+
     location / {
         index index.html;
         try_files $uri $uri/ /index.html;

docs/guides/agentic_rag/_category_.json

@@ -0,0 +1,8 @@
+{
+  "label": "Agents",
+  "position": 3,
+  "link": {
+    "type": "generated-index",
+    "description": "RAGFlow v0.8.0 introduces an agent mechanism, featuring a no-code workflow editor on the front end and a comprehensive graph-based task orchestration framework on the back end."
+  }
+}

docs/guides/agentic_rag/agent_introduction.md

@@ -0,0 +1,78 @@
+---
+sidebar_position: 1
+slug: /agent_introduction
+---
+
+# Introduction to agents
+
+Agents and RAG are complementary techniques, each enhancing the other’s capabilities in business applications. RAGFlow v0.8.0 introduces an agent mechanism, featuring a no-code workflow editor on the front end and a comprehensive graph-based task orchestration framework on the back end. This mechanism is built on top of RAGFlow's existing RAG solutions and aims to orchestrate search technologies such as query intent classification, conversation leading, and query rewriting to:
+
+- Provide higher retrievals and,
+- Accommodate more complex scenarios.
+
+## Create an agent
+
+:::tip NOTE
+
+Before proceeding, ensure that:  
+
+1. You have properly set the LLM to use. See the guides on [Configure your API key](../llm_api_key_setup.md) or [Deploy a local LLM](../deploy_local_llm.mdx) for more information.
+2. You have a knowledge base configured and the corresponding files properly parsed. See the guide on [Configure a knowledge base](../configure_knowledge_base.md) for more information.
+
+:::
+
+Click the **Agent** tab in the middle top of the page to show the **Agent** page. As shown in the screenshot below, the cards on this page represent the created agents, which you can continue to edit.
+
+![agents](https://github.com/user-attachments/assets/5e10758b-ec43-49ae-bf91-ff7d04c56e9d)
+
+We also provide templates catered to different business scenarios. You can either generate your agent from one of our agent templates or create one from scratch:
+
+1. Click **+ Create agent** to show the **agent template** page:
+
+   ![agent_templates](https://github.com/user-attachments/assets/73bd476c-4bab-4c8c-82f8-6b00fb2cd044)
+
+2. To create an agent from scratch, click the **Blank** card. Alternatively, to create an agent from one of our templates, hover over the desired card, such as **General-purpose chatbot**, click **Use this template**, name your agent in the pop-up dialogue, and click **OK** to confirm.  
+
+   *You are now taken to the **no-code workflow editor** page. The left panel lists the components (operators): Above the dividing line are the RAG-specific components; below the line are tools. We are still working to expand the component list.*
+
+   ![workflow_editor](https://github.com/user-attachments/assets/9fc6891c-7784-43b8-ab4a-3b08a9e551c4)
+
+4. General speaking, now you can do the following:
+   - Drag and drop a desired component to your workflow,
+   - Select the knowledge base to use,
+   - Update settings of specific components,
+   - Update LLM settings
+   - Sets the input and output for a specific component, and more.
+5. Click **Save** to apply changes to your agent and **Run** to test it.
+
+## Components
+
+Please review the flowing description of the RAG-specific components before you proceed:
+
+| Component      | Description                                                  |
+| -------------- | ------------------------------------------------------------ |
+| **Retrieval**  | A component that retrieves information from specified knowledge bases and returns 'Empty response' if no information is found. Ensure the correct knowledge bases are selected. |
+| **Generate**   | A component that prompts the LLM to generate responses. You must ensure the prompt is set correctly. |
+| **Answer**     | A component that serves as the interface between human and the bot, receiving user inputs and displaying the agent's responses. |
+| **Categorize** | A component that uses the LLM to classify user inputs into predefined categories. Ensure you specify the name, description, and examples for each category, along with the corresponding next component. |
+| **Message**    | A component that sends out a static message. If multiple messages are supplied, it randomly selects one to send. Ensure its downstream is **Answer**, the interface component. |
+| **Relevant**   | A component that uses the LLM to assess whether the upstream output is relevant to the user's latest query. Ensure you specify the next component for each judge result. |
+| **Rewrite**    | A component that refines a user query if it fails to retrieve relevant information from the knowledge base. It repeats this process until the predefined looping upper limit is reached. Ensure its upstream is **Relevant** and downstream is **Retrieval**. |
+| **Keyword**    | A component that retrieves top N search results from wikipedia.org. Ensure the TopN value is set properly before use. |
+
+:::caution NOTE
+
+- Ensure **Rewrite**'s upstream component is **Relevant** and downstream component is **Retrieval**.
+- Ensure the downstream component of **Message** is **Answer**.
+- The downstream component of **Begin** is always **Answer**.
+
+:::
+
+## Basic operations
+
+| Operation                 | Description                                                  |
+| ------------------------- | ------------------------------------------------------------ |
+| Add a component           | Drag and drop the desired component from the left panel onto the canvas. |
+| Delete a component        | On the canvas, hover over the three dots (...) of the component to display the delete option, then select it to remove the component. |
+| Copy a component          | On the canvas, hover over the three dots (...) of the component to display the copy option, then select it to make a copy the component. |
+| Update component settings | On the canvas, click the desired component to display the component settings. |

docs/guides/agentic_rag/general_purpose_chatbot.md

@@ -0,0 +1,102 @@
+---
+sidebar_position: 2
+slug: /general_purpose_chatbot
+---
+
+# Create a general-purpose chatbot
+
+Chatbot is one of the most common AI scenarios. However, effectively understanding user queries and responding appropriately remains a challenge. RAGFlow's general-purpose chatbot agent is our attempt to tackle this longstanding issue.  
+
+This chatbot closely resembles the chatbot introduced in [Start an AI chat](../start_chat.md), but with a key difference - it introduces a reflective mechanism that allows it to improve the retrieval from the target knowledge bases by rewriting the user's query.
+
+This document provides guides on creating such a chatbot using our chatbot template.
+
+## Prerequisites
+
+1. Ensure you have properly set the LLM to use. See the guides on [Configure your API key](../llm_api_key_setup.md) or [Deploy a local LLM](../deploy_local_llm.mdx) for more information.
+2. Ensure you have a knowledge base configured and the corresponding files properly parsed. See the guide on [Configure a knowledge base](../configure_knowledge_base.md) for more information.
+3. Make sure you have read the [Introduction to Agentic RAG](./agentic_rag_introduction.md).
+
+## Create a chatbot agent from template
+
+To create a general-purpose chatbot agent using our template:
+
+1. Click the **Agent** tab in the middle top of the page to show the **Agent** page.
+2. Click **+ Create agent** on the top right of the page to show the **agent template** page.
+3. On the **agent template** page, hover over the card on **General-purpose chatbot** and click **Use this template**.  
+   *You are now directed to the **no-code workflow editor** page.*
+
+   ![workflow_editor](https://github.com/user-attachments/assets/9fc6891c-7784-43b8-ab4a-3b08a9e551c4)
+
+:::tip NOTE
+RAGFlow's no-code editor spares you the trouble of coding, making agent development effortless.
+:::
+
+## Understand each component in the template
+
+Here’s a breakdown of each component and its role and requirements in the chatbot template:
+
+- **Begin**
+  - Function: Sets the opening greeting for the user.
+  - Purpose: Establishes a welcoming atmosphere and prepares the user for interaction.
+
+- **Answer**
+  - Function: Serves as the interface between human and the bot.
+  - Role: Acts as the downstream component of **Begin**.  
+  - Note: Though named "Answer", it does not engage with the LLM.
+
+- **Retrieval**
+  - Function: Retrieves information from specified knowledge base(s).
+  - Requirement: Must have `knowledgebases` set up to function.
+
+- **Relevant**
+  - Function: Assesses the relevance of the retrieved information from the **Retrieval** component to the user query.
+  - Process:  
+    - If relevant, it directs the data to the **Generate** component for final response generation.
+    - Otherwise, it triggers the **Rewrite** component to refine the user query and redo the retrival process.
+
+- **Generate**
+  - Function: Prompts the LLM to generate responses based on the retrieved information.  
+  - Note: The prompt settings allow you to control the way in which the LLM generates responses. Be sure to review the prompts and make necessary changes.
+
+- **Rewrite**:  
+  - Function: Refines a user query when no relevant information from the knowledge base is retrieved.  
+  - Usage: Often used in conjunction with **Relevant** and **Retrieval** to create a reflective/feedback loop.  
+
+## Configure your chatbot agent
+
+1. Click **Begin** to set an opening greeting:  
+   ![opener](https://github.com/user-attachments/assets/4416bc16-2a84-4f24-a19b-6dc8b1de0908)
+
+2. Click **Retrieval** to select the right knowledge base(s) and make any necessary adjustments:  
+   ![setting_knowledge_bases](https://github.com/user-attachments/assets/5f694820-5651-45bc-afd6-cf580ca0228d)
+
+3. Click **Generate** to configure the LLM's summarization behavior:  
+   3.1. Confirm the model.  
+   3.2. Review the prompt settings. If there are variables, ensure they match the correct component IDs:  
+   ![prompt_settings](https://github.com/user-attachments/assets/19e94ea7-7f62-4b73-b526-32fcfa62f1e9)
+
+4. Click **Relevant** to review or change its settings:  
+   *You may retain the current settings, but feel free to experiment with changes to understand how the agent operates.*
+   ![relevant_settings](https://github.com/user-attachments/assets/f582cc1c-0dd5-499c-813a-294dbfb941dd)
+
+5. Click **Rewrite** to select a different model for query rewriting or update the maximum loop times for query rewriting:  
+   ![choose_model](https://github.com/user-attachments/assets/2bac1d6c-c4f1-42ac-997b-102858c3f550)
+   ![loop_time](https://github.com/user-attachments/assets/09a4ce34-7aac-496f-aa59-d8aa33bf0b1f)
+
+:::danger NOTE
+Increasing the maximum loop times may significantly extend the time required to receive the final response.
+:::
+
+1. Update your workflow where you see necessary.
+
+2. Click to **Save** to apply your changes.  
+   *Your agent appears as one of the agent cards on the **Agent** page.*
+
+## Test your chatbot agent
+
+1. Find your chatbot agent on the **Agent** page:  
+   ![find_chatbot](https://github.com/user-attachments/assets/6e6382c6-9a86-4190-9fdd-e363b7f64ba9)
+
+2. Experiment with your questions to verify if this chatbot functions as intended:  
+   ![test_chatbot](https://github.com/user-attachments/assets/c074d3bd-4c39-4b05-a68b-1fd361f256b3)

docs/guides/configure_knowledge_base.md

@@ -128,7 +128,7 @@ RAGFlow uses multiple recall of both full-text search and vector search in its c
 
 ## Search for knowledge base
 
-As of RAGFlow v0.10.0, the search feature is still in a rudimentary form, supporting only knowledge base search by name.
+As of RAGFlow v0.11.0, the search feature is still in a rudimentary form, supporting only knowledge base search by name.
 
 ![search knowledge base](https://github.com/infiniflow/ragflow/assets/93570324/836ae94c-2438-42be-879e-c7ad2a59693e)
 

docs/guides/deploy_local_llm.mdx

@@ -1,5 +1,5 @@
 ---
-sidebar_position: 5
+sidebar_position: 6
 slug: /deploy_local_llm
 ---
 
@@ -7,7 +7,7 @@ slug: /deploy_local_llm
 import Tabs from '@theme/Tabs';
 import TabItem from '@theme/TabItem';
 
-RAGFlow supports deploying models locally using Ollama or Xinference. If you have locally deployed models to leverage or wish to enable GPU or CUDA for inference acceleration, you can bind Ollama or Xinference into RAGFlow and use either of them as a local "server" for interacting with your local models.
+RAGFlow supports deploying models locally using Ollama, Xinference, IPEX-LLM, or jina. If you have locally deployed models to leverage or wish to enable GPU or CUDA for inference acceleration, you can bind Ollama or Xinference into RAGFlow and use either of them as a local "server" for interacting with your local models.
 
 RAGFlow seamlessly integrates with Ollama and Xinference, without the need for further environment configurations. You can use them to deploy two types of local models in RAGFlow: chat models and embedding models.
 
@@ -15,40 +15,6 @@ RAGFlow seamlessly integrates with Ollama and Xinference, without the need for f
 This user guide does not intend to cover much of the installation or configuration details of Ollama or Xinference; its focus is on configurations inside RAGFlow. For the most current information, you may need to check out the official site of Ollama or Xinference.
 :::
 
-# Deploy a local model using jina 
-
-[Jina](https://github.com/jina-ai/jina) lets you build AI services and pipelines that communicate via gRPC, HTTP and WebSockets, then scale them up and deploy to production.
-
-To deploy a local model, e.g., **gpt2**, using Jina:
-
-### 1. Check firewall settings
-
-Ensure that your host machine's firewall allows inbound connections on port 12345.
-
-```bash
-sudo ufw allow 12345/tcp
-```
-
-### 2.install jina package
-
-```bash
-pip install jina
-```
-
-### 3. deployment local model
-
-Step 1: Navigate to the rag/svr directory.
-
-```bash
-cd rag/svr
-```
-
-Step 2: Use Python to run the jina_server.py script and pass in the model name or the local path of the model (the script only supports loading models downloaded from Huggingface)
-
-```bash
-python jina_server.py  --model_name gpt2
-```
-
 ## Deploy a local model using Ollama
 
 [Ollama](https://github.com/ollama/ollama) enables you to run open-source large language models that you deployed locally. It bundles model weights, configurations, and data into a single package, defined by a Modelfile, and optimizes setup and configurations, including GPU usage.
@@ -347,3 +313,36 @@ To enable IPEX-LLM accelerated Ollama in RAGFlow, you must also complete the con
 2. [Complete basic Ollama settings](#5-complete-basic-ollama-settings)
 3. [Update System Model Settings](#6-update-system-model-settings)
 4. [Update Chat Configuration](#7-update-chat-configuration)
+
+## Deploy a local model using jina 
+
+To deploy a local model, e.g., **gpt2**, using jina:
+
+### 1. Check firewall settings
+
+Ensure that your host machine's firewall allows inbound connections on port 12345.
+
+```bash
+sudo ufw allow 12345/tcp
+```
+
+### 2. Install jina package
+
+```bash
+pip install jina
+```
+
+### 3. Deploy a local model
+
+Step 1: Navigate to the **rag/svr** directory.
+
+```bash
+cd rag/svr
+```
+
+Step 2: Run **jina_server.py**, specifying either the model's name or its local directory: 
+
+```bash
+python jina_server.py  --model_name gpt2
+```
+> The script only supports models downloaded from Hugging Face.

docs/guides/llm_api_key_setup.md

@@ -1,5 +1,5 @@
 ---
-sidebar_position: 4
+sidebar_position: 5
 slug: /llm_api_key_setup
 ---
 

docs/guides/manage_files.md

@@ -1,5 +1,5 @@
 ---
-sidebar_position: 3
+sidebar_position: 4
 slug: /manage_files
 ---
 
@@ -49,7 +49,7 @@ You can link your file to one knowledge base or multiple knowledge bases at one
 
 ## Search files or folders
 
-As of RAGFlow v0.10.0, the search feature is still in a rudimentary form, supporting only file and folder search in the current directory by name (files or folders in the child directory will not be retrieved).
+As of RAGFlow v0.11.0, the search feature is still in a rudimentary form, supporting only file and folder search in the current directory by name (files or folders in the child directory will not be retrieved).
 
 ![search file](https://github.com/infiniflow/ragflow/assets/93570324/77ffc2e5-bd80-4ed1-841f-068e664efffe)
 
@@ -81,4 +81,4 @@ RAGFlow's file management allows you to download an uploaded file:
 
 ![download_file](https://github.com/infiniflow/ragflow/assets/93570324/cf3b297f-7d9b-4522-bf5f-4f45743e4ed5)
 
-> As of RAGFlow v0.10.0, bulk download is not supported, nor can you download an entire folder. 
+> As of RAGFlow v0.11.0, bulk download is not supported, nor can you download an entire folder. 

docs/quickstart.mdx

@@ -34,7 +34,7 @@ This section provides instructions on setting up the RAGFlow server on Linux. If
 
    `vm.max_map_count`. This value sets the maximum number of memory map areas a process may have. Its default value is 65530. While most applications require fewer than a thousand maps, reducing this value can result in abmornal behaviors, and the system will throw out-of-memory errors when a process reaches the limitation. 
 
-   RAGFlow v0.10.0 uses Elasticsearch for multiple recall. Setting the value of `vm.max_map_count` correctly is crucial to the proper functioning of the Elasticsearch component.
+   RAGFlow v0.11.0 uses Elasticsearch for multiple recall. Setting the value of `vm.max_map_count` correctly is crucial to the proper functioning of the Elasticsearch component.
 
 <Tabs
   defaultValue="linux"
@@ -177,7 +177,7 @@ This section provides instructions on setting up the RAGFlow server on Linux. If
 
 3. Build the pre-built Docker images and start up the server:
 
-   > Running the following commands automatically downloads the *dev* version RAGFlow Docker image. To download and run a specified Docker version, update `RAGFLOW_VERSION` in **docker/.env** to the intended version, for example `RAGFLOW_VERSION=v0.10.0`, before running the following commands.
+   > Running the following commands automatically downloads the *dev* version RAGFlow Docker image. To download and run a specified Docker version, update `RAGFLOW_VERSION` in **docker/.env** to the intended version, for example `RAGFLOW_VERSION=v0.11.0`, before running the following commands.
 
    ```bash
    $ cd ragflow/docker
@@ -298,7 +298,7 @@ Once you have selected an embedding model and used it to parse a file, you are n
 
    _When the file parsing completes, its parsing status changes to **SUCCESS**._
 
-:::alert NOTE 
+:::caution NOTE 
 - If your file parsing gets stuck at below 1%, see [FAQ 4.3](https://ragflow.io/docs/dev/faq#43-why-does-my-document-parsing-stall-at-under-one-percent). 
 - If your file parsing gets stuck at near completion, see [FAQ 4.4](https://ragflow.io/docs/dev/faq#44-why-does-my-pdf-parsing-stall-near-completion-while-the-log-does-not-show-any-error)
 :::

docs/references/api.md

@@ -398,7 +398,43 @@ This method uploads a specific file to a specified knowledge base.
     "retmsg": "success"
 }
 ```
+### Demo for Upload File(Python) 
 
+```python
+# upload_to_kb.py  
+import requests
+
+
+def upload_file_to_kb(file_path, kb_name, token='ragflow-xxxxxxxxxxxxx', parser_id='naive'):
+    """  
+    Uploads a file to a knowledge base.  
+
+    Args:  
+    - file_path: Path to the file to upload.  
+    - kb_name: Name of the target knowledge base.  
+    - parser_id: ID of the chosen file parser (defaults to 'naive').  
+    - token: API token for authentication.  
+    """
+    url = 'http://127.0.0.1/v1/api/document/upload'  # Replace with your actual API URL  
+    files = {'file': open(file_path, 'rb')}  # The file to upload  
+    data = {'kb_name': kb_name, 'parser_id': parser_id, 'run': '1'}  # Additional form data  
+    headers = {'Authorization': f'Bearer {token}'}  # Replace with your actual Bearer token  
+
+    response = requests.post(url, files=files, data=data, headers=headers)
+
+    if response.status_code == 200:
+        print("File uploaded successfully:", response.json())
+    else:
+        print("Failed to upload file:", response.status_code, response.text)
+
+file_to_upload = './ai_intro.pdf'  # For example: './documents/report.pdf'  
+knowledge_base_name = 'AI_knowledge_base'
+# Assume you have already obtained your token and set it here  
+token = 'ragflow-xxxxxxxxxxxxx'
+
+# Call the function to upload the file  
+upload_file_to_kb(file_to_upload, knowledge_base_name, token=token)
+```
 ## Get document chunks
 
 This method retrieves the chunks of a specific document by `doc_name` or `doc_id`.

docs/references/faq.md

@@ -340,8 +340,8 @@ This exception occurs when starting up the RAGFlow server. Try the following:
 ![hint102](https://github.com/infiniflow/ragflow/assets/93570324/6633d892-b4f8-49b5-9a0a-37a0a8fba3d2)
 
 1. Ensure that the RAGFlow server can access the base URL.
-2. Do not forget to append **/v1/** to **http://IP:port**: 
+2. Do not forget to append `/v1/` to `http://IP:port`:  
-   **http://IP:port/v1/**
+   `http://IP:port/v1/`
 
 #### 4.16 `FileNotFoundError: [Errno 2] No such file or directory`
 
@@ -370,7 +370,7 @@ You limit what the system responds to what you specify in **Empty response** if
 
 ### 4. How to run RAGFlow with a locally deployed LLM?
 
-You can use Ollama to deploy local LLM. See [here](https://github.com/infiniflow/ragflow/blob/main/docs/guides/deploy_local_llm.md) for more information. 
+You can use Ollama to deploy local LLM. See [here](../guides/deploy_local_llm.mdx) for more information.
 
 ### 5. How to link up ragflow and ollama servers?
 

graphrag/claim_extractor.py

@@ -1,22 +1,10 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
 """
+
 import argparse
 import json
 import logging
@@ -182,7 +170,7 @@ class ClaimExtractor:
                     }
         text = perform_variable_replacements(self._extraction_prompt, variables=variables)
         gen_conf = {"temperature": 0.5}
-        results = self._llm.chat(text, [], gen_conf)
+        results = self._llm.chat(text, [{"role": "user", "content": "Output:"}], gen_conf)
         claims = results.strip().removesuffix(completion_delimiter)
         history = [{"role": "system", "content": text}, {"role": "assistant", "content": results}]
 

graphrag/claim_prompt.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

graphrag/community_report_prompt.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

graphrag/community_reports_extractor.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
@@ -89,7 +76,7 @@ class CommunityReportsExtractor:
                 text = perform_variable_replacements(self._extraction_prompt, variables=prompt_variables)
                 gen_conf = {"temperature": 0.3}
                 try:
-                    response = self._llm.chat(text, [], gen_conf)
+                    response = self._llm.chat(text, [{"role": "user", "content": "Output:"}], gen_conf)
                     token_count += num_tokens_from_string(text + response)
                     response = re.sub(r"^[^\{]*", "", response)
                     response = re.sub(r"[^\}]*$", "", response)
@@ -138,4 +125,5 @@ class CommunityReportsExtractor:
         report_sections = "\n\n".join(
             f"## {finding_summary(f)}\n\n{finding_explanation(f)}" for f in findings
         )
+     
         return f"# {title}\n\n{summary}\n\n{report_sections}"

graphrag/description_summary.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

graphrag/entity_embedding.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

graphrag/entity_resolution.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+
 import logging
 import re
 import traceback
@@ -124,7 +125,7 @@ class EntityResolution:
                     }
                     text = perform_variable_replacements(self._resolution_prompt, variables=variables)
 
-                    response = self._llm.chat(text, [], gen_conf)
+                    response = self._llm.chat(text, [{"role": "user", "content": "Output:"}], gen_conf)
                     result = self._process_results(len(candidate_resolution_i[1]), response,
                                                    prompt_variables.get(self._record_delimiter_key,
                                                                         DEFAULT_RECORD_DELIMITER),

graphrag/graph_extractor.py

@@ -1,21 +1,10 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
 """
+
 import logging
 import numbers
 import re
@@ -142,7 +131,7 @@ class GraphExtractor:
                 total_token_count += token_count
                 if callback: callback(msg=f"{doc_index+1}/{total}, elapsed: {timer() - st}s, used tokens: {total_token_count}")
             except Exception as e:
-                if callback: callback("Knowledge graph extraction error:{}".format(str(e)))
+                if callback: callback(msg="Knowledge graph extraction error:{}".format(str(e)))
                 logging.exception("error extracting graph")
                 self._on_error(
                     e,
@@ -174,7 +163,7 @@ class GraphExtractor:
         token_count = 0
         text = perform_variable_replacements(self._extraction_prompt, variables=variables)
         gen_conf = {"temperature": 0.3}
-        response = self._llm.chat(text, [], gen_conf)
+        response = self._llm.chat(text, [{"role": "user", "content": "Output:"}], gen_conf)
         token_count = num_tokens_from_string(text + response)
 
         results = response or ""

graphrag/graph_prompt.py

@@ -1,22 +1,10 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)
 """
+
 GRAPH_EXTRACTION_PROMPT = """
 -Goal-
 Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.

graphrag/leiden.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

graphrag/mind_map_extractor.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+
 import collections
 import logging
 import re
@@ -65,7 +66,7 @@ class MindMapExtractor:
             obj = [obj]
         if isinstance(obj, list):
             for i in obj: keyset.add(i)
-            return [{"id": re.sub(r"\*+", "", i), "children": []} for i in obj]
+            return [{"id": re.sub(r"\*+", "", i), "children": []} for i in obj if re.sub(r"\*+", "", i)]
         arr = []
         for k, v in obj.items():
             k = self._key(k)
@@ -106,7 +107,7 @@ class MindMapExtractor:
                 res.append(_.result())
 
             if not res:
-                return MindMapResult(output={"root":{}})
+                return MindMapResult(output={"id": "root", "children": []})
 
             merge_json = reduce(self._merge, res)
             if len(merge_json.keys()) > 1:
@@ -179,7 +180,7 @@ class MindMapExtractor:
         }
         text = perform_variable_replacements(self._mind_map_prompt, variables=variables)
         gen_conf = {"temperature": 0.5}
-        response = self._llm.chat(text, [], gen_conf)
+        response = self._llm.chat(text, [{"role": "user", "content": "Output:"}], gen_conf)
         response = re.sub(r"```[^\n]*", "", response)
         print(response)
         print("---------------------------------------------------\n", self._todict(markdown_to_json.dictify(response)))

graphrag/mind_map_prompt.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+
 MIND_MAP_EXTRACTION_PROMPT = """
 - Role: You're a talent text processor to summarize a piece of text into a mind map.
 

graphrag/search.py

@@ -23,7 +23,7 @@ from rag.nlp.search import Dealer
 
 
 class KGSearch(Dealer):
-    def search(self, req, idxnm, emb_mdl=None):
+    def search(self, req, idxnm, emb_mdl=None, highlight=False):
         def merge_into_first(sres, title=""):
             df,texts = [],[]
             for d in sres["hits"]["hits"]:
@@ -93,9 +93,9 @@ class KGSearch(Dealer):
         s = s.query(bqry)[0: 6]
         s = s.to_dict()
         txt_res = self.es.search(deepcopy(s), idxnm=idxnm, timeout="600s", src=src)
-        txt_ids = self.es.getDocIds(comm_res)
+        txt_ids = self.es.getDocIds(txt_res)
         if merge_into_first(txt_res, "-Original Content-"):
-            txt_ids = comm_ids[0:1]
+            txt_ids = txt_ids[0:1]
 
         return self.SearchResult(
             total=len(ent_ids) + len(comm_ids) + len(txt_ids),

graphrag/utils.py

@@ -1,18 +1,5 @@
-#
+# Copyright (c) 2024 Microsoft Corporation.
-#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+# Licensed under the MIT License
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  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.
-#
 """
 Reference:
  - [graphrag](https://github.com/microsoft/graphrag)

rag/benchmark.py

@@ -0,0 +1,234 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+import json
+import os
+from collections import defaultdict
+from api.db import LLMType
+from api.db.services.llm_service import LLMBundle
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.settings import retrievaler
+from api.utils import get_uuid
+from rag.nlp import tokenize, search
+from rag.utils.es_conn import ELASTICSEARCH
+from ranx import evaluate
+import pandas as pd
+from tqdm import tqdm
+
+
+class Benchmark:
+    def __init__(self, kb_id):
+        e, kb = KnowledgebaseService.get_by_id(kb_id)
+        self.similarity_threshold = kb.similarity_threshold
+        self.vector_similarity_weight = kb.vector_similarity_weight
+        self.embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id, lang=kb.language)
+
+    def _get_benchmarks(self, query, dataset_idxnm, count=16):
+        req = {"question": query, "size": count, "vector": True, "similarity": self.similarity_threshold}
+        sres = retrievaler.search(req, search.index_name(dataset_idxnm), self.embd_mdl)
+        return sres
+
+    def _get_retrieval(self, qrels, dataset_idxnm):
+        run = defaultdict(dict)
+        query_list = list(qrels.keys())
+        for query in query_list:
+            sres = self._get_benchmarks(query, dataset_idxnm)
+            sim, _, _ = retrievaler.rerank(sres, query, 1 - self.vector_similarity_weight,
+                                           self.vector_similarity_weight)
+            for index, id in enumerate(sres.ids):
+                run[query][id] = sim[index]
+        return run
+
+    def embedding(self, docs, batch_size=16):
+        vects = []
+        cnts = [d["content_with_weight"] for d in docs]
+        for i in range(0, len(cnts), batch_size):
+            vts, c = self.embd_mdl.encode(cnts[i: i + batch_size])
+            vects.extend(vts.tolist())
+        assert len(docs) == len(vects)
+        for i, d in enumerate(docs):
+            v = vects[i]
+            d["q_%d_vec" % len(v)] = v
+        return docs
+
+    def ms_marco_index(self, file_path, index_name):
+        qrels = defaultdict(dict)
+        texts = defaultdict(dict)
+        docs = []
+        filelist = os.listdir(file_path)
+        for dir in filelist:
+            data = pd.read_parquet(os.path.join(file_path, dir))
+            for i in tqdm(range(len(data)), colour="green", desc="Indexing:" + dir):
+
+                query = data.iloc[i]['query']
+                for rel, text in zip(data.iloc[i]['passages']['is_selected'], data.iloc[i]['passages']['passage_text']):
+                    d = {
+                        "id": get_uuid()
+                    }
+                    tokenize(d, text, "english")
+                    docs.append(d)
+                    texts[d["id"]] = text
+                    qrels[query][d["id"]] = int(rel)
+                if len(docs) >= 32:
+                    docs = self.embedding(docs)
+                    ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+                    docs = []
+
+        docs = self.embedding(docs)
+        ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+        return qrels, texts
+
+    def trivia_qa_index(self, file_path, index_name):
+        qrels = defaultdict(dict)
+        texts = defaultdict(dict)
+        docs = []
+        filelist = os.listdir(file_path)
+        for dir in filelist:
+            data = pd.read_parquet(os.path.join(file_path, dir))
+            for i in tqdm(range(len(data)), colour="green", desc="Indexing:" + dir):
+                query = data.iloc[i]['question']
+                for rel, text in zip(data.iloc[i]["search_results"]['rank'],
+                                     data.iloc[i]["search_results"]['search_context']):
+                    d = {
+                        "id": get_uuid()
+                    }
+                    tokenize(d, text, "english")
+                    docs.append(d)
+                    texts[d["id"]] = text
+                    qrels[query][d["id"]] = int(rel)
+                if len(docs) >= 32:
+                    docs = self.embedding(docs)
+                    ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+                    docs = []
+
+        docs = self.embedding(docs)
+        ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+        return qrels, texts
+
+    def miracl_index(self, file_path, corpus_path, index_name):
+
+        corpus_total = {}
+        for corpus_file in os.listdir(corpus_path):
+            tmp_data = pd.read_json(os.path.join(corpus_path, corpus_file), lines=True)
+            for index, i in tmp_data.iterrows():
+                corpus_total[i['docid']] = i['text']
+
+        topics_total = {}
+        for topics_file in os.listdir(os.path.join(file_path, 'topics')):
+            if 'test' in topics_file:
+                continue
+            tmp_data = pd.read_csv(os.path.join(file_path, 'topics', topics_file), sep='\t', names=['qid', 'query'])
+            for index, i in tmp_data.iterrows():
+                topics_total[i['qid']] = i['query']
+
+        qrels = defaultdict(dict)
+        texts = defaultdict(dict)
+        docs = []
+        for qrels_file in os.listdir(os.path.join(file_path, 'qrels')):
+            if 'test' in qrels_file:
+                continue
+
+            tmp_data = pd.read_csv(os.path.join(file_path, 'qrels', qrels_file), sep='\t',
+                                   names=['qid', 'Q0', 'docid', 'relevance'])
+            for i in tqdm(range(len(tmp_data)), colour="green", desc="Indexing:" + qrels_file):
+                query = topics_total[tmp_data.iloc[i]['qid']]
+                text = corpus_total[tmp_data.iloc[i]['docid']]
+                rel = tmp_data.iloc[i]['relevance']
+                d = {
+                    "id": get_uuid()
+                }
+                tokenize(d, text, 'english')
+                docs.append(d)
+                texts[d["id"]] = text
+                qrels[query][d["id"]] = int(rel)
+                if len(docs) >= 32:
+                    docs = self.embedding(docs)
+                    ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+                    docs = []
+
+        docs = self.embedding(docs)
+        ELASTICSEARCH.bulk(docs, search.index_name(index_name))
+
+        return qrels, texts
+
+    def save_results(self, qrels, run, texts, dataset, file_path):
+        keep_result = []
+        run_keys = list(run.keys())
+        for run_i in tqdm(range(len(run_keys)), desc="Calculating ndcg@10 for single query"):
+            key = run_keys[run_i]
+            keep_result.append({'query': key, 'qrel': qrels[key], 'run': run[key],
+                                'ndcg@10': evaluate({key: qrels[key]}, {key: run[key]}, "ndcg@10")})
+        keep_result = sorted(keep_result, key=lambda kk: kk['ndcg@10'])
+        with open(os.path.join(file_path, dataset + 'result.md'), 'w', encoding='utf-8') as f:
+            f.write('## Score For Every Query\n')
+            for keep_result_i in keep_result:
+                f.write('### query: ' + keep_result_i['query'] + ' ndcg@10:' + str(keep_result_i['ndcg@10']) + '\n')
+                scores = [[i[0], i[1]] for i in keep_result_i['run'].items()]
+                scores = sorted(scores, key=lambda kk: kk[1])
+                for score in scores[:10]:
+                    f.write('- text: ' + str(texts[score[0]]) + '\t qrel: ' + str(score[1]) + '\n')
+        print(os.path.join(file_path, dataset + '_result.md'), 'Saved!')
+
+    def __call__(self, dataset, file_path, miracl_corpus=''):
+        if dataset == "ms_marco_v1.1":
+            qrels, texts = self.ms_marco_index(file_path, "benchmark_ms_marco_v1.1")
+            run = self._get_retrieval(qrels, "benchmark_ms_marco_v1.1")
+            print(dataset, evaluate(qrels, run, ["ndcg@10", "map@5", "mrr"]))
+            self.save_results(qrels, run, texts, dataset, file_path)
+        if dataset == "trivia_qa":
+            qrels, texts = self.trivia_qa_index(file_path, "benchmark_trivia_qa")
+            run = self._get_retrieval(qrels, "benchmark_trivia_qa")
+            print(dataset, evaluate(qrels, run, ["ndcg@10", "map@5", "mrr"]))
+            self.save_results(qrels, run, texts, dataset, file_path)
+        if dataset == "miracl":
+            for lang in ['ar', 'bn', 'de', 'en', 'es', 'fa', 'fi', 'fr', 'hi', 'id', 'ja', 'ko', 'ru', 'sw', 'te', 'th',
+                         'yo', 'zh']:
+                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang)):
+                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang) + ' not found!')
+                    continue
+                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang, 'qrels')):
+                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang, 'qrels') + 'not found!')
+                    continue
+                if not os.path.isdir(os.path.join(file_path, 'miracl-v1.0-' + lang, 'topics')):
+                    print('Directory: ' + os.path.join(file_path, 'miracl-v1.0-' + lang, 'topics') + 'not found!')
+                    continue
+                if not os.path.isdir(os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang)):
+                    print('Directory: ' + os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang) + ' not found!')
+                    continue
+                qrels, texts = self.miracl_index(os.path.join(file_path, 'miracl-v1.0-' + lang),
+                                                 os.path.join(miracl_corpus, 'miracl-corpus-v1.0-' + lang),
+                                                 "benchmark_miracl_" + lang)
+                run = self._get_retrieval(qrels, "benchmark_miracl_" + lang)
+                print(dataset, evaluate(qrels, run, ["ndcg@10", "map@5", "mrr"]))
+                self.save_results(qrels, run, texts, dataset, file_path)
+
+
+if __name__ == '__main__':
+    print('*****************RAGFlow Benchmark*****************')
+    kb_id = input('Please input kb_id:\n')
+    ex = Benchmark(kb_id)
+    dataset = input(
+        'RAGFlow Benchmark Support:\n\tms_marco_v1.1:<https://huggingface.co/datasets/microsoft/ms_marco>\n\ttrivia_qa:<https://huggingface.co/datasets/mandarjoshi/trivia_qa>\n\tmiracl:<https://huggingface.co/datasets/miracl/miracl>\nPlease input dataset choice:\n')
+    if dataset in ['ms_marco_v1.1', 'trivia_qa']:
+        if dataset == "ms_marco_v1.1":
+            print("Notice: Please provide the ms_marco_v1.1 dataset only. ms_marco_v2.1 is not supported!")
+        dataset_path = input('Please input ' + dataset + ' dataset path:\n')
+        ex(dataset, dataset_path)
+    elif dataset == 'miracl':
+        dataset_path = input('Please input ' + dataset + ' dataset path:\n')
+        corpus_path = input('Please input ' + dataset + '-corpus dataset path:\n')
+        ex(dataset, dataset_path, miracl_corpus=corpus_path)
+    else:
+        print("Dataset: ", dataset, "not supported!")

rag/llm/__init__.py

@@ -18,6 +18,7 @@ from .chat_model import *
 from .cv_model import *
 from .rerank_model import *
 from .sequence2txt_model import *
+from .tts_model import * 
 
 EmbeddingModel = {
     "Ollama": OllamaEmbed,
@@ -44,7 +45,8 @@ EmbeddingModel = {
     "Upstage": UpstageEmbed,
     "SILICONFLOW": SILICONFLOWEmbed,
     "Replicate": ReplicateEmbed,
-    "BaiduYiyan": BaiduYiyanEmbed
+    "BaiduYiyan": BaiduYiyanEmbed,
+    "Voyage AI": VoyageEmbed
 }
 
 
@@ -103,7 +105,9 @@ ChatModel = {
     "Replicate": ReplicateChat,
     "Tencent Hunyuan": HunyuanChat,
     "XunFei Spark": SparkChat,
-    "BaiduYiyan": BaiduYiyanChat
+    "BaiduYiyan": BaiduYiyanChat,
+    "Anthropic": AnthropicChat,
+    "Google Cloud": GoogleChat,
 }
 
 
@@ -118,7 +122,8 @@ RerankModel = {
     "cohere": CoHereRerank,
     "TogetherAI": TogetherAIRerank,
     "SILICONFLOW": SILICONFLOWRerank,
-    "BaiduYiyan": BaiduYiyanRerank
+    "BaiduYiyan": BaiduYiyanRerank,
+    "Voyage AI": VoyageRerank
 }
 
 
@@ -127,5 +132,11 @@ Seq2txtModel = {
     "Tongyi-Qianwen": QWenSeq2txt,
     "Ollama": OllamaSeq2txt,
     "Azure-OpenAI": AzureSeq2txt,
-    "Xinference": XinferenceSeq2txt
+    "Xinference": XinferenceSeq2txt,
+    "Tencent Cloud": TencentCloudSeq2txt
+}
+
+TTSModel = {
+    "Fish Audio": FishAudioTTS,
+    "Tongyi-Qianwen": QwenTTS
 }

rag/llm/chat_model.py

@@ -450,72 +450,16 @@ class LocalLLM(Base):
 
 
 class VolcEngineChat(Base):
-    def __init__(self, key, model_name, base_url):
+    def __init__(self, key, model_name, base_url='https://ark.cn-beijing.volces.com/api/v3'):
         """
         Since do not want to modify the original database fields, and the VolcEngine authentication method is quite special,
-        Assemble ak, sk, ep_id into api_key, store it as a dictionary type, and parse it for use
+        Assemble ark_api_key, ep_id into api_key, store it as a dictionary type, and parse it for use
         model_name is for display only
         """
-        self.client = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing')
+        base_url = base_url if base_url else 'https://ark.cn-beijing.volces.com/api/v3'
-        self.volc_ak = eval(key).get('volc_ak', '')
+        ark_api_key = json.loads(key).get('ark_api_key', '')
-        self.volc_sk = eval(key).get('volc_sk', '')
+        model_name = json.loads(key).get('ep_id', '')
-        self.client.set_ak(self.volc_ak)
+        super().__init__(ark_api_key, model_name, base_url)
-        self.client.set_sk(self.volc_sk)
-        self.model_name = eval(key).get('ep_id', '')
-
-    def chat(self, system, history, gen_conf):
-        if system:
-            history.insert(0, {"role": "system", "content": system})
-        try:
-            req = {
-                "parameters": {
-                    "min_new_tokens": gen_conf.get("min_new_tokens", 1),
-                    "top_k": gen_conf.get("top_k", 0),
-                    "max_prompt_tokens": gen_conf.get("max_prompt_tokens", 30000),
-                    "temperature": gen_conf.get("temperature", 0.1),
-                    "max_new_tokens": gen_conf.get("max_tokens", 1000),
-                    "top_p": gen_conf.get("top_p", 0.3),
-                },
-                "messages": history
-            }
-            response = self.client.chat(self.model_name, req)
-            ans = response.choices[0].message.content.strip()
-            if response.choices[0].finish_reason == "length":
-                ans += "...\nFor the content length reason, it stopped, continue?" if is_english(
-                    [ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？"
-            return ans, response.usage.total_tokens
-        except Exception as e:
-            return "**ERROR**: " + str(e), 0
-
-    def chat_streamly(self, system, history, gen_conf):
-        if system:
-            history.insert(0, {"role": "system", "content": system})
-        ans = ""
-        tk_count = 0
-        try:
-            req = {
-                "parameters": {
-                    "min_new_tokens": gen_conf.get("min_new_tokens", 1),
-                    "top_k": gen_conf.get("top_k", 0),
-                    "max_prompt_tokens": gen_conf.get("max_prompt_tokens", 30000),
-                    "temperature": gen_conf.get("temperature", 0.1),
-                    "max_new_tokens": gen_conf.get("max_tokens", 1000),
-                    "top_p": gen_conf.get("top_p", 0.3),
-                },
-                "messages": history
-            }
-            stream = self.client.stream_chat(self.model_name, req)
-            for resp in stream:
-                if not resp.choices[0].message.content:
-                    continue
-                ans += resp.choices[0].message.content
-                if resp.choices[0].finish_reason == "stop":
-                    tk_count = resp.usage.total_tokens
-                yield ans
-
-        except Exception as e:
-            yield ans + "\n**ERROR**: " + str(e)
-        yield tk_count
 
 
 class MiniMaxChat(Base):
@@ -658,9 +602,9 @@ class BedrockChat(Base):
 
     def __init__(self, key, model_name, **kwargs):
         import boto3
-        self.bedrock_ak = eval(key).get('bedrock_ak', '')
+        self.bedrock_ak = json.loads(key).get('bedrock_ak', '')
-        self.bedrock_sk = eval(key).get('bedrock_sk', '')
+        self.bedrock_sk = json.loads(key).get('bedrock_sk', '')
-        self.bedrock_region = eval(key).get('bedrock_region', '')
+        self.bedrock_region = json.loads(key).get('bedrock_region', '')
         self.model_name = model_name
         self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
                                    aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
@@ -757,9 +701,13 @@ class GeminiChat(Base):
         self.model = GenerativeModel(model_name=self.model_name)
         self.model._client = _client
         
+        
     def chat(self,system,history,gen_conf):
+        from google.generativeai.types import content_types
+        
         if system:
-            history.insert(0, {"role": "user", "parts": system})
+            self.model._system_instruction = content_types.to_content(system)
+            
         if 'max_tokens' in gen_conf:
             gen_conf['max_output_tokens'] = gen_conf['max_tokens']
         for k in list(gen_conf.keys()):
@@ -781,8 +729,10 @@ class GeminiChat(Base):
             return "**ERROR**: " + str(e), 0
 
     def chat_streamly(self, system, history, gen_conf):
+        from google.generativeai.types import content_types
+        
         if system:
-            history.insert(0, {"role": "user", "parts": system})
+            self.model._system_instruction = content_types.to_content(system)
         if 'max_tokens' in gen_conf:
             gen_conf['max_output_tokens'] = gen_conf['max_tokens']
         for k in list(gen_conf.keys()):
@@ -1249,3 +1199,218 @@ class BaiduYiyanChat(Base):
             return ans + "\n**ERROR**: " + str(e), 0
 
         yield total_tokens
+
+
+class AnthropicChat(Base):
+    def __init__(self, key, model_name, base_url=None):
+        import anthropic
+
+        self.client = anthropic.Anthropic(api_key=key)
+        self.model_name = model_name
+        self.system = ""
+
+    def chat(self, system, history, gen_conf):
+        if system:
+            self.system = system
+        if "max_tokens" not in gen_conf:
+            gen_conf["max_tokens"] = 4096
+
+        try:
+            response = self.client.messages.create(
+                model=self.model_name,
+                messages=history,
+                system=self.system,
+                stream=False,
+                **gen_conf,
+            ).json()
+            ans = response["content"][0]["text"]
+            if response["stop_reason"] == "max_tokens":
+                ans += (
+                    "...\nFor the content length reason, it stopped, continue?"
+                    if is_english([ans])
+                    else "······\n由于长度的原因，回答被截断了，要继续吗？"
+                )
+            return (
+                ans,
+                response["usage"]["input_tokens"] + response["usage"]["output_tokens"],
+            )
+        except Exception as e:
+            return ans + "\n**ERROR**: " + str(e), 0
+
+    def chat_streamly(self, system, history, gen_conf):
+        if system:
+            self.system = system
+        if "max_tokens" not in gen_conf:
+            gen_conf["max_tokens"] = 4096
+
+        ans = ""
+        total_tokens = 0
+        try:
+            response = self.client.messages.create(
+                model=self.model_name,
+                messages=history,
+                system=self.system,
+                stream=True,
+                **gen_conf,
+            )
+            for res in response.iter_lines():
+                res = res.decode("utf-8")
+                if "content_block_delta" in res and "data" in res:
+                    text = json.loads(res[6:])["delta"]["text"]
+                    ans += text
+                    total_tokens += num_tokens_from_string(text)
+        except Exception as e:
+            yield ans + "\n**ERROR**: " + str(e)
+
+        yield total_tokens
+
+
+class GoogleChat(Base):
+    def __init__(self, key, model_name, base_url=None):
+        from google.oauth2 import service_account
+        import base64
+
+        key = json.load(key)
+        access_token = json.loads(
+            base64.b64decode(key.get("google_service_account_key", ""))
+        )
+        project_id = key.get("google_project_id", "")
+        region = key.get("google_region", "")
+
+        scopes = ["https://www.googleapis.com/auth/cloud-platform"]
+        self.model_name = model_name
+        self.system = ""
+
+        if "claude" in self.model_name:
+            from anthropic import AnthropicVertex
+            from google.auth.transport.requests import Request
+
+            if access_token:
+                credits = service_account.Credentials.from_service_account_info(
+                    access_token, scopes=scopes
+                )
+                request = Request()
+                credits.refresh(request)
+                token = credits.token
+                self.client = AnthropicVertex(
+                    region=region, project_id=project_id, access_token=token
+                )
+            else:
+                self.client = AnthropicVertex(region=region, project_id=project_id)
+        else:
+            from google.cloud import aiplatform
+            import vertexai.generative_models as glm
+
+            if access_token:
+                credits = service_account.Credentials.from_service_account_info(
+                    access_token
+                )
+                aiplatform.init(
+                    credentials=credits, project=project_id, location=region
+                )
+            else:
+                aiplatform.init(project=project_id, location=region)
+            self.client = glm.GenerativeModel(model_name=self.model_name)
+
+    def chat(self, system, history, gen_conf):
+        if system:
+            self.system = system
+
+        if "claude" in self.model_name:
+            if "max_tokens" not in gen_conf:
+                gen_conf["max_tokens"] = 4096
+            try:
+                response = self.client.messages.create(
+                    model=self.model_name,
+                    messages=history,
+                    system=self.system,
+                    stream=False,
+                    **gen_conf,
+                ).json()
+                ans = response["content"][0]["text"]
+                if response["stop_reason"] == "max_tokens":
+                    ans += (
+                        "...\nFor the content length reason, it stopped, continue?"
+                        if is_english([ans])
+                        else "······\n由于长度的原因，回答被截断了，要继续吗？"
+                    )
+                return (
+                    ans,
+                    response["usage"]["input_tokens"]
+                    + response["usage"]["output_tokens"],
+                )
+            except Exception as e:
+                return ans + "\n**ERROR**: " + str(e), 0
+        else:
+            self.client._system_instruction = self.system
+            if "max_tokens" in gen_conf:
+                gen_conf["max_output_tokens"] = gen_conf["max_tokens"]
+            for k in list(gen_conf.keys()):
+                if k not in ["temperature", "top_p", "max_output_tokens"]:
+                    del gen_conf[k]
+            for item in history:
+                if "role" in item and item["role"] == "assistant":
+                    item["role"] = "model"
+                if "content" in item:
+                    item["parts"] = item.pop("content")
+            try:
+                response = self.client.generate_content(
+                    history, generation_config=gen_conf
+                )
+                ans = response.text
+                return ans, response.usage_metadata.total_token_count
+            except Exception as e:
+                return "**ERROR**: " + str(e), 0
+
+    def chat_streamly(self, system, history, gen_conf):
+        if system:
+            self.system = system
+
+        if "claude" in self.model_name:
+            if "max_tokens" not in gen_conf:
+                gen_conf["max_tokens"] = 4096
+            ans = ""
+            total_tokens = 0
+            try:
+                response = self.client.messages.create(
+                    model=self.model_name,
+                    messages=history,
+                    system=self.system,
+                    stream=True,
+                    **gen_conf,
+                )
+                for res in response.iter_lines():
+                    res = res.decode("utf-8")
+                    if "content_block_delta" in res and "data" in res:
+                        text = json.loads(res[6:])["delta"]["text"]
+                        ans += text
+                        total_tokens += num_tokens_from_string(text)
+            except Exception as e:
+                yield ans + "\n**ERROR**: " + str(e)
+
+            yield total_tokens
+        else:
+            self.client._system_instruction = self.system
+            if "max_tokens" in gen_conf:
+                gen_conf["max_output_tokens"] = gen_conf["max_tokens"]
+            for k in list(gen_conf.keys()):
+                if k not in ["temperature", "top_p", "max_output_tokens"]:
+                    del gen_conf[k]
+            for item in history:
+                if "role" in item and item["role"] == "assistant":
+                    item["role"] = "model"
+                if "content" in item:
+                    item["parts"] = item.pop("content")
+            ans = ""
+            try:
+                response = self.model.generate_content(
+                    history, generation_config=gen_conf, stream=True
+                )
+                for resp in response:
+                    ans += resp.text
+                    yield ans
+
+            except Exception as e:
+                yield ans + "\n**ERROR**: " + str(e)
+
+            yield response._chunks[-1].usage_metadata.total_token_count

rag/llm/cv_model.py

@@ -293,9 +293,12 @@ class Zhipu4V(Base):
     def describe(self, image, max_tokens=1024):
         b64 = self.image2base64(image)
 
+        prompt = self.prompt(b64)
+        prompt[0]["content"][1]["type"] = "text"
+        
         res = self.client.chat.completions.create(
             model=self.model_name,
-            messages=self.prompt(b64),
+            messages=prompt,
             max_tokens=max_tokens,
         )
         return res.choices[0].message.content.strip(), res.usage.total_tokens

rag/llm/embedding_model.py

@@ -15,7 +15,7 @@
 #
 import re
 from typing import Optional
-import  threading
+import threading
 import requests
 from huggingface_hub import snapshot_download
 from openai.lib.azure import AzureOpenAI
@@ -155,6 +155,7 @@ class QWenEmbed(Base):
 
     def encode(self, texts: list, batch_size=10):
         import dashscope
+        batch_size = min(batch_size, 4)
         try:
             res = []
             token_count = 0
@@ -403,9 +404,9 @@ class BedrockEmbed(Base):
     def __init__(self, key, model_name,
                  **kwargs):
         import boto3
-        self.bedrock_ak = eval(key).get('bedrock_ak', '')
+        self.bedrock_ak = json.loads(key).get('bedrock_ak', '')
-        self.bedrock_sk = eval(key).get('bedrock_sk', '')
+        self.bedrock_sk = json.loads(key).get('bedrock_sk', '')
-        self.bedrock_region = eval(key).get('bedrock_region', '')
+        self.bedrock_region = json.loads(key).get('bedrock_region', '')
         self.model_name = model_name
         self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
                                    aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
@@ -577,11 +578,40 @@ class UpstageEmbed(OpenAIEmbed):
         super().__init__(key, model_name, base_url)
 
 
-class SILICONFLOWEmbed(OpenAIEmbed):
+class SILICONFLOWEmbed(Base):
-    def __init__(self, key, model_name, base_url="https://api.siliconflow.cn/v1"):
+    def __init__(
+        self, key, model_name, base_url="https://api.siliconflow.cn/v1/embeddings"
+    ):
         if not base_url:
-            base_url = "https://api.siliconflow.cn/v1"
+            base_url = "https://api.siliconflow.cn/v1/embeddings"
-        super().__init__(key, model_name, base_url)
+        self.headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "authorization": f"Bearer {key}",
+        }
+        self.base_url = base_url
+        self.model_name = model_name
+
+    def encode(self, texts: list, batch_size=32):
+        payload = {
+            "model": self.model_name,
+            "input": texts,
+            "encoding_format": "float",
+        }
+        res = requests.post(self.base_url, json=payload, headers=self.headers).json()
+        return (
+            np.array([d["embedding"] for d in res["data"]]),
+            res["usage"]["total_tokens"],
+        )
+
+    def encode_queries(self, text):
+        payload = {
+            "model": self.model_name,
+            "input": text,
+            "encoding_format": "float",
+        }
+        res = requests.post(self.base_url, json=payload, headers=self.headers).json()
+        return np.array(res["data"][0]["embedding"]), res["usage"]["total_tokens"]
 
 
 class ReplicateEmbed(Base):
@@ -623,3 +653,24 @@ class BaiduYiyanEmbed(Base):
             np.array([r["embedding"] for r in res["data"]]),
             res["usage"]["total_tokens"],
         )
+
+
+class VoyageEmbed(Base):
+    def __init__(self, key, model_name, base_url=None):
+        import voyageai
+
+        self.client = voyageai.Client(api_key=key)
+        self.model_name = model_name
+
+    def encode(self, texts: list, batch_size=32):
+        res = self.client.embed(
+            texts=texts, model=self.model_name, input_type="document"
+        )
+        return np.array(res.embeddings), res.total_tokens
+
+    def encode_queries(self, text):
+        res = self.client.embed
+        res = self.client.embed(
+            texts=text, model=self.model_name, input_type="query"
+            )
+        return np.array(res.embeddings), res.total_tokens

rag/llm/rerank_model.py

@@ -311,3 +311,19 @@ class BaiduYiyanRerank(Base):
         rank = np.array([d["relevance_score"] for d in res["results"]])
         indexs = [d["index"] for d in res["results"]]
         return rank[indexs], res["usage"]["total_tokens"]
+
+
+class VoyageRerank(Base):
+    def __init__(self, key, model_name, base_url=None):
+        import voyageai
+
+        self.client = voyageai.Client(api_key=key)
+        self.model_name = model_name
+
+    def similarity(self, query: str, texts: list):
+        res = self.client.rerank(
+            query=query, documents=texts, model=self.model_name, top_k=len(texts)
+        )
+        rank = np.array([r.relevance_score for r in res.results])
+        indexs = [r.index for r in res.results]
+        return rank[indexs], res.total_tokens

rag/llm/sequence2txt_model.py

@@ -22,7 +22,8 @@ from openai import OpenAI
 import os
 import json
 from rag.utils import num_tokens_from_string
-
+import base64
+import re
 
 class Base(ABC):
     def __init__(self, key, model_name):
@@ -36,6 +37,13 @@ class Base(ABC):
         )
         return transcription.text.strip(), num_tokens_from_string(transcription.text.strip())
     
+    def audio2base64(self,audio):
+        if isinstance(audio, bytes):
+            return base64.b64encode(audio).decode("utf-8")
+        if isinstance(audio, io.BytesIO):
+            return base64.b64encode(audio.getvalue()).decode("utf-8")
+        raise TypeError("The input audio file should be in binary format.")
+
 
 class GPTSeq2txt(Base):
     def __init__(self, key, model_name="whisper-1", base_url="https://api.openai.com/v1"):
@@ -63,7 +71,7 @@ class QWenSeq2txt(Base):
         ans = ""
         if result.status_code == HTTPStatus.OK:
             for sentence in result.get_sentence():
-                ans += str(sentence + '\n')
+                ans += sentence.text.decode('utf-8') + '\n'
             return ans, num_tokens_from_string(ans)
 
         return "**ERROR**: " + result.message, 0
@@ -87,3 +95,66 @@ class XinferenceSeq2txt(Base):
     def __init__(self, key, model_name="", base_url=""):
         self.client = OpenAI(api_key="xxx", base_url=base_url)
         self.model_name = model_name
+
+
+class TencentCloudSeq2txt(Base):
+    def __init__(
+        self, key, model_name="16k_zh", base_url="https://asr.tencentcloudapi.com"
+    ):
+        from tencentcloud.common import credential
+        from tencentcloud.asr.v20190614 import asr_client
+
+        key = json.loads(key)
+        sid = key.get("tencent_cloud_sid", "")
+        sk = key.get("tencent_cloud_sk", "")
+        cred = credential.Credential(sid, sk)
+        self.client = asr_client.AsrClient(cred, "")
+        self.model_name = model_name
+
+    def transcription(self, audio, max_retries=60, retry_interval=5):
+        from tencentcloud.common.exception.tencent_cloud_sdk_exception import (
+            TencentCloudSDKException,
+        )
+        from tencentcloud.asr.v20190614 import models
+        import time
+
+        b64 = self.audio2base64(audio)
+        try:
+            # dispatch disk
+            req = models.CreateRecTaskRequest()
+            params = {
+                "EngineModelType": self.model_name,
+                "ChannelNum": 1,
+                "ResTextFormat": 0,
+                "SourceType": 1,
+                "Data": b64,
+            }
+            req.from_json_string(json.dumps(params))
+            resp = self.client.CreateRecTask(req)
+
+            # loop query
+            req = models.DescribeTaskStatusRequest()
+            params = {"TaskId": resp.Data.TaskId}
+            req.from_json_string(json.dumps(params))
+            retries = 0
+            while retries < max_retries:
+                resp = self.client.DescribeTaskStatus(req)
+                if resp.Data.StatusStr == "success":
+                    text = re.sub(
+                        r"\[\d+:\d+\.\d+,\d+:\d+\.\d+\]\s*", "", resp.Data.Result
+                    ).strip()
+                    return text, num_tokens_from_string(text)
+                elif resp.Data.StatusStr == "failed":
+                    return (
+                        "**ERROR**: Failed to retrieve speech recognition results.",
+                        0,
+                    )
+                else:
+                    time.sleep(retry_interval)
+                    retries += 1
+            return "**ERROR**: Max retries exceeded. Task may still be processing.", 0
+
+        except TencentCloudSDKException as e:
+            return "**ERROR**: " + str(e), 0
+        except Exception as e:
+            return "**ERROR**: " + str(e), 0

rag/llm/tts_model.py

@@ -0,0 +1,156 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  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.
+#
+
+from typing import Annotated, Literal
+from abc import ABC
+import httpx
+import ormsgpack
+from pydantic import BaseModel, conint
+from rag.utils import num_tokens_from_string
+import json
+import re
+import time
+class ServeReferenceAudio(BaseModel):
+    audio: bytes
+    text: str
+
+
+class ServeTTSRequest(BaseModel):
+    text: str
+    chunk_length: Annotated[int, conint(ge=100, le=300, strict=True)] = 200
+    # Audio format
+    format: Literal["wav", "pcm", "mp3"] = "mp3"
+    mp3_bitrate: Literal[64, 128, 192] = 128
+    # References audios for in-context learning
+    references: list[ServeReferenceAudio] = []
+    # Reference id
+    # For example, if you want use https://fish.audio/m/7f92f8afb8ec43bf81429cc1c9199cb1/
+    # Just pass 7f92f8afb8ec43bf81429cc1c9199cb1
+    reference_id: str | None = None
+    # Normalize text for en & zh, this increase stability for numbers
+    normalize: bool = True
+    # Balance mode will reduce latency to 300ms, but may decrease stability
+    latency: Literal["normal", "balanced"] = "normal"
+
+
+class Base(ABC):
+    def __init__(self, key, model_name, base_url):
+        pass
+
+    def tts(self, audio):
+        pass
+    
+    def normalize_text(self, text):
+        return re.sub(r'(\*\*|##\d+\$\$|#)', '', text)
+
+
+class FishAudioTTS(Base):
+    def __init__(self, key, model_name, base_url="https://api.fish.audio/v1/tts"):
+        if not base_url:
+            base_url = "https://api.fish.audio/v1/tts"
+        key = json.loads(key)
+        self.headers = {
+            "api-key": key.get("fish_audio_ak"),
+            "content-type": "application/msgpack",
+        }
+        self.ref_id = key.get("fish_audio_refid")
+        self.base_url = base_url
+
+    def tts(self, text):
+        from http import HTTPStatus
+
+        text = self.normalize_text(text)
+        request = ServeTTSRequest(text=text, reference_id=self.ref_id)
+
+        with httpx.Client() as client:
+            try:
+                with client.stream(
+                    method="POST",
+                    url=self.base_url,
+                    content=ormsgpack.packb(
+                        request, option=ormsgpack.OPT_SERIALIZE_PYDANTIC
+                    ),
+                    headers=self.headers,
+                    timeout=None,
+                ) as response:
+                    if response.status_code == HTTPStatus.OK:
+                        for chunk in response.iter_bytes():
+                            yield chunk
+                    else:
+                        response.raise_for_status()
+
+                yield num_tokens_from_string(text)
+
+            except httpx.HTTPStatusError as e:
+                raise RuntimeError(f"**ERROR**: {e}")
+
+
+class QwenTTS(Base):
+    def __init__(self, key, model_name, base_url=""):
+        import dashscope
+        
+        self.model_name = model_name
+        dashscope.api_key = key
+
+    def tts(self, text):
+        from dashscope.api_entities.dashscope_response import SpeechSynthesisResponse
+        from dashscope.audio.tts import ResultCallback, SpeechSynthesizer, SpeechSynthesisResult
+        from collections import deque
+        
+        class Callback(ResultCallback):
+            def __init__(self) -> None:
+                self.dque = deque()
+                   
+            def _run(self):
+                while True:
+                    if not self.dque:
+                        time.sleep(0)
+                        continue
+                    val = self.dque.popleft()
+                    if val:
+                        yield val
+                    else:
+                        break
+
+            def on_open(self):
+                pass
+
+            def on_complete(self):
+                self.dque.append(None)
+
+            def on_error(self, response: SpeechSynthesisResponse):
+                raise RuntimeError(str(response))
+
+            def on_close(self):
+                pass
+
+            def on_event(self, result: SpeechSynthesisResult):
+                if result.get_audio_frame() is not None:
+                    self.dque.append(result.get_audio_frame())
+
+        text = self.normalize_text(text)
+        callback = Callback()
+        SpeechSynthesizer.call(model=self.model_name,
+                                text=text,
+                                callback=callback,
+                                format="mp3")
+        try:
+            for data in callback._run():
+                yield data
+            yield num_tokens_from_string(text)
+            
+        except Exception as e:
+            raise RuntimeError(f"**ERROR**: {e}")

rag/nlp/__init__.py

@@ -214,7 +214,7 @@ def is_english(texts):
     eng = 0
     if not texts: return False
     for t in texts:
-        if re.match(r"[a-zA-Z]{2,}", t.strip()):
+        if re.match(r"[ `a-zA-Z.,':;/\"?<>!\(\)-]", t.strip()):
             eng += 1
     if eng / len(texts) > 0.8:
         return True