feat: 优化批量上传脚本,支持多种文件格式上传并添加任务状态检查功能

This commit is contained in:
Wenjie Zhang 2025-10-26 21:10:42 +08:00
parent 2a6be3e2ac
commit 107ec04dba
5 changed files with 261 additions and 144 deletions

View File

@ -116,38 +116,36 @@ docker compose up -d api
### 文件上传脚本
使用 `scripts/batch_upload.py upload` 批量上传文件到知识库:
使用 `scripts/batch_upload.py` 批量上传文件到知识库:
```bash
# 批量上传文档
uv run scripts/batch_upload.py upload \
--db-id your_kb_id \
--directory path/to/your/data \
--pattern "*.docx" \
--base-url http://127.0.0.1:5050/api \
--username your_username \
--password your_password \
--concurrency 4 \
# 批量上传文档(多种格式)
uv run scripts/batch_upload.py \
--db-id kb_b2730ad6801b149694021106c7eddd38 \
--directory data.nogit/农业农村局 \
--pattern "*.docx" --pattern "*.txt" --pattern "*.html" \
--base-url http://172.19.13.6:5050/api \
--username admin \
--password admin123 \
--batch-size 20 \
--wait-for-completion \
--poll-interval 5 \
--recursive \
--record-file scripts/tmp/batch_processed_files.txt
--record-file scripts/tmp/batch_processed_files_1029.txt
```
**参数说明**:
- `--db-id`: 目标知识库 ID
- `--directory`: 文件目录路径
- `--pattern`: 文件匹配模式
- `--concurrency`: 并发处理数量
- `--pattern`: 文件匹配模式,可以多次指定以支持多种格式(例如:`--pattern "*.docx" --pattern "*.pdf" --pattern "*.html"`
- `--batch-size`: 每批处理的文件数量默认20
- `--wait-for-completion`: 是否等待任务完成再处理下一批(默认开启)
- `--poll-interval`: 任务状态检查间隔单位秒默认5秒
- `--recursive`: 递归处理子目录
- `--record-file`: 处理记录文件路径
提示系统按“内容哈希”进行去重同一知识库已存在相同内容的文件会被拒绝409
### 脚本功能
- **进度跟踪**: 实时显示处理进度
- **错误处理**: 自动跳过无法处理的文件
- **断点续传**: 支持中断后继续处理
- **日志记录**: 详细记录处理过程
- **结果统计**: 处理完成后显示统计信息
更多关于“入库参数、导出数据、支持类型”等,请参阅:介绍 → 知识库与知识图谱 → 文档管理。
**注意事项**:
- 系统按"内容哈希"进行去重同一知识库已存在相同内容的文件会被拒绝409
- 建议根据系统性能调整批次大小
- 大量文件处理时建议开启分批等待功能
- 先上传后处理的机制更稳定,适合大批量文档导入

View File

@ -7,13 +7,12 @@
## Bugs
- [x] 修复本地知识库的 metadata 和 向量数据库中不一致的情况。
- [ ] v1 版本的 LangGraph 的工具渲染有问题
- [ ] upload 接口会阻塞主进程
## Next
- [x] 修改现有的智能体Demo并尽量将默认助手的特性兼容到 LangGraph 的 [`create_agent`](https://docs.langchain.com/oss/python/langchain/agents) 中
- [x] 基于 create_agent 创建 SQL Viewer 智能体 <Badge type="info" text="0.3.5" />
- [x] 优化 MCP 逻辑,支持 common + special 创建方式 <Badge type="info" text="0.3.5" />
- [ ] 新建 DeepAgents 智能体
- [ ] 添加对于上传文件的支持
- [ ] 统一图谱数据结构,优化可视化方式 [#298](https://github.com/xerrors/Yuxi-Know/issues/298) <Badge type="info" text="0.4" />
- [ ] 集成智能体评估,首先使用命令行来实现,然后考虑放在 UI 里面展示
@ -37,3 +36,7 @@
- [x] 支持 MinerU 2.5 的解析方法 <Badge type="info" text="0.3.5" />
- [x] 文件管理1文件选择的时候会跨数据库2文件校验会算上失败的文件
- [x] Tasker 中获取历史任务的时候,仅获取 top100 个 task。
- [x] 修改现有的智能体Demo并尽量将默认助手的特性兼容到 LangGraph 的 [`create_agent`](https://docs.langchain.com/oss/python/langchain/agents) 中
- [x] 基于 create_agent 创建 SQL Viewer 智能体 <Badge type="info" text="0.3.5" />
- [x] 优化 MCP 逻辑,支持 common + special 创建方式 <Badge type="info" text="0.3.5" />
- [x] 修复本地知识库的 metadata 和 向量数据库中不一致的情况。

View File

@ -29,6 +29,21 @@ async def login(client: httpx.AsyncClient, base_url: str, username: str, passwor
return None
async def check_task_status(client: httpx.AsyncClient, base_url: str, task_id: str) -> str | None:
"""Check the status of a task. Returns status string or None if failed."""
try:
response = await client.get(f"{base_url}/tasks/{task_id}")
response.raise_for_status()
task_data = response.json().get("task", {})
return task_data.get("status")
except httpx.HTTPStatusError as e:
console.print(f"[bold yellow]Warning: Failed to check task {task_id}: {e.response.status_code}[/bold yellow]")
return None
except httpx.RequestError as e:
console.print(f"[bold yellow]Warning: Failed to check task {task_id}: {e}[/bold yellow]")
return None
async def upload_file(
client: httpx.AsyncClient,
base_url: str,
@ -67,7 +82,7 @@ async def process_document(
chunk_overlap: int = 200,
use_qa_split: bool = False,
qa_separator: str = "\n\n\n",
) -> bool:
) -> tuple[bool, str | None]:
"""Triggers the processing of an uploaded file in the knowledge base."""
# Prepare processing parameters
params = {
@ -97,26 +112,26 @@ async def process_document(
f"[bold cyan]Ingestion queued for {server_file_path}{extra}. "
"Track progress in the task center.[/bold cyan]"
)
return True
return True, task_id
# Check if the overall request was successful for synchronous responses
if overall_status != "success":
console.print(
f"[bold yellow]Processing warning for {server_file_path}: {result.get('message')}[/bold yellow]"
)
return False
return False, None
# Check the specific file's processing status in the items array
items = result.get("items", [])
if not items:
console.print(f"[bold red]No processing result for {server_file_path}[/bold red]")
return False
return False, None
# Since we only sent one file, check the first item
item = items[0]
# Check for both 'success' and 'done' status (different APIs might use different status values)
if item.get("status") in ["success", "done"]:
return True
return True, None
else:
# Get more detailed error information
error_msg = item.get("message", "")
@ -140,59 +155,122 @@ async def process_document(
# Also log the full item for debugging
console.print(f"[dim]Debug - Full item response: {item}[/dim]")
return False
return False, None
except httpx.HTTPStatusError as e:
console.print(
f"[bold red]Failed to process {server_file_path}: {e.response.status_code} - {e.response.text}[/bold red]"
)
return False
return False, None
except httpx.RequestError as e:
console.print(f"[bold red]Failed to process {server_file_path}: {e}[/bold red]")
return False
return False, None
async def worker(
semaphore: asyncio.Semaphore,
async def upload_single_file(
client: httpx.AsyncClient,
base_url: str,
db_id: str,
file_path: pathlib.Path,
file_hash: str,
progress: Progress,
upload_task_id: int,
process_task_id: int,
task_id: int,
) -> str | None:
"""Upload a single file and return server file path."""
server_file_path = await upload_file(client, base_url, db_id, file_path)
if server_file_path:
progress.update(task_id, advance=1, postfix=f"Uploaded {file_path.name}")
else:
progress.update(task_id, advance=1, postfix=f"Failed: {file_path.name}")
return server_file_path
async def add_batch_to_knowledge_base(
client: httpx.AsyncClient,
base_url: str,
db_id: str,
server_file_paths: list[str],
enable_ocr: str = "paddlex_ocr",
chunk_size: int = 1000,
chunk_overlap: int = 200,
use_qa_split: bool = False,
qa_separator: str = "\n\n\n",
):
"""A worker task that uploads and then processes a single file."""
async with semaphore:
# 1. Upload file
server_file_path = await upload_file(client, base_url, db_id, file_path)
progress.update(upload_task_id, advance=1, postfix=f"Uploaded {file_path.name}")
) -> tuple[bool, str | None]:
"""Add a batch of files to knowledge base and return task_id."""
if not server_file_paths:
return True, None
if not server_file_path:
progress.update(process_task_id, advance=1) # Mark as processed to not hang the progress bar
return file_path, file_hash, "upload_failed"
# Prepare processing parameters
params = {
"chunk_size": chunk_size,
"chunk_overlap": chunk_overlap,
"enable_ocr": enable_ocr,
"use_qa_split": use_qa_split,
"qa_separator": qa_separator,
"content_type": "file",
}
# 2. Process file
success = await process_document(
client,
base_url,
db_id,
server_file_path,
enable_ocr=enable_ocr,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
use_qa_split=use_qa_split,
qa_separator=qa_separator,
try:
response = await client.post(
f"{base_url}/knowledge/databases/{db_id}/documents",
json={"items": server_file_paths, "params": params},
timeout=600, # 10 minutes timeout for processing
)
progress.update(process_task_id, advance=1, postfix=f"Processed {file_path.name}")
response.raise_for_status()
result = response.json()
return file_path, file_hash, "success" if success else "processing_failed"
overall_status = result.get("status")
if overall_status == "queued":
task_id = result.get("task_id")
extra = f" (task id: {task_id})" if task_id else ""
console.print(
f"[bold cyan]Batch of {len(server_file_paths)} files queued for processing{extra}. "
"Track progress in the task center.[/bold cyan]"
)
return True, task_id
elif overall_status == "success":
console.print(f"[bold green]Batch of {len(server_file_paths)} files processed successfully[/bold green]")
return True, None
else:
console.print(f"[bold yellow]Batch processing warning: {result.get('message')}[/bold yellow]")
return False, None
except httpx.HTTPStatusError as e:
console.print(f"[bold red]Failed to process batch: {e.response.status_code} - {e.response.text}[/bold red]")
return False, None
except httpx.RequestError as e:
console.print(f"[bold red]Failed to process batch: {e}[/bold red]")
return False, None
async def wait_for_tasks_completion(
client: httpx.AsyncClient,
base_url: str,
task_ids: list[str],
poll_interval: int = 5,
) -> dict[str, str]:
"""Wait for all tasks to complete and return their final statuses."""
if not task_ids:
return {}
console.print(f"[bold cyan]Waiting for {len(task_ids)} tasks to complete...[/bold cyan]")
pending_tasks = task_ids.copy()
completed_tasks = {}
while pending_tasks:
for task_id in pending_tasks.copy():
status = await check_task_status(client, base_url, task_id)
if status:
if status in ["success", "failed", "cancelled"]:
completed_tasks[task_id] = status
pending_tasks.remove(task_id)
console.print(f"[dim]Task {task_id} completed with status: {status}[/dim]")
if pending_tasks:
await asyncio.sleep(poll_interval)
console.print(f"[bold green]All {len(task_ids)} tasks completed[/bold green]")
return completed_tasks
def get_file_hash(file_path: pathlib.Path) -> str:
@ -236,11 +314,10 @@ def upload(
directory: pathlib.Path = typer.Option(
..., help="The directory containing files to upload.", exists=True, file_okay=False
),
pattern: str = typer.Option("*.md", help="The glob pattern for files to upload (e.g., '*.pdf', '**/*.txt')."),
pattern: list[str] = typer.Option(["*.md"], help="The glob patterns for files to upload (e.g., '*.pdf', '**/*.txt'). Can be specified multiple times."),
base_url: str = typer.Option("http://127.0.0.1:5050/api", help="The base URL of the API server."),
username: str = typer.Option(..., help="Admin username for login."),
password: str = typer.Option(..., help="Admin password for login."),
concurrency: int = typer.Option(1, help="The number of concurrent upload/process tasks."),
recursive: bool = typer.Option(False, "--recursive", "-r", help="Search for files recursively in subdirectories."),
record_file: pathlib.Path = typer.Option(
"scripts/tmp/batch_processed_files.txt", help="File to store processed files record."
@ -252,6 +329,9 @@ def upload(
),
use_qa_split: bool = typer.Option(False, help="Whether to use QA splitting."),
qa_separator: str = typer.Option("\n\n\n", help="Separator for QA splitting."),
batch_size: int = typer.Option(20, help="Number of files to process in each batch."),
wait_for_completion: bool = typer.Option(True, help="Whether to wait for tasks to complete before next batch."),
poll_interval: int = typer.Option(5, help="Polling interval in seconds for checking task status."),
):
"""
Batch upload and process files into a Yuxi-Know knowledge base.
@ -262,11 +342,19 @@ def upload(
processed_files = load_processed_files(record_file)
console.print(f"Loaded {len(processed_files)} previously processed files from record.")
# Discover files
# Discover files from multiple patterns
glob_method = directory.rglob if recursive else directory.glob
all_files = list(glob_method(pattern))
all_files = []
for pat in pattern:
files_for_pat = list(glob_method(pat))
all_files.extend(files_for_pat)
# Remove duplicates
all_files = list(set(all_files))
if not all_files:
console.print(f"[bold yellow]No files found in '{directory}' matching '{pattern}'. Aborting.[/bold yellow]")
patterns_str = "', '".join(pattern)
console.print(f"[bold yellow]No files found in '{directory}' matching patterns: '{patterns_str}'. Aborting.[/bold yellow]")
raise typer.Exit()
# 过滤掉macos的隐藏文件
@ -302,80 +390,106 @@ def upload(
client.headers = {"Authorization": f"Bearer {token}"}
# Setup concurrency and tasks
semaphore = asyncio.Semaphore(concurrency)
tasks = []
# Process files in batches: upload 20 -> process 20 -> wait -> repeat
total_processed_files = []
total_upload_failures = []
total_processing_failures = []
all_successful_hashes = set()
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeElapsedColumn(),
TextColumn("{task.fields[postfix]}"),
console=console,
transient=True,
) as progress:
upload_task_id = progress.add_task("[bold blue]Uploading...", total=len(files_to_upload), postfix="")
process_task_id = progress.add_task("[bold cyan]Processing...", total=len(files_to_upload), postfix="")
# Split all files into batches
for batch_num in range(0, len(files_to_upload), batch_size):
batch_files = files_to_upload[batch_num:batch_num + batch_size]
batch_start = batch_num + 1
batch_end = min(batch_num + batch_size, len(files_to_upload))
for file_path, file_hash in files_to_upload:
task = asyncio.create_task(
worker(
semaphore,
client,
base_url,
db_id,
file_path,
file_hash,
progress,
upload_task_id,
process_task_id,
enable_ocr=enable_ocr,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
use_qa_split=use_qa_split,
qa_separator=qa_separator,
console.print(f"\n[bold yellow]=== Batch {batch_start}-{batch_end} of {len(files_to_upload)} ===[/bold yellow]")
# Step 1: Upload this batch of files sequentially
console.print(f"[blue]Step 1: Uploading {len(batch_files)} files...[/blue]")
successful_uploads = []
batch_upload_failures = []
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeElapsedColumn(),
TextColumn("{task.fields[postfix]}"),
console=console,
transient=True,
) as progress:
upload_task_id = progress.add_task(f"Uploading batch {batch_start}-{batch_end}...", total=len(batch_files), postfix="")
for file_path, file_hash in batch_files:
server_file_path = await upload_single_file(
client, base_url, db_id, file_path, progress, upload_task_id
)
if server_file_path:
successful_uploads.append((file_path, file_hash, server_file_path))
all_successful_hashes.add(file_hash)
else:
batch_upload_failures.append(file_path)
# Step 2: Process this batch if uploads succeeded
if successful_uploads:
console.print(f"[green]Step 2: Processing {len(successful_uploads)} uploaded files...[/green]")
# Extract server file paths
server_file_paths = [item[2] for item in successful_uploads]
# Submit batch to knowledge base
success, task_id = await add_batch_to_knowledge_base(
client,
base_url,
db_id,
server_file_paths,
enable_ocr=enable_ocr,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
use_qa_split=use_qa_split,
qa_separator=qa_separator,
)
tasks.append(task)
results = await asyncio.gather(*tasks)
if success:
total_processed_files.extend([item[0] for item in successful_uploads])
# Summarize results and update processed files record
successful_files = []
upload_failures = []
processing_failures = []
newly_processed_hashes = set()
# Step 3: Wait for this batch to complete
if wait_for_completion and task_id:
console.print(f"[cyan]Step 3: Waiting for batch {batch_start}-{batch_end} to complete...[/cyan]")
await wait_for_tasks_completion(client, base_url, [task_id], poll_interval)
console.print(f"[green]Batch {batch_start}-{batch_end} completed![/green]")
else:
console.print(f"[green]Batch {batch_start}-{batch_end} submitted successfully![/green]")
else:
total_processing_failures.extend([item[0] for item in successful_uploads])
console.print(f"[red]Batch {batch_start}-{batch_end} processing failed[/red]")
for file_path, file_hash, status in results:
if status == "success":
successful_files.append(file_path)
newly_processed_hashes.add(file_hash)
elif status == "upload_failed":
upload_failures.append(file_path)
elif status == "processing_failed":
processing_failures.append(file_path)
# Don't add to processed files if processing failed
# Record batch failures
total_upload_failures.extend(batch_upload_failures)
# Save newly processed files to record
if newly_processed_hashes:
all_processed_files = processed_files | newly_processed_hashes
save_processed_files(record_file, all_processed_files)
console.print(
f"[bold green]Updated processed files record with "
f"{len(newly_processed_hashes)} new entries.[/bold green]"
)
# Update processed files record after each batch
if all_successful_hashes:
all_processed_files = processed_files | all_successful_hashes
save_processed_files(record_file, all_processed_files)
console.print("[bold green]Batch operation complete.[/bold green]")
console.print(f" - [green]Successful:[/green] {len(successful_files)}")
console.print(f" - [red]Upload Failed:[/red] {len(upload_failures)}")
if upload_failures:
for f in upload_failures:
# Small delay between batches
if batch_end < len(files_to_upload):
console.print("[dim]Waiting 2 seconds before next batch...[/dim]")
await asyncio.sleep(2)
# Final summary
console.print("\n[bold green]=== All Batches Complete ===[/bold green]")
console.print(f" - [green]Files successfully processed:[/green] {len(total_processed_files)}")
console.print(f" - [red]Upload failures:[/red] {len(total_upload_failures)}")
if total_upload_failures:
for f in total_upload_failures:
console.print(f" - {f}")
console.print(f" - [yellow]Processing Failed:[/yellow] {len(processing_failures)}")
if processing_failures:
for f in processing_failures:
console.print(f" - [yellow]Processing failures:[/yellow] {len(total_processing_failures)}")
if total_processing_failures:
for f in total_processing_failures:
console.print(f" - {f}")
asyncio.run(run())
@ -386,11 +500,13 @@ def upload(
uv run scripts/batch_upload.py upload \
--db-id your_kb_id \
--directory path/to/your/data \
--pattern "*.docx" \
--pattern "*.docx" --pattern "*.pdf" --pattern "*.html" \
--base-url http://127.0.0.1:5050/api \
--username your_username \
--password your_password \
--concurrency 4 \
--batch-size 20 \
--wait-for-completion \
--poll-interval 5 \
--recursive \
--record-file scripts/tmp/batch_processed_files.txt
"""

View File

@ -73,18 +73,19 @@
<div class="task-card-footer">
<div class="task-card-timestamps">
<span v-if="task.started_at">开始: {{ formatTime(task.started_at, 'short') }}</span>
<span v-if="task.completed_at">完成: {{ formatTime(task.completed_at, 'short') }}</span>
<span v-if="task.started_at">开始: {{ formatTime(task.started_at) }}</span>
<span v-if="task.completed_at">完成: {{ formatTime(task.completed_at) }}</span>
<span v-if="!task.started_at">创建: {{ formatTime(task.created_at, 'short') }}</span>
</div>
<div class="task-card-actions">
<a-button type="link" size="small" @click="handleDetail(task.id)">
<a-button type="link" size="small" @click="handleDetail(task.id)" style="color: var(--gray-500);">
详情
</a-button>
<a-button
type="link"
size="small"
danger
v-if="canCancel(task)"
:disabled="!canCancel(task)"
@click="handleCancel(task.id)"
>
@ -394,18 +395,18 @@ function canCancel(task) {
.task-card {
background: #ffffff;
border: 1px solid var(--gray-200);
border: 1px solid var(--gray-100);
border-radius: 12px;
padding: 16px 18px;
transition: all 0.2s ease;
display: flex;
flex-direction: column;
gap: 12px;
gap: 8px;
// box-shadow: 0 1px 3px rgba(0, 0, 0, 0.05);
}
.task-card:hover {
border-color: rgba(59, 130, 246, 0.3);
border-color: var(--gray-200);;
}
.task-card--active {
@ -413,7 +414,7 @@ function canCancel(task) {
}
.task-card--success {
background: linear-gradient(to bottom, #ffffff, #f7fff9);
background: linear-gradient(to bottom, #ffffff, #fafffb);
}
.task-card--failed {
@ -516,7 +517,7 @@ function canCancel(task) {
.task-card-timestamps {
display: flex;
flex-direction: row;
gap: 6px;
gap: 10px;
font-size: 12px;
color: #94a3b8;
}
@ -555,7 +556,7 @@ function canCancel(task) {
width: 22px;
height: 22px;
border-radius: 50%;
background: #dcfce7;
background: #f3fff8;
font-size: 14px;
}

View File

@ -145,7 +145,6 @@ const mainList = [{
class="nav-item task-center"
:class="{ active: isDrawerOpen }"
@click="taskerStore.openDrawer()"
v-if="activeTaskCount > 0"
>
<a-tooltip placement="right">
<template #title>任务中心</template>