fix chunking

src/.env.template

@@ -5,4 +5,3 @@ QIANFAN_ACCESS_KEY=optional
 QIANFAN_SECRET_KEY=optional
 HF_ENDPOINT='https://hf-mirror.com'
 CUDA_VISIBLE_DEVICES=0
-MODEL_ROOT_DIR=MODEL_ROOT_DIR_UNSET

src/core/database.py

@@ -115,21 +115,16 @@ class DataBaseManager:
             db.files.append(new_file)
             new_files.append(new_file)
 
-        from src.core.indexing import chunk_file, chunk_text
+        from src.core.indexing import chunk
         for new_file in new_files:
             file_id = new_file["file_id"]
             idx = self.get_idx_by_fileid(db, file_id)
             db.files[idx]["status"] = "processing"
 
             try:
-                if new_file["type"] in ["txt", "docx", "md"]:
+                nodes = chunk(new_file["path"], params=params)
-                    nodes = chunk_file(new_file["path"], params=params)
-                else:
-                    text = self.read_text(new_file["path"])
-                    nodes = chunk_text(text, params)
-
                 self.knowledge_base.add_documents(
-                    docs=[node.to_dict() for node in nodes],
+                    docs=[node.text for node in nodes],
                     collection_name=db.metaname,
                     file_id=file_id)
 

src/core/indexing.py

@@ -1,37 +1,38 @@
 import os
 from pathlib import Path
+from llama_index.core import Document
 from llama_index.core.node_parser import SimpleFileNodeParser
-from llama_index.readers.file import FlatReader
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.readers.file import FlatReader, DocxReader
 
 from src.utils import hashstr
 
-def chunk_text(text, params=None):
+def chunk(text_or_path, params=None):
     params = params or {}
-    from llama_index.core import Document
-    from llama_index.core.node_parser import SentenceSplitter
     chunk_size = int(params.get("chunk_size", 500))
     chunk_overlap = int(params.get("chunk_overlap", 20))
     splitter = SentenceSplitter(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,
     )
-    doc = Document(id_=hashstr(text), text=text)
-    nodes = splitter.get_nodes_from_documents([doc])
-    return nodes
 
-def chunk_file(file, params=None):
+    if os.path.isfile(text_or_path) and "uploads" in text_or_path:
-    parser = SimpleFileNodeParser()
+        parser = SimpleFileNodeParser()
-    if file.endswith(".txt"):
+        file_type = Path(text_or_path).suffix
-        from llama_index.readers.file import FlatReader
+        if file_type in [".txt", ".json", ".md"]:
-        docs = FlatReader().load_data(Path(file))
+            docs = FlatReader().load_data(Path(text_or_path))
-    elif file.endswith(".docx"):
+        elif file_type in [".docx"]:
-        from llama_index.readers.file import DocxReader
+            docs = DocxReader().load_data(Path(text_or_path))
-        docs = DocxReader().load_data(Path(file))
+        else:
-    elif file.endswith(".md"):
+            raise ValueError(f"Unsupported file type `{file_type}`")
-        from llama_index.readers.file import MarkdownReader
+
-        docs = MarkdownReader().load_data(Path(file))
+        if params.get("use_parser"):
+            nodes = parser.get_nodes_from_documents(docs)
+        else:
+            nodes = splitter.get_nodes_from_documents(docs)
+
     else:
-        raise ValueError("Unsupported file type")
+        docs = [Document(id_=hashstr(text_or_path), text=text_or_path)]
+        nodes = splitter.get_nodes_from_documents(docs)
 
-    nodes = parser.get_nodes_from_documents(docs)
+    return nodes
-    return nodes  # 返回节点

src/views/tools_view.py

@@ -16,21 +16,10 @@ def route_index():
     tools = [
         {
             "name": "text_chunking",
-            "title": "Text Chunking",
+            "title": "文本分块",
-            "description": "Chunking text into smaller pieces for better understanding.",
+            "description": "将文本分块以更好地理解。可以输入文本或者上传文件。",
             "url": "/tools/text_chunking",
             "method": "POST",
-            "params": [
-                {
-                    "name": "text",
-                    "type": "string",
-                    "description": "Text to be chunked."
-                },
-                {
-                    "name": "chunk_size",
-                    "type": "int",
-                }
-            ]
         }
     ]
 
@@ -39,7 +28,7 @@ def route_index():
 
 @tools.route("/text_chunking", methods=["POST"])
 def text_chunking():
-    from src.core.indexing import chunk_text
+    from src.core.indexing import chunk
     text = request.json.get("text")
-    nodes = chunk_text(text, params=request.json)
+    nodes = chunk(text, params=request.json)
     return jsonify({"nodes": [node.to_dict() for node in nodes]})

web/src/assets/base.css

@@ -25,9 +25,9 @@
   --gray-500: #A7ACAF;
   --gray-400: #C2C7CA;
   --gray-300: #DEE3E6;
-  --gray-200: #EDF1F4;
+  --gray-200: #EFF1F4;
-  --gray-100: #F5FAFD;
+  --gray-100: #F8FAFD;
-  --gray-50: #F9FDFF;
+  --gray-50: #FBFDFF;
 
   --main-color: #1c6586;
   --main-color-dark: #004d5c;

web/src/components/TextChunkingComponent.vue

@@ -11,10 +11,13 @@
           @finish="chunkText"
         >
           <a-form-item label="Chunk Size" name="chunkSize" >
-            <a-input v-model:value="params.chunkSize" />
+            <a-input v-model:value="params.chunkSize" :disabled="params.useParser" />
           </a-form-item>
           <a-form-item label="Chunk Overlap" name="chunkOverlap" >
-            <a-input v-model:value="params.chunkOverlap" />
+            <a-input v-model:value="params.chunkOverlap" :disabled="params.useParser" />
+          </a-form-item>
+          <a-form-item label="使用文件节点解析器" name="useParser" v-if="state.useFile">
+            <a-switch v-model:checked="params.useParser" />
           </a-form-item>
 
           <a-form-item>
@@ -26,16 +29,38 @@
     </div>
     <div class="result-container">
       <div class="input-container">
-        <textarea v-model="text" placeholder="Enter text to chunk"></textarea>
+        <div class="actions">
-        <div class="infos">
+          <span :class="{'active': !state.useFile}" @click="state.useFile = false">上传文本</span>
+          <span :class="{'active': state.useFile}" @click="state.useFile = true">上传文件</span>
+        </div>
+        <div class="upload" v-if="state.useFile">
+          <a-upload-dragger
+            class="upload-dragger"
+            v-model:fileList="fileList"
+            name="file"
+            :max-count="1"
+            :disabled="state.uploading"
+            action="/api/database/upload"
+            @change="handleFileUpload"
+            @drop="handleDrop"
+          >
+            <p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
+            <p class="ant-upload-hint">
+              目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加
+            </p>
+          </a-upload-dragger>
+        </div>
+        <textarea v-if="!state.useFile" v-model="text" placeholder="Enter text to chunk"></textarea>
+        <div class="infos" v-if="!state.useFile">
           <!-- <span>字数: {{ wordCount }}</span> -->
           <span>字符数: {{ charCount }}</span>
           <span>Token 数：{{ estimatedTokenCount }}</span>
         </div>
       </div>
+      <!-- <div>{{ chunks[0] }}</div> -->
       <div id="result-cards" class="result-cards">
         <div v-for="(chunk, index) in chunks" :key="index" class="chunk">
-          <p>{{ chunk }}</p>
+          <p><strong>#{{ index + 1 }}</strong> {{ chunk.text }}</p>
         </div>
       </div>
     </div>
@@ -49,13 +74,17 @@ import { message } from 'ant-design-vue'
 const text = ref('');
 const state = reactive({
   loading: true,
+  uploading: false,
+  useFile: false,
 })
 
 const params = reactive({
   chunkSize: 500,
-  chunkOverlap: 20
+  chunkOverlap: 20,
+  useParser: false,
 })
 const chunks = ref([]);
+const fileList = ref([]);
 
 const wordCount = computed(() => text.value.split(/\s+/).filter(Boolean).length);
 const charCount = computed(() => text.value.length)
@@ -78,19 +107,32 @@ const estimatedTokenCount = computed(() => {
 })
 
 const chunkText = async () => {
-  if (text.value.length === 0) {
+  let text_or_file = ''
-    message.error("请输入文本")
+  if (state.useFile) {
-    return
+    if (fileList.value.length === 0) {
+      message.error("请上传文件")
+      return
+    }
+    console.log(fileList.value)
+    text_or_file = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path)[0]
+  } else {
+    if (text.value.length === 0) {
+      message.error("请输入文本")
+      return
+    }
+    text_or_file = text.value
   }
+
   try {
     state.loading = true
     const response = await fetch('/api/tools/text_chunking', {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify({
-        text: text.value,
+        text: text_or_file,
         chunk_size: params.chunkSize,
-        chunk_overlap: params.chunkOverlap
+        chunk_overlap: params.chunkOverlap,
+        use_parser: params.useParser
       })
     });
 
@@ -99,7 +141,7 @@ const chunkText = async () => {
     }
 
     const data = await response.json();
-    chunks.value = data.nodes.map(node => node.text);
+    chunks.value = data.nodes;
     state.loading = false
   } catch (error) {
     console.error('Error chunking text:', error);
@@ -141,6 +183,28 @@ const chunkText = async () => {
       flex-direction: column;
       margin-bottom: 15px;
 
+      .actions {
+        display: flex;
+        justify-content: space-between;
+        margin-bottom: 10px;
+        width: fit-content;
+        background-color: var(--gray-200);
+        padding: 8px 6px;
+        border-radius: 8px;
+
+        span {
+          color: var(--gray-900);
+          cursor: pointer;
+          padding: 4px 10px;
+          border-radius: 8px;
+          transition: background-color 0.3s;
+
+          &.active {
+            background-color: white;
+          }
+        }
+      }
+
       textarea {
         resize: vertical;
         height: 300px;
@@ -149,7 +213,7 @@ const chunkText = async () => {
         border-radius: 8px;
         font-size: 1rem;
         transition: border-color 0.3s;
-        background-color: var(--gray-50);
+        background-color: var(--gray-100);
 
         &:focus {
           border-color: var(--main-color);
@@ -174,38 +238,36 @@ const chunkText = async () => {
     }
 
     .chunk {
-      background-color: #f9fcff;
+      background-color: var(--main-5);
-      border-radius: 4px;
+      border: 1px solid var(--main-light-3);
-      padding: 10px;
+      border-radius: 8px;
+      padding: 16px;
       margin-bottom: 10px;
-      box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+      box-shadow: 0 0 10px 2px rgba(0, 0, 0, 0.01);
       break-inside: avoid; /* 避免元素被分割到不同列 */
+      // 强制换行
+      word-wrap: break-word;
+      word-break: break-all;
     }
   }
 }
 
 
-@media (max-width: 768px) {
+@media (max-width: 980px) {
   #result-cards {
     column-count: 1;
   }
 }
 
-@media (min-width: 769px) and (max-width: 1200px) {
+@media (min-width: 981px) and (max-width: 1500px) {
   #result-cards {
     column-count: 2;
   }
 }
 
-@media (min-width: 1201px) and (max-width: 1900px) {
+@media (min-width: 1501px){
   #result-cards {
     column-count: 3;
   }
 }
-
-@media (min-width: 1901px){
-  #result-cards {
-    column-count: 4;
-  }
-}
 </style>

web/src/views/ToolsView.vue

@@ -107,6 +107,7 @@ onMounted(() => {
 
       p {
         margin: 0;
+        color: var(--gray-800);
       }
     }
   }