Merge pull request #69 from xerrors/68-database-chunk-info

添加 chunk 信息

src/core/database.py

@@ -1,7 +1,8 @@
 import os
 import json
 import time
-from src.utils import hashstr, logger, is_text_pdf
+from src.utils import hashstr, logger
+from src.core.indexing import chunk
 from src.models.embedding import get_embedding_model
 
 
@@ -140,7 +141,6 @@ class DataBaseManager:
         # 先保存一次数据库状态，确保waiting状态被记录
         self._save_databases()
 
-        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)
@@ -156,9 +156,10 @@ class DataBaseManager:
                     nodes = chunk(new_file["path"], params=params)
 
                 self.knowledge_base.add_documents(
-                    docs=[node.text for node in nodes],
+                    file_id=file_id,
                     collection_name=db.metaname,
-                    file_id=file_id)
+                    docs=[node.text for node in nodes],
+                    chunk_infos=[node.dict() for node in nodes])
 
                 idx = self.get_idx_by_fileid(db, file_id)
                 db.files[idx]["status"] = "done"
@@ -219,20 +220,16 @@ class DataBaseManager:
         lines = self.knowledge_base.client.query(
             collection_name=db.metaname,
             filter=f"file_id == '{file_id}'",
-            output_fields=["id", "text", "file_id", "hash"]
+            output_fields=None
         )
+        # 删除 vector 字段
+        for line in lines:
+            line.pop("vector")
+
+        lines.sort(key=lambda x: x.get("start_char_idx", 0))
+        logger.debug(f"lines[0]: {lines[0]}")
         return {"lines": lines}
 
-    def chunking(self, text, params=None):
-        chunk_method = params.get("chunk_method", "fixed")
-        chunk_size = params.get("chunk_size", 500)
-
-        """将文本切分成固定大小的块"""
-        chunks = []
-        for i in range(0, len(text), chunk_size):
-            chunks.append(text[i:i + chunk_size])
-        return chunks
-
     def delete_database(self, db_id):
         db = self.get_kb_by_id(db_id)
         if db is None:

src/core/indexing.py

@@ -8,9 +8,21 @@ from llama_index.readers.file import FlatReader, DocxReader
 from src.utils import hashstr
 
 def chunk(text_or_path, params=None):
+    """
+    将文本或文件切分成固定大小的块
+    
+    Args:
+        text_or_path: 文本或文件路径
+        params: 参数
+            chunk_size: 块大小
+            chunk_overlap: 块重叠大小
+            use_parser: 是否使用文件解析器
+    Returns:
+        nodes: 节点列表
+    """
     params = params or {}
     chunk_size = int(params.get("chunk_size", 500))
-    chunk_overlap = int(params.get("chunk_overlap", 20))
+    chunk_overlap = int(params.get("chunk_overlap", 100))
     splitter = SentenceSplitter(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,

src/core/knowledgebase.py

@@ -59,7 +59,7 @@ class KnowledgeBase:
             dimension= dimension,  # The vectors we will use in this demo has 768 dimensions
         )
 
-    def add_documents(self, docs, collection_name, **kwargs):
+    def add_documents(self, docs, collection_name, chunk_infos=None, **kwargs):
         """添加已经分块之后的文本"""
         # 检查 collection 是否存在
         import random
@@ -67,6 +67,8 @@ class KnowledgeBase:
             logger.error(f"Collection {collection_name} not found, create it")
             # self.add_collection(collection_name)
 
+        chunk_infos = chunk_infos or [{}] * len(docs)
+
         vectors = self.embed_model.batch_encode(docs)
 
         data = [{
@@ -74,7 +76,9 @@ class KnowledgeBase:
             "vector": vectors[i],
             "text": docs[i],
             "hash": hashstr(docs[i], with_salt=True),
-            **kwargs} for i in range(len(vectors))]
+            **kwargs,
+            **chunk_infos[i]
+        } for i in range(len(vectors))]
 
         res = self.client.insert(collection_name=collection_name, data=data)
         return res

web/src/views/DataBaseInfoView.vue

@@ -94,8 +94,8 @@
             @after-open-change="afterOpenChange"
           >
             <h2>共 {{ selectedFile?.lines.length }} 个片段</h2>
-            <p v-for="line in selectedFile?.lines" :key="line.id">
+            <p v-for="line in selectedFile?.lines" :key="line.id" class="line-text">
-              <strong>Chunk #{{ line.id }}</strong>   {{ line.text }}
+              {{ line.text }}
             </p>
           </a-drawer>
         </div>
@@ -871,6 +871,15 @@ onUnmounted(() => {
     }
   }
 }
+
+.custom-class .line-text {
+  padding: 10px;
+  border-radius: 4px;
+  
+  &:hover {
+    background-color: var(--main-light-4);
+  }
+}
 </style>
 
 <style lang="less">