优化文档添加逻辑
This commit is contained in:
Wenjie Zhang
parent
55b05df6dd
commit
7563c116fc
@ -29,7 +29,8 @@ def chunk(text_or_path, params=None):
|
||||
chunk_overlap=chunk_overlap,
|
||||
)
|
||||
|
||||
if os.path.isfile(text_or_path) and "uploads" in text_or_path:
|
||||
# 如果文件存在，并且是当前目录下的文件，则使用文件解析器
|
||||
if os.path.isfile(text_or_path) and os.path.exists(text_or_path) and os.path.abspath(text_or_path).startswith(os.getcwd()):
|
||||
parser = SimpleFileNodeParser()
|
||||
file_type = Path(text_or_path).suffix.lower()
|
||||
if file_type in [".txt", ".json", ".md"]:
|
||||
|
||||
@ -18,6 +18,12 @@ class KnowledgeBase:
|
||||
self.client = None
|
||||
self.work_dir = os.path.join(config.save_dir, "data")
|
||||
self.database_path = os.path.join(self.work_dir, "database.json")
|
||||
|
||||
# Configuration
|
||||
self.default_distance_threshold = 0.5
|
||||
self.default_rerank_threshold = 0.1
|
||||
self.default_max_query_count = 20
|
||||
|
||||
self._load_models()
|
||||
self._load_databases()
|
||||
|
||||
@ -29,6 +35,10 @@ class KnowledgeBase:
|
||||
from src.models.embedding import get_embedding_model
|
||||
self.embed_model = get_embedding_model(config)
|
||||
|
||||
if config.enable_reranker:
|
||||
from src.models.rerank_model import get_reranker
|
||||
self.reranker = get_reranker(config)
|
||||
|
||||
if not self.connect_to_milvus():
|
||||
raise ConnectionError("Failed to connect to Milvus")
|
||||
|
||||
@ -126,29 +136,87 @@ class KnowledgeBase:
|
||||
return {"lines": lines}
|
||||
|
||||
def get_kb_by_id(self, db_id):
|
||||
if not config.enable_knowledge_base:
|
||||
return None
|
||||
|
||||
return next((db for db in self.data if db.db_id == db_id), None)
|
||||
|
||||
def file_to_chunk(self, files, params=None):
|
||||
"""将文件转换为分块
|
||||
|
||||
这里主要是将文件转换为分块，但并不保存到数据库，仅仅返回分块后的信息，返回的信息里面也包含文件的id，文件名，文件类型，文件路径，文件状态，文件创建时间等。
|
||||
files: list of file path
|
||||
params: params for chunking
|
||||
|
||||
return: list of chunk info
|
||||
"""
|
||||
file_infos = {}
|
||||
for file in files:
|
||||
file_id = "file_" + hashstr(file + str(time.time()))
|
||||
|
||||
file_type = file.split(".")[-1].lower()
|
||||
|
||||
if file_type == "pdf":
|
||||
texts = read_text(file)
|
||||
nodes = chunk(texts, params=params)
|
||||
else:
|
||||
nodes = chunk(file, params=params)
|
||||
|
||||
file_infos[file_id] = {
|
||||
"file_id": file_id,
|
||||
"filename": os.path.basename(file),
|
||||
"path": file,
|
||||
"type": file_type,
|
||||
"status": "waiting",
|
||||
"created_at": time.time(),
|
||||
"nodes": [node.dict() for node in nodes]
|
||||
}
|
||||
|
||||
return file_infos
|
||||
|
||||
def url_to_chunk(self, url, params=None):
|
||||
"""将url转换为分块，读取url的内容，并转换为分块"""
|
||||
raise NotImplementedError("Not implemented")
|
||||
|
||||
def add_chunks(self, db_id, file_chunks):
|
||||
"""添加分块"""
|
||||
db = self.get_kb_by_id(db_id)
|
||||
|
||||
if db.embed_model != config.embed_model:
|
||||
logger.error(f"Embed model not match, {db.embed_model} != {config.embed_model}")
|
||||
return {"message": f"Embed model not match, cur: {config.embed_model}, req: {db.embed_model}", "status": "failed"}
|
||||
|
||||
db.files.update(file_chunks)
|
||||
self._save_databases()
|
||||
|
||||
for file_id, chunk in file_chunks.items():
|
||||
db.files[file_id]["status"] = "processing"
|
||||
self._save_databases()
|
||||
|
||||
try:
|
||||
self.add_documents(
|
||||
file_id=file_id,
|
||||
collection_name=db.db_id,
|
||||
docs=[node["text"] for node in chunk["nodes"]],
|
||||
chunk_infos=chunk["nodes"])
|
||||
|
||||
db.files[file_id]["status"] = "done"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to add documents to collection {db.db_id}, {e}, {traceback.format_exc()}")
|
||||
db.files[file_id]["status"] = "failed"
|
||||
|
||||
self._save_databases()
|
||||
|
||||
def add_files(self, db_id, files, params=None):
|
||||
db = self.get_kb_by_id(db_id)
|
||||
|
||||
if db.embed_model != config.embed_model:
|
||||
logger.error(f"Embed model not match, {db.embed_model} != {config.embed_model}")
|
||||
return {"message": f"Embed model not match, cur: {config.embed_model}", "status": "failed"}
|
||||
return {"message": f"Embed model not match, cur: {config.embed_model}, req: {db.embed_model}", "status": "failed"}
|
||||
|
||||
# Preprocessing the files to the queue
|
||||
new_files = {}
|
||||
for file in files:
|
||||
file_id = "file_" + hashstr(file + str(time.time()))
|
||||
new_file = {
|
||||
"file_id": file_id,
|
||||
"filename": os.path.basename(file),
|
||||
"path": file,
|
||||
"type": file.split(".")[-1].lower(),
|
||||
"status": "waiting",
|
||||
"created_at": time.time()
|
||||
}
|
||||
new_files[file_id] = new_file
|
||||
|
||||
new_files = self.file_to_chunk(files, params=params)
|
||||
db.files.update(new_files) # 更新数据库状态
|
||||
|
||||
# 先保存一次数据库状态，确保waiting状态被记录
|
||||
@ -160,17 +228,11 @@ class KnowledgeBase:
|
||||
self._save_databases()
|
||||
|
||||
try:
|
||||
if new_file["type"] == "pdf":
|
||||
texts = read_text(new_file["path"])
|
||||
nodes = chunk(texts, params=params)
|
||||
else:
|
||||
nodes = chunk(new_file["path"], params=params)
|
||||
|
||||
self.add_documents(
|
||||
file_id=file_id,
|
||||
collection_name=db.db_id,
|
||||
docs=[node.text for node in nodes],
|
||||
chunk_infos=[node.dict() for node in nodes])
|
||||
docs=[node["text"] for node in new_file["nodes"]],
|
||||
chunk_infos=new_file["nodes"])
|
||||
|
||||
db.files[file_id]["status"] = "done"
|
||||
|
||||
@ -204,8 +266,63 @@ class KnowledgeBase:
|
||||
self._load_models()
|
||||
self._load_databases()
|
||||
|
||||
###################################
|
||||
#* Below is the code for retriever #
|
||||
###################################
|
||||
|
||||
def get_retriever(self, db_id):
|
||||
db = self.get_kb_by_id(db_id)
|
||||
if db is None:
|
||||
raise Exception(f"database not found, {db_id}")
|
||||
|
||||
return db.retriever
|
||||
|
||||
def query(self, query, db_id, **kwargs):
|
||||
db = self.get_kb_by_id(db_id)
|
||||
|
||||
distance_threshold = kwargs.get("distance_threshold", self.default_distance_threshold)
|
||||
rerank_threshold = kwargs.get("rerank_threshold", self.default_rerank_threshold)
|
||||
max_query_count = kwargs.get("max_query_count", self.default_max_query_count)
|
||||
|
||||
all_db_result = self.search(query, db_id, limit=max_query_count)
|
||||
for res in all_db_result:
|
||||
res["file"] = db.files[res["entity"]["file_id"]]
|
||||
|
||||
db_result = [r for r in all_db_result if r["distance"] > distance_threshold]
|
||||
|
||||
if config.enable_reranker and len(db_result) > 0 and self.reranker:
|
||||
texts = [r["entity"]["text"] for r in db_result]
|
||||
rerank_scores = self.reranker.compute_score([query, texts], normalize=False)
|
||||
for i, r in enumerate(db_result):
|
||||
r["rerank_score"] = rerank_scores[i]
|
||||
db_result.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||||
db_result = [_res for _res in db_result if _res["rerank_score"] > rerank_threshold]
|
||||
|
||||
if kwargs.get("top_k", None):
|
||||
db_result = db_result[:kwargs["top_k"]]
|
||||
|
||||
return {
|
||||
"results": db_result,
|
||||
"all_results": all_db_result,
|
||||
}
|
||||
|
||||
def get_retriever(self, db_id):
|
||||
retriever_params = {
|
||||
"distance_threshold": self.default_distance_threshold,
|
||||
"rerank_threshold": self.default_rerank_threshold,
|
||||
"max_query_count": self.default_max_query_count,
|
||||
"top_k": 10,
|
||||
}
|
||||
|
||||
def retriever(query):
|
||||
response = self.query(query, db_id, **retriever_params)
|
||||
return response["results"]
|
||||
|
||||
return retriever
|
||||
|
||||
|
||||
################################
|
||||
# Below is the code for milvus #
|
||||
#* Below is the code for milvus #
|
||||
################################
|
||||
def connect_to_milvus(self):
|
||||
"""
|
||||
@ -276,7 +393,7 @@ class KnowledgeBase:
|
||||
return res
|
||||
|
||||
def search(self, query, collection_name, limit=3):
|
||||
|
||||
"""搜索数据库"""
|
||||
query_vectors = self.embed_model.batch_encode([query])
|
||||
return self.search_by_vector(query_vectors[0], collection_name, limit)
|
||||
|
||||
|
||||
@ -85,48 +85,35 @@ class Retriever:
|
||||
def query_knowledgebase(self, query, history, refs):
|
||||
"""查询知识库"""
|
||||
|
||||
kb_res = []
|
||||
final_res = []
|
||||
response = {
|
||||
"results": [],
|
||||
"all_results": [],
|
||||
"rw_query": query,
|
||||
"message": "",
|
||||
}
|
||||
|
||||
db_id = refs["meta"].get("db_id")
|
||||
meta = refs["meta"]
|
||||
|
||||
db_id = meta.get("db_id")
|
||||
if not db_id or not config.enable_knowledge_base:
|
||||
return {
|
||||
"results": final_res,
|
||||
"all_results": kb_res,
|
||||
"rw_query": query,
|
||||
"message": "Knowledge base is disabled",
|
||||
}
|
||||
response["message"] = "知识库未启用、或未指定知识库、或知识库不存在"
|
||||
return response
|
||||
|
||||
rw_query = self.rewrite_query(query, history, refs)
|
||||
|
||||
kb = knowledge_base.id2db[db_id]
|
||||
logger.debug(f"{refs['meta']=}")
|
||||
logger.debug(f"{meta=}")
|
||||
query_result = knowledge_base.query(query=rw_query,
|
||||
db_id=db_id,
|
||||
distance_threshold=meta.get("distanceThreshold", 0.5),
|
||||
rerank_threshold=meta.get("rerankThreshold", 0.1),
|
||||
max_query_count=meta.get("maxQueryCount", 20),
|
||||
top_k=meta.get("topK", 10))
|
||||
|
||||
meta = refs["meta"]
|
||||
max_query_count = meta.get("maxQueryCount", 10)
|
||||
rerank_threshold = meta.get("rerankThreshold", 0.1)
|
||||
distance_threshold = meta.get("distanceThreshold", 0)
|
||||
top_k = meta.get("topK", 5)
|
||||
response["results"] = query_result["results"]
|
||||
response["all_results"] = query_result["all_results"]
|
||||
response["rw_query"] = rw_query
|
||||
|
||||
# 检索
|
||||
all_kb_res = knowledge_base.search(rw_query, db_id, limit=max_query_count)
|
||||
for r in all_kb_res:
|
||||
r["file"] = kb.files[r["entity"]["file_id"]]
|
||||
|
||||
kb_res = [r for r in all_kb_res if r["distance"] > distance_threshold]
|
||||
|
||||
# 重排序
|
||||
if config.enable_reranker and len(kb_res) > 0:
|
||||
texts = [r["entity"]["text"] for r in kb_res]
|
||||
rerank_scores = self.reranker.compute_score([rw_query, texts], normalize=True)
|
||||
for i, r in enumerate(kb_res):
|
||||
r["rerank_score"] = rerank_scores[i]
|
||||
kb_res.sort(key=lambda x: x["rerank_score"], reverse=True)
|
||||
kb_res = [_res for _res in kb_res if _res["rerank_score"] > rerank_threshold]
|
||||
|
||||
kb_res = kb_res[:top_k]
|
||||
|
||||
return {"results": kb_res, "all_results": all_kb_res, "rw_query": rw_query}
|
||||
return response
|
||||
|
||||
def query_web(self, query, history, refs):
|
||||
"""查询网络"""
|
||||
@ -144,7 +131,9 @@ class Retriever:
|
||||
|
||||
def rewrite_query(self, query, history, refs):
|
||||
"""重写查询"""
|
||||
model = select_model(config)
|
||||
model_provider = config.model_provider_lite
|
||||
model_name = config.model_name_lite
|
||||
model = select_model(config, model_provider=model_provider, model_name=model_name)
|
||||
if refs["meta"].get("mode") == "search": # 如果是搜索模式，就使用 meta 的配置，否则就使用全局的配置
|
||||
rewrite_query_span = refs["meta"].get("use_rewrite_query", "off")
|
||||
else:
|
||||
@ -168,7 +157,9 @@ class Retriever:
|
||||
def reco_entities(self, query, history, refs):
|
||||
"""识别句子中的实体"""
|
||||
query = refs.get("rewritten_query", query)
|
||||
model = select_model(config)
|
||||
model_provider = config.model_provider_lite
|
||||
model_name = config.model_name_lite
|
||||
model = select_model(config, model_provider=model_provider, model_name=model_name)
|
||||
|
||||
entities = []
|
||||
if refs["meta"].get("use_graph"):
|
||||
|
||||
@ -49,6 +49,12 @@ async def query_test(query: str = Body(...), meta: dict = Body(...)):
|
||||
result = retriever.query_knowledgebase(query, history=None, refs={"meta": meta})
|
||||
return result
|
||||
|
||||
@data.post("/file-to-chunk")
|
||||
async def file_to_chunk(files: List[str] = Body(...), params: dict = Body(...)):
|
||||
logger.debug(f"File to chunk: {files}")
|
||||
result = knowledge_base.file_to_chunk(files, params=params)
|
||||
return result
|
||||
|
||||
@data.post("/add-by-file")
|
||||
async def create_document_by_file(db_id: str = Body(...), files: List[str] = Body(...)):
|
||||
logger.debug(f"Add document in {db_id} by file: {files}")
|
||||
@ -64,6 +70,20 @@ async def create_document_by_file(db_id: str = Body(...), files: List[str] = Bod
|
||||
logger.error(f"添加文件失败: {e}, {traceback.format_exc()}")
|
||||
return {"message": f"添加文件失败: {e}", "status": "failed"}
|
||||
|
||||
@data.post("/add-by-chunks")
|
||||
async def add_by_chunks(db_id: str = Body(...), file_chunks: dict = Body(...)):
|
||||
logger.debug(f"Add chunks in {db_id}: {file_chunks}")
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(
|
||||
executor, # 使用与chat_router相同的线程池
|
||||
lambda: knowledge_base.add_chunks(db_id, file_chunks)
|
||||
)
|
||||
return {"message": "分块添加完成", "status": "success"}
|
||||
except Exception as e:
|
||||
logger.error(f"添加分块失败: {e}, {traceback.format_exc()}")
|
||||
return {"message": f"添加分块失败: {e}", "status": "failed"}
|
||||
|
||||
@data.get("/info")
|
||||
async def get_database_info(db_id: str):
|
||||
# logger.debug(f"Get database {db_id} info")
|
||||
|
||||
@ -7,7 +7,7 @@
|
||||
<div class="database-info">
|
||||
<a-tag color="blue" v-if="database.embed_model">{{ database.embed_model }}</a-tag>
|
||||
<a-tag color="green" v-if="database.dimension">{{ database.dimension }}</a-tag>
|
||||
<span class="row-count">{{ database.metadata?.row_count }} 行 · {{ database.files?.length || 0 }} 文件</span>
|
||||
<span class="row-count">{{ database.metadata?.row_count }} 行 · {{ database.files ? Object.keys(database.files).length : 0 }} 文件</span>
|
||||
</div>
|
||||
</template>
|
||||
<template #actions>
|
||||
@ -22,37 +22,12 @@
|
||||
<a-alert v-if="database.embed_model != configStore.config.embed_model" message="向量模型不匹配，请重新选择" type="warning" style="margin: 10px 20px;" />
|
||||
<div class="db-main-container">
|
||||
<a-tabs v-model:activeKey="state.curPage" class="atab-container" type="card">
|
||||
<a-tab-pane key="add">
|
||||
<template #tab><span><CloudUploadOutlined />添加文件</span></template>
|
||||
|
||||
<a-tab-pane key="files">
|
||||
<template #tab><span><ReadOutlined />文件列表</span></template>
|
||||
<div class="db-tab-container">
|
||||
<div class="upload">
|
||||
<a-upload-dragger
|
||||
class="upload-dragger"
|
||||
v-model:fileList="fileList"
|
||||
name="file"
|
||||
:multiple="true"
|
||||
:disabled="state.loading"
|
||||
:action="'/api/data/upload?db_id=' + databaseId"
|
||||
@change="handleFileUpload"
|
||||
@drop="handleDrop"
|
||||
>
|
||||
<p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
|
||||
<p class="ant-upload-hint">
|
||||
目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加
|
||||
</p>
|
||||
</a-upload-dragger>
|
||||
</div>
|
||||
<div class="actions">
|
||||
<a-button
|
||||
type="primary"
|
||||
@click="addDocumentByFile"
|
||||
:loading="state.loading"
|
||||
:disabled="fileList.length === 0"
|
||||
style="margin: 0px 20px 20px 0;"
|
||||
>
|
||||
添加到知识库
|
||||
</a-button>
|
||||
<a-button @click="handleRefresh" :loading="state.refrashing">刷新状态</a-button>
|
||||
<a-button @click="handleRefresh" :loading="state.refrashing">刷新</a-button>
|
||||
</div>
|
||||
<a-table :columns="columns" :data-source="Object.values(database.files || {})" row-key="file_id" class="my-table">
|
||||
<template #bodyCell="{ column, text, record }">
|
||||
@ -60,7 +35,7 @@
|
||||
<a-button class="main-btn" type="link" @click="openFileDetail(record)">{{ text }}</a-button>
|
||||
</template>
|
||||
<template v-else-if="column.key === 'type'">
|
||||
<span :class="['span-type', text]">{{ text.toUpperCase() }}</span>
|
||||
<span :class="['span-type', text]">{{ text?.toUpperCase() }}</span>
|
||||
</template>
|
||||
<template v-else-if="column.key === 'status' && text === 'done'">
|
||||
<CheckCircleFilled style="color: #41A317;"/>
|
||||
@ -100,13 +75,102 @@
|
||||
</a-drawer>
|
||||
</div>
|
||||
</a-tab-pane>
|
||||
|
||||
<a-tab-pane key="add">
|
||||
<template #tab><span><CloudUploadOutlined />添加文件</span></template>
|
||||
<div class="db-tab-container">
|
||||
<div class="upload-section">
|
||||
<div class="upload-sidebar">
|
||||
<div class="chunking-params">
|
||||
<div class="params-info">
|
||||
<p>调整分块参数可以控制文本的切分方式，影响检索质量和文档加载效率。</p>
|
||||
</div>
|
||||
<a-form
|
||||
:model="chunkParams"
|
||||
name="basic"
|
||||
autocomplete="off"
|
||||
layout="vertical"
|
||||
>
|
||||
<a-form-item label="Chunk Size" name="chunk_size">
|
||||
<a-input-number v-model:value="chunkParams.chunk_size" :min="100" :max="10000" />
|
||||
<p class="param-description">每个文本片段的最大字符数</p>
|
||||
</a-form-item>
|
||||
<a-form-item label="Chunk Overlap" name="chunk_overlap">
|
||||
<a-input-number v-model:value="chunkParams.chunk_overlap" :min="0" :max="1000" />
|
||||
<p class="param-description">相邻文本片段间的重叠字符数</p>
|
||||
</a-form-item>
|
||||
<a-form-item label="使用文件节点解析器" name="use_parser">
|
||||
<a-switch v-model:checked="chunkParams.use_parser" />
|
||||
<p class="param-description">启用特定文件格式的智能分析</p>
|
||||
</a-form-item>
|
||||
</a-form>
|
||||
</div>
|
||||
</div>
|
||||
<div class="upload-main">
|
||||
<div class="upload">
|
||||
<a-upload-dragger
|
||||
class="upload-dragger"
|
||||
v-model:fileList="fileList"
|
||||
name="file"
|
||||
:multiple="true"
|
||||
:disabled="state.loading"
|
||||
:action="'/api/data/upload?db_id=' + databaseId"
|
||||
@change="handleFileUpload"
|
||||
@drop="handleDrop"
|
||||
>
|
||||
<p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
|
||||
<p class="ant-upload-hint">
|
||||
目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加
|
||||
</p>
|
||||
</a-upload-dragger>
|
||||
</div>
|
||||
<div class="actions">
|
||||
<a-button
|
||||
type="primary"
|
||||
@click="chunkFiles"
|
||||
:loading="state.loading"
|
||||
:disabled="fileList.length === 0"
|
||||
style="margin: 0px 20px 20px 0;"
|
||||
>
|
||||
生成分块
|
||||
</a-button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 分块结果预览区域 -->
|
||||
<div class="chunk-preview" v-if="chunkResults.length > 0">
|
||||
<div class="preview-header">
|
||||
<h3>分块预览 (共 {{ chunkResults.length }} 个文件，{{ getTotalChunks() }} 个分块)</h3>
|
||||
<a-button
|
||||
type="primary"
|
||||
@click="addToDatabase"
|
||||
:loading="state.adding"
|
||||
>
|
||||
添加到数据库
|
||||
</a-button>
|
||||
</div>
|
||||
|
||||
<a-collapse v-model:activeKey="activeFileKeys">
|
||||
<a-collapse-panel v-for="(file, fileIdx) in chunkResults" :key="fileIdx" :header="file.filename + ' (' + file.nodes.length + ' 个分块)'">
|
||||
<div id="result-cards" class="result-cards">
|
||||
<div v-for="(chunk, index) in file.nodes" :key="index" class="chunk">
|
||||
<p><strong>#{{ index + 1 }}</strong> {{ chunk.text }}</p>
|
||||
</div>
|
||||
</div>
|
||||
</a-collapse-panel>
|
||||
</a-collapse>
|
||||
</div>
|
||||
</div>
|
||||
</a-tab-pane>
|
||||
|
||||
<a-tab-pane key="query-test" force-render>
|
||||
<template #tab><span><SearchOutlined />检索测试</span></template>
|
||||
<div class="query-test-container db-tab-container">
|
||||
<div class="sider">
|
||||
<div class="sider-top">
|
||||
<div class="query-params" v-if="state.curPage == 'query-test'">
|
||||
<!-- <h3 class="params-title">参数配置</h3> -->
|
||||
<!-- <h3 class="params-title">查询参数</h3> -->
|
||||
<div class="params-group">
|
||||
<div class="params-item">
|
||||
<p>检索数量：</p>
|
||||
@ -120,6 +184,13 @@
|
||||
<p>筛选 TopK：</p>
|
||||
<a-input-number size="small" v-model:value="meta.topK" :min="1" :max="meta.maxQueryCount" />
|
||||
</div>
|
||||
<div class="params-item" v-if="configStore.config.enable_reranker">
|
||||
<p>排序方式：</p>
|
||||
<a-radio-group v-model:value="meta.sortBy" button-style="solid" size="small">
|
||||
<a-radio-button value="rerank_score">重排序分</a-radio-button>
|
||||
<a-radio-button value="distance">相似度</a-radio-button>
|
||||
</a-radio-group>
|
||||
</div>
|
||||
</div>
|
||||
<div class="params-group">
|
||||
<div class="params-item w100" v-if="configStore.config.enable_reranker">
|
||||
@ -162,12 +233,14 @@
|
||||
</div>
|
||||
|
||||
<!-- 新增示例按钮 -->
|
||||
<span>示例查询：</span>
|
||||
<div class="query-examples">
|
||||
<a-button v-for="example in queryExamples" :key="example" @click="useQueryExample(example)">
|
||||
{{ example }}
|
||||
</a-button>
|
||||
</div>
|
||||
<!-- <div class="query-examples-container">
|
||||
<div class="examples-title">示例查询：</div>
|
||||
<div class="query-examples">
|
||||
<a-button v-for="example in queryExamples" :key="example" @click="useQueryExample(example)">
|
||||
{{ example }}
|
||||
</a-button>
|
||||
</div>
|
||||
</div> -->
|
||||
<div class="query-test" v-if="queryResult">
|
||||
<div class="results-overview">
|
||||
<div class="results-stats">
|
||||
@ -180,8 +253,11 @@
|
||||
<span class="stat-item">
|
||||
<strong>TopK:</strong> {{ meta.topK }}
|
||||
</span>
|
||||
<span class="stat-item">
|
||||
<strong>排序:</strong> {{ meta.sortBy === 'rerank_score' ? '重排序分' : '相似度' }}
|
||||
</span>
|
||||
</div>
|
||||
<div class="rewritten-query">
|
||||
<div class="rewritten-query" v-if="queryResult.rw_query">
|
||||
<strong>重写后查询:</strong>
|
||||
<span class="query-text">{{ queryResult.rw_query }}</span>
|
||||
</div>
|
||||
@ -190,8 +266,8 @@
|
||||
<p>
|
||||
<strong>#{{ idx + 1 }} </strong>
|
||||
<span>{{ result.file.filename }} </span>
|
||||
<span><strong>距离</strong>：{{ result.distance.toFixed(4) }} </span>
|
||||
<span v-if="result.rerank_score"><strong>重排序</strong>：{{ result.rerank_score.toFixed(4) }}</span>
|
||||
<span><strong>相似度</strong>：{{ result.distance.toFixed(4) }} </span>
|
||||
<span v-if="result.rerank_score"><strong>重排序分</strong>：{{ result.rerank_score.toFixed(4) }}</span>
|
||||
</p>
|
||||
<p class="query-text">{{ result.entity.text }}</p>
|
||||
</div>
|
||||
@ -206,13 +282,13 @@
|
||||
</template>
|
||||
|
||||
<script setup>
|
||||
import { onMounted, reactive, ref, watch, toRaw, onUnmounted } from 'vue';
|
||||
import { onMounted, reactive, ref, watch, toRaw, onUnmounted, computed } from 'vue';
|
||||
import { message, Modal } from 'ant-design-vue';
|
||||
import { useRoute, useRouter } from 'vue-router';
|
||||
import { useConfigStore } from '@/stores/config'
|
||||
import HeaderComponent from '@/components/HeaderComponent.vue';
|
||||
import {
|
||||
ReadFilled,
|
||||
ReadOutlined,
|
||||
LeftOutlined,
|
||||
CheckCircleFilled,
|
||||
HourglassFilled,
|
||||
@ -221,7 +297,8 @@ import {
|
||||
DeleteOutlined,
|
||||
CloudUploadOutlined,
|
||||
SearchOutlined,
|
||||
LoadingOutlined
|
||||
LoadingOutlined,
|
||||
CaretUpOutlined
|
||||
} from '@ant-design/icons-vue'
|
||||
|
||||
|
||||
@ -241,12 +318,13 @@ const configStore = useConfigStore()
|
||||
|
||||
const state = reactive({
|
||||
loading: false,
|
||||
adding: false,
|
||||
refrashing: false,
|
||||
searchLoading: false,
|
||||
lock: false,
|
||||
drawer: false,
|
||||
refreshInterval: null,
|
||||
curPage: "add",
|
||||
curPage: "files",
|
||||
});
|
||||
|
||||
const meta = reactive({
|
||||
@ -257,6 +335,7 @@ const meta = reactive({
|
||||
rerankThreshold: 0.1,
|
||||
distanceThreshold: 0.3,
|
||||
topK: 10,
|
||||
sortBy: 'rerank_score',
|
||||
});
|
||||
|
||||
const use_rewrite_queryOptions = ref([
|
||||
@ -276,11 +355,24 @@ const filterQueryResults = () => {
|
||||
if (meta.filter) {
|
||||
results = results.filter(r => r.distance >= meta.distanceThreshold);
|
||||
console.log("before", results);
|
||||
|
||||
// 根据排序方式决定排序逻辑
|
||||
if (configStore.config.enable_reranker) {
|
||||
results = results
|
||||
.filter(r => r.rerank_score >= meta.rerankThreshold)
|
||||
.sort((a, b) => b.rerank_score - a.rerank_score);
|
||||
// 先过滤掉低于阈值的结果
|
||||
results = results.filter(r => r.rerank_score >= meta.rerankThreshold);
|
||||
|
||||
// 根据选择的排序方式进行排序
|
||||
if (meta.sortBy === 'rerank_score') {
|
||||
results = results.sort((a, b) => b.rerank_score - a.rerank_score);
|
||||
} else {
|
||||
// 按距离排序 (数值越大表示越相似)
|
||||
results = results.sort((a, b) => b.distance - a.distance);
|
||||
}
|
||||
} else {
|
||||
// 没有启用重排序时，默认按距离排序
|
||||
results = results.sort((a, b) => b.distance - a.distance);
|
||||
}
|
||||
|
||||
console.log("after", results);
|
||||
|
||||
results = results.slice(0, meta.topK);
|
||||
@ -489,42 +581,114 @@ const deleteFile = (fileId) => {
|
||||
})
|
||||
}
|
||||
|
||||
const addDocumentByFile = () => {
|
||||
|
||||
const chunkParams = ref({
|
||||
chunk_size: 1000,
|
||||
chunk_overlap: 200,
|
||||
use_parser: false,
|
||||
})
|
||||
|
||||
const chunkResults = ref([]);
|
||||
const activeFileKeys = ref([]);
|
||||
|
||||
// 获取所有分块的总数
|
||||
const getTotalChunks = () => {
|
||||
return chunkResults.value.reduce((total, file) => total + file.nodes.length, 0);
|
||||
}
|
||||
|
||||
// 分块预览
|
||||
const chunkFiles = () => {
|
||||
console.log(fileList.value)
|
||||
const files = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path)
|
||||
console.log(files)
|
||||
|
||||
state.loading = true
|
||||
state.lock = true
|
||||
if (files.length === 0) {
|
||||
message.error('请先上传文件')
|
||||
return
|
||||
}
|
||||
|
||||
fetch('/api/data/add-by-file', {
|
||||
state.loading = true
|
||||
|
||||
// 调用file-to-chunk接口获取分块信息
|
||||
fetch('/api/data/file-to-chunk', {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json" // 添加 Content-Type 头
|
||||
"Content-Type": "application/json"
|
||||
},
|
||||
body: JSON.stringify({
|
||||
files: files,
|
||||
params: chunkParams.value
|
||||
}),
|
||||
})
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
console.log('文件分块信息:', data)
|
||||
chunkResults.value = Object.values(data);
|
||||
activeFileKeys.value = chunkResults.value.length > 0 ? [0] : []; // 默认展开第一个文件
|
||||
})
|
||||
.catch(error => {
|
||||
console.error(error)
|
||||
message.error(error.message)
|
||||
})
|
||||
.finally(() => {
|
||||
state.loading = false
|
||||
})
|
||||
}
|
||||
|
||||
// 添加到数据库
|
||||
const addToDatabase = () => {
|
||||
if (chunkResults.value.length === 0) {
|
||||
message.error('没有可添加的分块')
|
||||
return
|
||||
}
|
||||
|
||||
state.adding = true
|
||||
state.lock = true
|
||||
|
||||
// 转换为API需要的格式
|
||||
const fileChunks = {};
|
||||
chunkResults.value.forEach(file => {
|
||||
fileChunks[file.file_id] = file;
|
||||
});
|
||||
|
||||
// 调用add-by-chunks接口将分块添加到数据库
|
||||
fetch('/api/data/add-by-chunks', {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json"
|
||||
},
|
||||
body: JSON.stringify({
|
||||
db_id: databaseId.value,
|
||||
files: files
|
||||
file_chunks: fileChunks
|
||||
}),
|
||||
})
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
console.log(data)
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
console.log(data)
|
||||
|
||||
if (data.status === 'failed') {
|
||||
message.error(data.message)
|
||||
} else {
|
||||
message.success(data.message)
|
||||
fileList.value = []
|
||||
if (data.status === 'failed') {
|
||||
message.error(data.message)
|
||||
} else {
|
||||
message.success(data.message)
|
||||
}
|
||||
})
|
||||
.catch(error => {
|
||||
console.error(error)
|
||||
message.error(error.message)
|
||||
})
|
||||
.finally(() => {
|
||||
getDatabaseInfo()
|
||||
state.loading = false
|
||||
})
|
||||
chunkResults.value = []
|
||||
activeFileKeys.value = []
|
||||
}
|
||||
})
|
||||
.catch(error => {
|
||||
console.error(error)
|
||||
message.error(error.message)
|
||||
})
|
||||
.finally(() => {
|
||||
getDatabaseInfo()
|
||||
state.adding = false
|
||||
state.lock = false
|
||||
})
|
||||
}
|
||||
|
||||
const addDocumentByFile = () => {
|
||||
// 此函数不再需要，由chunkFiles和addToDatabase替代
|
||||
console.log('此功能已被拆分为两个步骤')
|
||||
}
|
||||
|
||||
const columns = [
|
||||
@ -549,9 +713,12 @@ watch(() => meta, () => {
|
||||
filterQueryResults()
|
||||
}, { deep: true })
|
||||
|
||||
// 添加示例查询
|
||||
// 添加更多示例查询
|
||||
const queryExamples = ref([
|
||||
'贾宝玉的丫鬟有哪些？',
|
||||
'请介绍一下红楼梦的主要人物',
|
||||
'林黛玉是什么性格？',
|
||||
'曹雪芹的创作背景',
|
||||
]);
|
||||
|
||||
// 使用示例查询的方法
|
||||
@ -616,14 +783,13 @@ onUnmounted(() => {
|
||||
box-sizing: border-box;
|
||||
font-size: 15px;
|
||||
gap: 12px;
|
||||
// background-color: var(--main-light-6);
|
||||
// padding: 16px;
|
||||
// box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
|
||||
// border: 1px solid var(--main-light-3);
|
||||
// border-radius: 8px;
|
||||
padding-top: 12px;
|
||||
padding-right: 16px;
|
||||
border-right: 1px solid var(--main-light-3);
|
||||
border: 1px solid var(--main-light-3);
|
||||
background-color: var(--main-light-6);
|
||||
border-radius: 8px;
|
||||
padding: 16px;
|
||||
margin-right: 8px;
|
||||
|
||||
.params-title {
|
||||
margin-top: 0;
|
||||
@ -631,6 +797,7 @@ onUnmounted(() => {
|
||||
color: var(--main-color);
|
||||
font-size: 18px;
|
||||
text-align: center;
|
||||
font-weight: bold;
|
||||
}
|
||||
|
||||
.params-group {
|
||||
@ -658,7 +825,6 @@ onUnmounted(() => {
|
||||
|
||||
p {
|
||||
margin: 0;
|
||||
font-size: 16px;
|
||||
color: var(--gray-900);
|
||||
}
|
||||
|
||||
@ -716,22 +882,36 @@ onUnmounted(() => {
|
||||
}
|
||||
}
|
||||
|
||||
.query-examples {
|
||||
display: flex;
|
||||
flex-wrap: wrap;
|
||||
gap: 8px;
|
||||
margin: 10px 0;
|
||||
.query-examples-container {
|
||||
margin-bottom: 20px;
|
||||
padding: 12px;
|
||||
background: var(--main-light-6);
|
||||
border-radius: 8px;
|
||||
border: 1px solid var(--main-light-3);
|
||||
|
||||
.ant-btn {
|
||||
font-size: 14px;
|
||||
padding: 4px 12px;
|
||||
height: auto;
|
||||
background-color: var(--gray-200);
|
||||
border: none;
|
||||
color: var(--gray-800);
|
||||
.examples-title {
|
||||
font-weight: bold;
|
||||
margin-bottom: 10px;
|
||||
color: var(--main-color);
|
||||
}
|
||||
|
||||
&:hover {
|
||||
background-color: var(--gray-300);
|
||||
.query-examples {
|
||||
display: flex;
|
||||
flex-wrap: wrap;
|
||||
gap: 8px;
|
||||
margin: 10px 0 0;
|
||||
|
||||
.ant-btn {
|
||||
font-size: 14px;
|
||||
padding: 4px 12px;
|
||||
height: auto;
|
||||
background-color: var(--gray-200);
|
||||
border: none;
|
||||
color: var(--gray-800);
|
||||
|
||||
&:hover {
|
||||
color: var(--main-color);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -816,6 +996,7 @@ onUnmounted(() => {
|
||||
}
|
||||
|
||||
.upload {
|
||||
margin-bottom: 20px;
|
||||
.upload-dragger {
|
||||
margin: 0px;
|
||||
}
|
||||
@ -885,6 +1066,140 @@ onUnmounted(() => {
|
||||
background-color: var(--main-light-4);
|
||||
}
|
||||
}
|
||||
|
||||
.upload-section {
|
||||
display: flex;
|
||||
gap: 20px;
|
||||
|
||||
.upload-sidebar {
|
||||
width: 280px;
|
||||
padding: 20px;
|
||||
background-color: var(--main-light-6);
|
||||
border-radius: 8px;
|
||||
border: 1px solid var(--main-light-3);
|
||||
// box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
|
||||
|
||||
.chunking-params {
|
||||
h4 {
|
||||
margin-top: 0;
|
||||
margin-bottom: 16px;
|
||||
color: var(--main-color);
|
||||
font-size: 18px;
|
||||
text-align: center;
|
||||
font-weight: bold;
|
||||
padding-bottom: 10px;
|
||||
border-bottom: 1px dashed var(--main-light-3);
|
||||
}
|
||||
|
||||
.params-info {
|
||||
background-color: var(--main-light-4);
|
||||
border-radius: 6px;
|
||||
padding: 10px 12px;
|
||||
margin-bottom: 16px;
|
||||
|
||||
p {
|
||||
margin: 0;
|
||||
font-size: 13px;
|
||||
line-height: 1.5;
|
||||
color: var(--gray-700);
|
||||
}
|
||||
}
|
||||
|
||||
.ant-form-item {
|
||||
margin-bottom: 16px;
|
||||
|
||||
.ant-form-item-label {
|
||||
padding-bottom: 6px;
|
||||
|
||||
label {
|
||||
color: var(--gray-800);
|
||||
font-weight: 500;
|
||||
font-size: 15px;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
.ant-input-number {
|
||||
width: 100%;
|
||||
border-radius: 6px;
|
||||
|
||||
&:hover, &:focus {
|
||||
border-color: var(--main-color);
|
||||
}
|
||||
}
|
||||
|
||||
.ant-switch {
|
||||
background-color: var(--gray-400);
|
||||
|
||||
&.ant-switch-checked {
|
||||
background-color: var(--main-color);
|
||||
}
|
||||
}
|
||||
|
||||
// 添加参数说明
|
||||
.param-description {
|
||||
color: var(--gray-600);
|
||||
font-size: 12px;
|
||||
margin-top: 4px;
|
||||
margin-bottom: 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
.upload-main {
|
||||
flex: 1;
|
||||
}
|
||||
}
|
||||
|
||||
.chunk-preview {
|
||||
margin-top: 20px;
|
||||
|
||||
.preview-header {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
margin-bottom: 16px;
|
||||
|
||||
h3 {
|
||||
margin: 0;
|
||||
color: var(--main-color);
|
||||
font-size: 18px;
|
||||
}
|
||||
}
|
||||
|
||||
.result-cards {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(auto-fill, minmax(600px, 1fr));
|
||||
gap: 12px;
|
||||
margin-top: 10px;
|
||||
}
|
||||
|
||||
.chunk {
|
||||
background-color: var(--main-light-5);
|
||||
border: 1px solid var(--main-light-3);
|
||||
border-radius: 8px;
|
||||
padding: 16px;
|
||||
word-wrap: break-word;
|
||||
word-break: break-all;
|
||||
transition: all 0.2s ease;
|
||||
|
||||
&:hover {
|
||||
background-color: var(--main-light-4);
|
||||
border-color: var(--main-light-2);
|
||||
box-shadow: 0 2px 6px rgba(0, 0, 0, 0.06);
|
||||
}
|
||||
|
||||
p {
|
||||
margin: 0;
|
||||
line-height: 1.6;
|
||||
|
||||
strong {
|
||||
color: var(--main-color);
|
||||
margin-right: 6px;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
<style lang="less">
|
||||
@ -954,3 +1269,4 @@ onUnmounted(() => {
|
||||
|
||||
|
||||
</style>
|
||||
|
||||
|
||||
@ -45,7 +45,7 @@
|
||||
<div class="icon"><ReadFilled /></div>
|
||||
<div class="info">
|
||||
<h3>{{ database.name }}</h3>
|
||||
<p><span>{{ database.metadata.row_count }} 行</span></p>
|
||||
<p><span>{{ database.metadata.row_count }} 行</span> · <span>{{ database.files ? Object.keys(database.files).length : 0 }} 文件</span></p>
|
||||
</div>
|
||||
</div>
|
||||
<p class="description">{{ database.description || '暂无描述' }}</p>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user