diff --git a/src/core/database.py b/src/core/database.py index 5676d4fb..c1288996 100644 --- a/src/core/database.py +++ b/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: diff --git a/src/core/indexing.py b/src/core/indexing.py index a20e480d..5dd74f7e 100644 --- a/src/core/indexing.py +++ b/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, diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py index 4b16b58d..de64eb65 100644 --- a/src/core/knowledgebase.py +++ b/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 diff --git a/web/src/views/DataBaseInfoView.vue b/web/src/views/DataBaseInfoView.vue index 85e772d7..ad676216 100644 --- a/web/src/views/DataBaseInfoView.vue +++ b/web/src/views/DataBaseInfoView.vue @@ -94,8 +94,8 @@ @after-open-change="afterOpenChange" >

共 {{ selectedFile?.lines.length }} 个片段

-

- Chunk #{{ line.id }} {{ line.text }} +

+ {{ line.text }}

@@ -871,6 +871,15 @@ onUnmounted(() => { } } } + +.custom-class .line-text { + padding: 10px; + border-radius: 4px; + + &:hover { + background-color: var(--main-light-4); + } +}