Merge pull request #12 from xerrors/bigchange

Bigchange Merge

README.md

@@ -15,7 +15,7 @@
 
 ## Dockers 启动
 
-**提醒**：此部分暂时依赖于前端打包之后的内容（后面考虑更新），同时会自动启动 neo4j 图数据库。
+**提醒**：下面的脚本会启动开发版本，源代码的修改会自动更新（含前端和后端）。如果生产环境部署，请使用 ` docker/docker-compose.yml` 启动。
 
 ```bash
 docker-compose -f docker/docker-compose.dev.yml up --build
@@ -40,98 +40,37 @@ docker-compose -f docker/docker-compose.dev.yml up --build
 docker-compose -f docker/docker-compose.dev.yml down
 ```
 
+查看日志：
+
+```bash
+docker logs <CONTAINER_NAME>  # 例如：docker logs api-dev
+```
+
 如果需要使用到本地模型，比如向量模型或者重排序模型，则需要将环境变量中设置的 `MODEL_ROOT_DIR` 做映射，比如本地模型都是存放在 `/hdd/models` 里面，则需要在 `docker-compose.yml` 中添加：
 
 ```yml
 services:
-  # 后端服务
+  api:
-  backend:
+    build:
-    image: pytorch/pytorch:2.4.1-cuda11.8-cudnn9-runtime  # 或者您可以自定义 Python 基础镜像
+      context: ..
-    container_name: backend
+      dockerfile: docker/api.Dockerfile
+    container_name: api-dev
     working_dir: /app
     volumes:
-      - ./src:/app/src  # 映射源代码
+      - ../src:/app/src
-      - ./requirements.txt:/app/requirements.txt
+      - ../saves:/app/saves
-      - ./saves:/app/saves
+      - /hdd/zwj/models:/hdd/zwj/models  # <== 修改这一行
-      - /hdd/models:/hdd/models  # <=== 修改这里
-...
 ```
 
-##
+## 更新日志
-
-1. 配置 python 环境 `pip install -r requirements.txt`，python 版本应当小于 `3.12`。
-2. 前端 UI 部分，需要安装 Node.js  环境，参考：[Download Node.js](https://nodejs.org/en/download/package-manager)。
-
-**如果不启用知识库，可以仅安装下面的依赖**
-
-```
-FlagEmbedding==1.2.10
-Flask==3.0.3
-Flask_Cors==4.0.1
-openai==1.35.10
-python-dotenv==1.0.1
-PyYAML==6.0.1
-zhipuai
-```
-
-> [!WARNING]
-> milvus-lite（向量数据库）不支持在 Windows 上运行，建议使用 WSL，详见：https://github.com/milvus-io/milvus-lite/issues/175
-
-### 配置图数据库 neo4j (可选)
-
-使用 docker 部署 neo4j 服务，配置文件见 [local_neo4j/docker-compose.yml](local_neo4j/docker-compose.yml).
-默认账号密码见最后一行，可以使用 `http://localhost:7474/` 在浏览器可视化访问。
-
-```bash
-cd local_neo4j
-docker compose up -d
-```
-
-可以使用 `python test_neo4j.py` 来测试是否正常启动。使用 `docker compose down` 可停止服务。
-如果想要管理 neo4j，也可以使用 `docker ps` 查看容器 id，然后使用 `docker exec -it  <CONTAINER_ID> /bin/bash` 进入容器。
-如果想要删除数据库中的文件，可以进入容器并停止 neo4j 后，执行 `rm -rf /data/databases`。
-
-## 启动
-
-推荐使用 docker 启动
-
-### 1. 手动启动
-
-```bash
-# 后端部分
-python -m src.main
-
-# 前端部分
-cd web
-npm install  # 首次运行需要
-npm run server
-```
-
-### 2. 脚本启动
-
-注意：此脚本不会启动图数据库 neo4j。
-
-```bash
-bash run.sh
-```
-
-### 3. Docker 启动
-
-## Changelog
 
 - 2024.10.12 后端修改为 FastAPI，并添加了 Milnvs 的独立部署。
 
-## 其余脚本
+## 相关问题
 
-镜像构建之后，单独启动后端服务
+### 镜像下载问题
 
-```bash
+如果无法直接下载相关镜像，参考 [DaoCloud/public-image-mirror](https://github.com/DaoCloud/public-image-mirror?tab=readme-ov-file#%E5%BF%AB%E9%80%9F%E5%BC%80%E5%A7%8B)，尝试将前缀替换为：
-docker run -w /app -v ./src:/app/src -v ./saves:/app/saves -p 5000:5000 yuxi-know-backend
-```
-
-## Docker 镜像相关
-
-如果无法直接下载相关镜像，尝试将前缀替换为：
 
 ```bash
 # 以 neo4j 为例，其余一样

src/models/README.md

@@ -13,24 +13,11 @@
 |`deepseek`|`deepseek-chat`|`DEEPSEEK_API_KEY`|
 |`siliconflow` | `meta-llama/Meta-Llama-3.1-8B-Instruct` | `SILICONFLOW_API_KEY`|
 
-vllm 的具体配置项可以参考[这里](https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#named-arguments), 部署参考脚本：
+同样支持以 OpenAI 的兼容模型运行模型，可以直接在设置里面添加。比如使用 vllm 和 Ollama 运行本地模型时。
-
-```bash
-python -m vllm.entrypoints.openai.api_server \
-	--model="/hdd/models/meta-llama/Meta-Llama-3.1-8B-Instruct" \
-	--tensor-parallel-size 1 \
-	--trust-remote-code \
-	--device auto \
-	--gpu-memory-utilization 0.98 \
-	--dtype half \
-	--served-model-name "vllm" \
-	--host 0.0.0.0 \
-	--port 8080
-```
 
 ### 2. 向量模型支持
 
-建议直接使用智谱 AI 的 embedding-3。
+建议直接使用智谱 AI 的 embedding-3，这样不需要做任何修改，且资费不贵。
 
 > [!Warning]
 > 需要注意，由于知识库和图数据库的构建都依赖于向量模型，如果中途更改向量模型，回导致知识库不可用。此外，知识图谱的向量索引的建立默认使用 embedding-3 构建，因此检索的时候必须使用 embedding-3（现阶段还不支持修改）
@@ -38,22 +25,23 @@ python -m vllm.entrypoints.openai.api_server \
 
 |模型名称(`config.embed_model`)|默认路径/模型|需要配置项目（`config.model_local_paths`）|
 |:-|:-|:-|
-|`bge-large-zh-v1.5`|`BAAI/bge-large-zh-v1.5`|`bge-large-zh-v1.5`（*修改为本地路径）|
+|`bge-large-zh-v1.5`|`BAAI/bge-large-zh-v1.5`|`bge-large-zh-v1.5`（*可选：修改为本地路径）|
 |`zhipu`|`embedding-2`, `embedding-3`|`ZHIPUAI_API_KEY` (`.env`)|
 
 
-例如（`saves/config/config.yaml`）：
-
-```yaml
-model_provider: qianfan
-model_name: null # for default
-
-## model dir 可以写**相对路径**和**绝对路径**
-### 相对路径是相对于环境变量 (.env) 中 MODEL_ROOT_DIR（若为空则是相对于 `pretrained_models`) 的路径
-model_local_paths:
-  bge-large-zh-v1.5: /models/bge-large-zh-v1.5
-```
-
 ### 3. 重排序模型支持
 
 目前仅支持 `BAAI/bge-reranker-v2-m3`。
+
+### 4. 本地模型支持
+
+对于**语言模型**，并不支持直接运行本地语言模型，请使用 vllm 或者 ollama 转成 API 服务之后使用。
+
+对于**向量模型**和**重排序模型**，可以不做修改会自动下载模型，如果下载过程中出现问题，请参考 [HF-Mirror](https://hf-mirror.com/) 配置相关内容。如果想要使用本地已经下载好的模型（不建议），可以在 `saves/config/config.yaml` 配置相关内容。同时注意要在 docker 中做映射，参考 README 中的 `docker/docker-compose.yml`。
+
+例如：
+
+```yaml
+model_local_paths:
+  bge-large-zh-v1.5: /models/bge-large-zh-v1.5
+```

src/routers/chat_router.py

@@ -28,14 +28,30 @@ def chat_post(
         cur_res_id: str = Body(...)):
 
     history_manager = HistoryManager(history)
-    new_query, refs = startup.retriever(query, history_manager.messages, meta)
-    refs_pool[cur_res_id] = refs
 
-    messages = history_manager.get_history_with_msg(new_query, max_rounds=meta.get('history_round'))
+    def make_chunk(content, status, history):
-    history_manager.add_user(query)
+        return json.dumps({
-    logger.debug(f"Web history: {history_manager.messages}")
+            "response": content,
+            "history": history,
+            "model_name": startup.config.model_name,
+            "status": status,
+        }, ensure_ascii=False).encode('utf-8') + b"\n"
 
     def generate_response():
+
+        if meta.get("enable_retrieval"):
+            chunk = make_chunk("", "searching", history=None)
+            yield chunk
+
+            new_query, refs = startup.retriever(query, history_manager.messages, meta)
+            refs_pool[cur_res_id] = refs
+        else:
+            new_query = query
+
+        messages = history_manager.get_history_with_msg(new_query, max_rounds=meta.get('history_round'))
+        history_manager.add_user(query)
+        logger.debug(f"Web history: {history_manager.messages}")
+
         content = ""
         for delta in startup.model.predict(messages, stream=True):
             if not delta.content:
@@ -47,12 +63,8 @@ def chat_post(
                 content += delta.content
 
             logger.debug(f"Response: {content}")
-
+            chunk = make_chunk(content, "loading", history=history_manager.update_ai(content))
-            _chunk = json.dumps({
+            yield chunk
-                "response": content,
-                "history": history_manager.update_ai(content),
-            }, ensure_ascii=False).encode('utf-8') + b"\n"
-            yield _chunk
 
     return StreamingResponse(generate_response(), media_type='application/json')
 

web/src/components/ChatComponent.vue

@@ -16,10 +16,10 @@
         </a-tooltip>
       </div>
       <div class="header__right">
-        <div class="nav-btn text metas" v-if="meta.use_graph">
+        <div class="nav-btn text metas" v-if="meta.use_graph && meta.enable_retrieval">
           <GoldOutlined /> 图数据库
         </div>
-        <a-dropdown v-if="meta.selectedKB !== null">
+        <a-dropdown v-if="meta.selectedKB !== null && meta.enable_retrieval">
           <a class="ant-dropdown-link nav-btn" @click.prevent>
             <!-- <component :is="meta.selectedKB === null ? BookOutlined : BookFilled" /> -->
              <BookOutlined />
@@ -40,7 +40,20 @@
           <component :is="opts.showPanel ? FolderOpenOutlined : FolderOutlined" /> <span class="text">选项</span>
         </div>
         <div v-if="opts.showPanel" class="my-panal r0 top100 swing-in-top-fwd" ref="panel">
-          <div class="flex-center" v-if="configStore.config.enable_knowledge_base">
+          <div class="flex-center" @click="meta.stream = !meta.stream">
+            流式输出 <div @click.stop><a-switch v-model:checked="meta.stream" /></div>
+          </div>
+          <div class="flex-center" @click="meta.summary_title = !meta.summary_title">
+            总结对话标题 <div @click.stop><a-switch v-model:checked="meta.summary_title" /></div>
+          </div>
+          <div class="flex-center" @click="meta.enable_retrieval = !meta.enable_retrieval">
+            启用检索 <div @click.stop><a-switch v-model:checked="meta.enable_retrieval" /></div>
+          </div>
+          <div class="flex-center">
+            最大历史轮数 <a-input-number id="inputNumber" v-model:value="meta.history_round" :min="1" :max="50" />
+          </div>
+          <a-divider v-if="meta.enable_retrieval"></a-divider>
+          <div class="flex-center" v-if="configStore.config.enable_knowledge_base && meta.enable_retrieval">
             知识库
             <div @click.stop>
               <a-dropdown>
@@ -62,24 +75,15 @@
               </a-dropdown>
             </div>
           </div>
-          <div class="flex-center" @click="meta.use_graph = !meta.use_graph" v-if="configStore.config.enable_knowledge_base">
+          <div class="flex-center" @click="meta.use_graph = !meta.use_graph" v-if="configStore.config.enable_knowledge_base && meta.enable_retrieval">
             图数据库 <div @click.stop><a-switch v-model:checked="meta.use_graph" /></div>
           </div>
-          <div class="flex-center" @click="meta.use_web = !meta.use_web" v-if="configStore.config.enable_search_engine">
+          <div class="flex-center" @click="meta.use_web = !meta.use_web" v-if="configStore.config.enable_search_engine && meta.enable_retrieval">
             搜索引擎（Bing） <div @click.stop><a-switch v-model:checked="meta.use_web" /></div>
           </div>
-          <div class="flex-center" @click="meta.stream = !meta.stream">
+          <div class="flex-center" v-if="configStore.config.enable_knowledge_base && meta.enable_retrieval">
-            流式输出 <div @click.stop><a-switch v-model:checked="meta.stream" /></div>
-          </div>
-          <div class="flex-center" @click="meta.summary_title = !meta.summary_title">
-            总结对话标题 <div @click.stop><a-switch v-model:checked="meta.summary_title" /></div>
-          </div>
-          <div class="flex-center" v-if="configStore.config.enable_knowledge_base">
             重写查询 <a-segmented v-model:value="meta.rewriteQuery" :options="['off', 'on', 'hyde']"/>
           </div>
-          <div class="flex-center">
-            最大历史轮数 <a-input-number id="inputNumber" v-model:value="meta.history_round" :min="1" :max="50" />
-          </div>
         </div>
       </div>
     </div>
@@ -104,12 +108,19 @@
         :class="message.role"
       >
         <p v-if="message.role=='sent'" style="white-space: pre-line" class="message-text">{{ message.text }}</p>
-        <div v-else-if="message.text.length == 0 && message.status=='querying'"  class="loading-dots">
+        <div v-else-if="message.text.length == 0 && message.status=='init'"  class="loading-dots">
           <div></div>
           <div></div>
           <div></div>
         </div>
-        <div v-else-if="message.text.length == 0 || message.status == 'error'" class="err-msg">请求错误，请重试</div>
+        <div v-else-if="message.status == 'searching' && isStreaming" class="searching-msg"><i>正在检索……</i></div>
+        <div
+          v-else-if="message.text.length == 0 || message.status == 'error' || (message.status != 'finished' && !isStreaming)"
+          class="err-msg"
+          @click="retryMessage(message.id)"
+        >
+          请求错误，请重试
+        </div>
         <div v-else
           v-html="renderMarkdown(message)"
           class="message-md"
@@ -195,6 +206,7 @@ const opts = reactive({
 })
 
 const meta = reactive(JSON.parse(localStorage.getItem('meta')) || {
+  enable_retrieval: false,
   use_graph: false,
   use_web: false,
   graph_name: "neo4j",
@@ -306,33 +318,32 @@ const appendAiMessage = (message, refs=null) => {
     role: 'received',
     text: message,
     refs,
-    status: "querying",
+    status: "init",
   })
   scrollToBottom()
 }
 
-const updateMessage = (text, id, refs, status) => {
+const updateMessage = (info) => {
-  const message = conv.value.messages.find((message) => message.id === id);
+  const message = conv.value.messages.find((message) => message.id === info.id);
 
   if (message) {
     // 只有在 text 不为空时更新
-    if (text !== null && text !== undefined && text !== '') {
+    if (info.text !== null && info.text !== undefined && info.text !== '') {
-      message.text = text;
+      message.text = info.text;
     }
 
     // 只有在 refs 不为空时更新
-    if (refs !== null && refs !== undefined) {
+    if (info.refs !== null && info.refs !== undefined) {
-      message.refs = refs;
+      message.refs = info.refs;
+    }
 
-      // 如果 refs 里面的 model_name 不为空时更新
+    if (info.model_name !== null && info.model_name !== undefined && info.model_name !== '') {
-      if (refs.model_name !== null && refs.model_name !== undefined && refs.model_name !== '') {
+      message.model_name = info.model_name;
-        message.model_name = refs.model_name;
-      }
     }
 
     // 只有在 status 不为空时更新
-    if (status !== null && status !== undefined && status !== '') {
+    if (info.status !== null && info.status !== undefined && info.status !== '') {
-      message.status = status;
+      message.status = info.status;
     }
   } else {
     console.error('Message not found');
@@ -412,11 +423,21 @@ const fetchChatResponse = (user_input, cur_res_id) => {
     const readChunk = () => {
       return reader.read().then(({ done, value }) => {
         if (done) {
-          fetchRefs(cur_res_id).then((data) => {
+          const message = conv.value.messages.find((message) => message.id === cur_res_id)
-            console.log(data)
+          if (message.refs && message.refs.meta.enable_retrieval) {
-            updateMessage(null, cur_res_id, data, "finished");
+            console.log("fetching refs")
+            fetchRefs(cur_res_id).then((data) => {
+              console.log(data)
+              updateMessage({
+                id: cur_res_id,
+                refs: data,
+                status: "finished",
+              });
+              updateStatus(cur_res_id, "finished");
+            })
+          } else {
             updateStatus(cur_res_id, "finished");
-          })
+          }
           isStreaming.value = false;
           if (conv.value.messages.length === 2) { renameTitle(); }
           return;
@@ -426,9 +447,16 @@ const fetchChatResponse = (user_input, cur_res_id) => {
         buffer += chunk;
         try {
           const data = JSON.parse(chunk);
-          updateMessage(data.response, cur_res_id, data.refs, "loading");
+          updateMessage({
+            id: cur_res_id,
+            text: data.response,
+            model_name: data.model_name,
+            status: data.status,
+          });
           console.debug(data.response)
-          conv.value.history = data.history;
+          if (data.history) {
+            conv.value.history = data.history;
+          }
         } catch (e) {
           // console.debug('JSON 解析错误:', e, chunk);
         }
@@ -444,9 +472,6 @@ const fetchChatResponse = (user_input, cur_res_id) => {
     updateStatus(cur_res_id, "error");
     isStreaming.value = false;
   })
-  .finally(() => {
-    isStreaming.value = false;
-  });
 }
 
 const fetchRefs = (cur_res_id) => {
@@ -484,6 +509,19 @@ const sendMessage = () => {
   }
 }
 
+const retryMessage = (id) => {
+  // 找到 id 对应的 message，然后删除包含 message 在内以及后面所有的 message
+  console.log("retryMessage", id)
+  const index = conv.value.messages.findIndex(message => message.id === id);
+  const pastMessage = conv.value.messages[index-1]
+  conv.value.inputText = pastMessage.text
+  if (index !== -1) {
+    conv.value.messages = conv.value.messages.slice(0, index-1);
+  }
+  console.log(conv.value.messages)
+  sendMessage();
+}
+
 const autoSend = (message) => {
   conv.value.inputText = message
   sendMessage()
@@ -677,6 +715,12 @@ watch(
       border-radius: 8px;
       text-align: center;
       background: #FFF0F0;
+      margin-bottom: 10px;
+      cursor: pointer;
+    }
+
+    .searching-msg {
+      color: var(--gray-500);
     }
   }