128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
|
"""导入主流程：下载 → 解析 → 清洗 → 切片 → embedding → Qdrant → AI 候选"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import uuid
|
|||
|
|
from parser import download_file, parse_document
|
|||
|
|
from chunker import chunk_document
|
|||
|
|
from embedder import embed_batch
|
|||
|
|
from indexer import upsert_points
|
|||
|
|
from candidate_generator import generate_candidates
|
|||
|
|
from api_client import (
|
|||
|
|
heartbeat as send_heartbeat,
|
|||
|
|
update_job_status,
|
|||
|
|
save_chunks,
|
|||
|
|
save_candidates,
|
|||
|
|
get_job_detail,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def run_import(job: dict):
|
|||
|
|
"""执行完整的文档导入流程"""
|
|||
|
|
job_id = job["id"]
|
|||
|
|
source_id = job.get("sourceId") or job.get("source_id")
|
|||
|
|
user_id = job["userId"] or job.get("user_id")
|
|||
|
|
kb_id = job["knowledgeBaseId"] or job.get("knowledge_base_id")
|
|||
|
|
file_id = job.get("fileId") or job.get("file_id")
|
|||
|
|
|
|||
|
|
if not source_id:
|
|||
|
|
raise ValueError(f"任务 {job_id} 缺少 sourceId")
|
|||
|
|
|
|||
|
|
# 获取 source 详情（从 NestJS）
|
|||
|
|
detail = await get_job_detail(job_id)
|
|||
|
|
source = (detail or {}).get("source", {}) if detail else {}
|
|||
|
|
mime_type = source.get("mimeType") or source.get("mime_type") or "text/plain"
|
|||
|
|
original_filename = source.get("originalFilename") or source.get("original_filename") or "unknown"
|
|||
|
|
|
|||
|
|
tmp_dir = f"/data/tmp/imports/{job_id}"
|
|||
|
|
file_path = os.path.join(tmp_dir, original_filename)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 1. 下载文件
|
|||
|
|
await update_job_status(job_id, "DOWNLOADING", {"progress": 5})
|
|||
|
|
file_url = source.get("downloadUrl") or (detail or {}).get("downloadUrl", "")
|
|||
|
|
if file_url:
|
|||
|
|
await download_file(file_url, file_path)
|
|||
|
|
|
|||
|
|
# 2. 解析
|
|||
|
|
await update_job_status(job_id, "PARSING", {"progress": 20})
|
|||
|
|
text = await parse_document(file_path, mime_type)
|
|||
|
|
|
|||
|
|
# 如果文件不在本地（纯文本导入），直接从 source/import 中取文本
|
|||
|
|
if not text and (job.get("rawText") or source.get("rawText")):
|
|||
|
|
text = job.get("rawText", "") or source.get("rawText", "")
|
|||
|
|
|
|||
|
|
if not text or len(text.strip()) < 10:
|
|||
|
|
raise ValueError("文档解析后内容过少，可能为空白或损坏文件")
|
|||
|
|
|
|||
|
|
# 3. 清洗
|
|||
|
|
await update_job_status(job_id, "CLEANING", {"progress": 40, "textLength": len(text)})
|
|||
|
|
|
|||
|
|
# 4. 切片
|
|||
|
|
await update_job_status(job_id, "CHUNKING", {"progress": 50})
|
|||
|
|
source_type = source.get("type") or "text"
|
|||
|
|
chunks = chunk_document(text, source_type)
|
|||
|
|
|
|||
|
|
# 5. Embedding
|
|||
|
|
await update_job_status(job_id, "EMBEDDING", {"progress": 60})
|
|||
|
|
texts = [c["content"] for c in chunks]
|
|||
|
|
vectors = await embed_batch(texts)
|
|||
|
|
|
|||
|
|
# 6. Qdrant 索引
|
|||
|
|
await update_job_status(job_id, "INDEXING", {"progress": 80})
|
|||
|
|
points = []
|
|||
|
|
chunk_records = []
|
|||
|
|
for i, (chunk, vec) in enumerate(zip(chunks, vectors)):
|
|||
|
|
chunk_id = f"chunk_{source_id}_{i}"
|
|||
|
|
points.append({
|
|||
|
|
"id": chunk_id,
|
|||
|
|
"vector": vec,
|
|||
|
|
"payload": {
|
|||
|
|
"userId": user_id,
|
|||
|
|
"knowledgeBaseId": kb_id,
|
|||
|
|
"sourceId": source_id,
|
|||
|
|
"chunkId": chunk_id,
|
|||
|
|
"pageNumber": chunk.get("pageNumber"),
|
|||
|
|
"sectionTitle": chunk.get("sectionTitle", ""),
|
|||
|
|
"deleted": False,
|
|||
|
|
},
|
|||
|
|
})
|
|||
|
|
chunk_records.append({
|
|||
|
|
"userId": user_id,
|
|||
|
|
"knowledgeBaseId": kb_id,
|
|||
|
|
"sourceId": source_id,
|
|||
|
|
"content": chunk["content"],
|
|||
|
|
"chunkIndex": chunk["chunkIndex"],
|
|||
|
|
"pageNumber": chunk.get("pageNumber"),
|
|||
|
|
"sectionTitle": chunk.get("sectionTitle", ""),
|
|||
|
|
"tokenCount": len(chunk["content"]),
|
|||
|
|
"externalVectorId": chunk_id,
|
|||
|
|
"embeddingModel": "bge-m3",
|
|||
|
|
"embeddingStatus": "COMPLETED",
|
|||
|
|
"metadataJson": {"chunkType": chunk.get("chunkType", "text")},
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
await upsert_points(points)
|
|||
|
|
await save_chunks(chunk_records)
|
|||
|
|
|
|||
|
|
# 7. 生成候选知识点
|
|||
|
|
await update_job_status(job_id, "GENERATING_CANDIDATES", {"progress": 90})
|
|||
|
|
candidates = await generate_candidates(text)
|
|||
|
|
if candidates:
|
|||
|
|
await save_candidates(user_id, kb_id, source_id, job_id, candidates)
|
|||
|
|
|
|||
|
|
# 8. 完成
|
|||
|
|
await update_job_status(job_id, "COMPLETED", {"progress": 100})
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
await update_job_status(job_id, "FAILED_RETRYABLE", {
|
|||
|
|
"errorCode": "WORKER_ERROR",
|
|||
|
|
"errorMessage": str(e)[:500],
|
|||
|
|
})
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
finally:
|
|||
|
|
# 清理临时文件
|
|||
|
|
if os.path.exists(tmp_dir):
|
|||
|
|
import shutil
|
|||
|
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|