kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

进一步优化并行处理能力：file-to-chunk 部分优化完成

Browse Source
This commit is contained in:
Wenjie Zhang 2025-05-09 23:47:16 +08:00
parent d474f7768f
commit 144632b703
3 changed files with 19 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
server/routers/data_router.py
View File

@ -55,25 +55,20 @@ async def query_test(query: str = Body(...), meta: dict = Body(...), current_use
@data.post("/file-to-chunk")
async def file_to_chunk(files: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
logger.debug(f"File to chunk: {files}")
result = knowledge_base.file_to_chunk(files, params=params)
result = await knowledge_base.file_to_chunk(files, params=params)
return result
@data.post("/url-to-chunk")
async def url_to_chunk(urls: List[str] = Body(...), params: dict = Body(...), current_user: User = Depends(get_admin_user)):
logger.debug(f"Url to chunk: {urls}")
result = knowledge_base.url_to_chunk(urls, params=params)
result = await knowledge_base.url_to_chunk(urls, params=params)
return result
@data.post("/add-by-file")
async def create_document_by_file(db_id: str = Body(...), files: List[str] = Body(...), current_user: User = Depends(get_admin_user)):
logger.debug(f"Add document in {db_id} by file: {files}")
try:
# 使用线程池执行耗时操作
loop = asyncio.get_event_loop()
await loop.run_in_executor(
executor, # 使用与chat_router相同的线程池
lambda: knowledge_base.add_files(db_id, files)
)
await knowledge_base.add_files(db_id, files)
return {"message": "文件添加完成", "status": "success"}
except Exception as e:
logger.error(f"添加文件失败: {e}, {traceback.format_exc()}")

3
src/core/indexing.py
View File

@ -1,4 +1,5 @@
import os
import asyncio
from pathlib import Path
from llama_index.core import Document
from llama_index.core.node_parser import SimpleFileNodeParser
@ -94,3 +95,5 @@ def read_text(file, params=None):
raise Exception(f"File format not supported, only support {support_format}")
async def read_text_async(file):
return await asyncio.to_thread(read_text, file)

26
src/core/knowledgebase.py
View File

@ -8,7 +8,7 @@ from pymilvus import MilvusClient, MilvusException
from src import config
from src.utils import logger, hashstr
from src.core.indexing import chunk, read_text
from src.core.indexing import chunk, read_text_async
from src.core.kb_db_manager import kb_db_manager
class KnowledgeBase:
@ -175,7 +175,7 @@ class KnowledgeBase:
return self.db_manager.get_database_by_id(db_id)
def file_to_chunk(self, files, params=None):
async def file_to_chunk(self, files, params=None):
"""将文件转换为分块
这里主要是将文件转换为分块，但并不保存到数据库，仅仅返回分块后的信息，返回的信息里面也包含文件的id，文件名，文件类型，文件路径，文件状态，文件创建时间等。
@ -191,7 +191,7 @@ class KnowledgeBase:
file_type = file.split(".")[-1].lower()
if file_type == "pdf":
texts = read_text(file)
texts = await read_text_async(file)
nodes = chunk(texts, params=params)
else:
nodes = chunk(file, params=params)
@ -208,7 +208,7 @@ class KnowledgeBase:
return file_infos
def url_to_chunk(self, urls, params=None):
async def url_to_chunk(self, urls, params=None):
"""将url转换为分块，读取url的内容，并转换为分块
Args:
@ -226,7 +226,7 @@ class KnowledgeBase:
file_infos = {}
# 使用UnstructuredURLLoader加载URL内容
loader = UnstructuredURLLoader(urls=urls, continue_on_failure=True)
# loader = UnstructuredURLLoader(urls=urls, continue_on_failure=True)
for url_idx, url in enumerate(urls):
file_id = "url_" + hashstr(url + str(time.time()))
@ -234,7 +234,7 @@ class KnowledgeBase:
try:
# 加载单个URL内容
single_loader = UnstructuredURLLoader(urls=[url], continue_on_failure=False)
documents = single_loader.load()
documents = await single_loader.aload()
# 将文档内容合并
text_content = "\n\n".join([doc.page_content for doc in documents])
@ -278,7 +278,7 @@ class KnowledgeBase:
return file_infos
def add_chunks(self, db_id, file_chunks):
async def add_chunks(self, db_id, file_chunks):
"""添加分块"""
db = self.get_kb_by_id(db_id)
@ -298,7 +298,7 @@ class KnowledgeBase:
)
try:
self.add_documents(
await self.add_documents(
file_id=file_id,
collection_name=db_id,
docs=[node["text"] for node in chunk_info["nodes"]],
@ -312,7 +312,7 @@ class KnowledgeBase:
# 更新文件状态为失败
self.db_manager.update_file_status(file_id, "failed")
def add_files(self, db_id, files, params=None):
async def add_files(self, db_id, files, params=None):
db = self.get_kb_by_id(db_id)
if not self.check_embed_model(db_id):
@ -320,7 +320,7 @@ class KnowledgeBase:
return {"message": f"Embed model not match, cur: {self.embed_model.embed_model_fullname}, req: {db['embed_model']}", "status": "failed"}
# Preprocessing the files to the queue
new_files = self.file_to_chunk(files, params=params)
new_files = await self.file_to_chunk(files, params=params)
for file_id, new_file in new_files.items():
# 在数据库中创建文件记录
@ -334,7 +334,7 @@ class KnowledgeBase:
)
try:
self.add_documents(
await self.add_documents(
file_id=file_id,
collection_name=db_id,
docs=[node["text"] for node in new_file["nodes"]],
@ -498,7 +498,7 @@ class KnowledgeBase:
dimension= dimension, # The vectors we will use in this demo has 768 dimensions
)
def add_documents(self, docs, collection_name, chunk_infos=None, **kwargs):
async def add_documents(self, docs, collection_name, chunk_infos=None, **kwargs):
"""添加已经分块之后的文本"""
# 检查 collection 是否存在
import random
@ -508,7 +508,7 @@ class KnowledgeBase:
chunk_infos = chunk_infos or [{}] * len(docs)
vectors = self.embed_model.batch_encode(docs)
vectors = await self.embed_model.abatch_encode(docs)
data = [{
"id": int(random.random() * 1e12),