306 lines
12 KiB
Python
306 lines
12 KiB
Python
"""
|
||
恢复处理中断的索引任务
|
||
用于服务重启后继续处理pending/processing状态的任务
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到路径
|
||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
||
from sqlalchemy import select
|
||
from app.core.database import async_session_maker
|
||
from app.models.entities import IndexJob, Document, IndexJobStatus, DocumentStatus
|
||
from app.api.admin.kb import _index_document
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
async def resume_pending_jobs():
|
||
"""恢复所有pending和processing状态的任务"""
|
||
async with async_session_maker() as session:
|
||
# 查询所有未完成的任务
|
||
result = await session.execute(
|
||
select(IndexJob).where(
|
||
IndexJob.status.in_([IndexJobStatus.PENDING.value, IndexJobStatus.PROCESSING.value])
|
||
)
|
||
)
|
||
pending_jobs = result.scalars().all()
|
||
|
||
if not pending_jobs:
|
||
logger.info("没有需要恢复的任务")
|
||
return
|
||
|
||
logger.info(f"发现 {len(pending_jobs)} 个未完成的任务")
|
||
|
||
for job in pending_jobs:
|
||
try:
|
||
# 获取关联的文档
|
||
doc_result = await session.execute(
|
||
select(Document).where(Document.id == job.doc_id)
|
||
)
|
||
doc = doc_result.scalar_one_or_none()
|
||
|
||
if not doc:
|
||
logger.error(f"找不到文档: {job.doc_id}")
|
||
continue
|
||
|
||
if not doc.file_path or not Path(doc.file_path).exists():
|
||
logger.error(f"文档文件不存在: {doc.file_path}")
|
||
# 标记为失败
|
||
job.status = IndexJobStatus.FAILED.value
|
||
job.error_msg = "文档文件不存在"
|
||
doc.status = DocumentStatus.FAILED.value
|
||
doc.error_msg = "文档文件不存在"
|
||
await session.commit()
|
||
continue
|
||
|
||
logger.info(f"恢复处理: job_id={job.id}, doc_id={doc.id}, file={doc.file_name}")
|
||
|
||
# 读取文件内容
|
||
with open(doc.file_path, 'rb') as f:
|
||
file_content = f.read()
|
||
|
||
# 重置任务状态为pending
|
||
job.status = IndexJobStatus.PENDING.value
|
||
job.progress = 0
|
||
job.error_msg = None
|
||
await session.commit()
|
||
|
||
# 启动后台任务处理
|
||
# 注意：这里我们直接调用，不使用background_tasks
|
||
await process_job(
|
||
tenant_id=job.tenant_id,
|
||
kb_id=doc.kb_id,
|
||
job_id=str(job.id),
|
||
doc_id=str(doc.id),
|
||
file_content=file_content,
|
||
filename=doc.file_name,
|
||
metadata=doc.doc_metadata or {}
|
||
)
|
||
|
||
logger.info(f"任务处理完成: job_id={job.id}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"处理任务失败: job_id={job.id}, error={e}")
|
||
# 标记为失败
|
||
job.status = IndexJobStatus.FAILED.value
|
||
job.error_msg = str(e)
|
||
if doc:
|
||
doc.status = DocumentStatus.FAILED.value
|
||
doc.error_msg = str(e)
|
||
await session.commit()
|
||
|
||
logger.info("所有任务处理完成")
|
||
|
||
|
||
async def process_job(tenant_id: str, kb_id: str, job_id: str, doc_id: str,
|
||
file_content: bytes, filename: str, metadata: dict):
|
||
"""
|
||
处理单个索引任务
|
||
复制自 _index_document 函数
|
||
"""
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
from qdrant_client.models import PointStruct
|
||
|
||
from app.core.qdrant_client import get_qdrant_client
|
||
from app.services.document import DocumentParseException, UnsupportedFormatError, parse_document
|
||
from app.services.embedding import get_embedding_provider
|
||
from app.services.kb import KBService
|
||
from app.api.admin.kb import chunk_text_by_lines, TextChunk
|
||
|
||
logger.info(f"[RESUME] Starting indexing: tenant={tenant_id}, kb_id={kb_id}, job_id={job_id}, doc_id={doc_id}")
|
||
|
||
async with async_session_maker() as session:
|
||
kb_service = KBService(session)
|
||
try:
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=10
|
||
)
|
||
await session.commit()
|
||
|
||
parse_result = None
|
||
text = None
|
||
file_ext = Path(filename or "").suffix.lower()
|
||
logger.info(f"[RESUME] File extension: {file_ext}, content size: {len(file_content)} bytes")
|
||
|
||
text_extensions = {".txt", ".md", ".markdown", ".rst", ".log", ".json", ".xml", ".yaml", ".yml"}
|
||
|
||
if file_ext in text_extensions or not file_ext:
|
||
logger.info("[RESUME] Treating as text file")
|
||
text = None
|
||
for encoding in ["utf-8", "gbk", "gb2312", "gb18030", "big5", "utf-16", "latin-1"]:
|
||
try:
|
||
text = file_content.decode(encoding)
|
||
logger.info(f"[RESUME] Successfully decoded with encoding: {encoding}")
|
||
break
|
||
except (UnicodeDecodeError, LookupError):
|
||
continue
|
||
|
||
if text is None:
|
||
text = file_content.decode("utf-8", errors="replace")
|
||
else:
|
||
logger.info("[RESUME] Binary file detected, will parse with document parser")
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=15
|
||
)
|
||
await session.commit()
|
||
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp_file:
|
||
tmp_file.write(file_content)
|
||
tmp_path = tmp_file.name
|
||
|
||
logger.info(f"[RESUME] Temp file created: {tmp_path}")
|
||
|
||
try:
|
||
logger.info(f"[RESUME] Starting document parsing for {file_ext}...")
|
||
parse_result = parse_document(tmp_path)
|
||
text = parse_result.text
|
||
logger.info(
|
||
f"[RESUME] Parsed document SUCCESS: {filename}, chars={len(text)}"
|
||
)
|
||
except UnsupportedFormatError as e:
|
||
logger.error(f"[RESUME] UnsupportedFormatError: {e}")
|
||
text = file_content.decode("utf-8", errors="ignore")
|
||
except DocumentParseException as e:
|
||
logger.error(f"[RESUME] DocumentParseException: {e}")
|
||
text = file_content.decode("utf-8", errors="ignore")
|
||
except Exception as e:
|
||
logger.error(f"[RESUME] Unexpected parsing error: {type(e).__name__}: {e}")
|
||
text = file_content.decode("utf-8", errors="ignore")
|
||
finally:
|
||
Path(tmp_path).unlink(missing_ok=True)
|
||
|
||
logger.info(f"[RESUME] Final text length: {len(text)} chars")
|
||
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=20
|
||
)
|
||
await session.commit()
|
||
|
||
logger.info("[RESUME] Getting embedding provider...")
|
||
embedding_provider = await get_embedding_provider()
|
||
logger.info(f"[RESUME] Embedding provider: {type(embedding_provider).__name__}")
|
||
|
||
all_chunks: list[TextChunk] = []
|
||
|
||
if parse_result and parse_result.pages:
|
||
logger.info(f"[RESUME] PDF with {len(parse_result.pages)} pages")
|
||
for page in parse_result.pages:
|
||
page_chunks = chunk_text_by_lines(
|
||
page.text,
|
||
min_line_length=10,
|
||
source=filename,
|
||
)
|
||
for pc in page_chunks:
|
||
pc.page = page.page
|
||
all_chunks.extend(page_chunks)
|
||
else:
|
||
logger.info("[RESUME] Using line-based chunking")
|
||
all_chunks = chunk_text_by_lines(
|
||
text,
|
||
min_line_length=10,
|
||
source=filename,
|
||
)
|
||
|
||
logger.info(f"[RESUME] Total chunks: {len(all_chunks)}")
|
||
|
||
qdrant = await get_qdrant_client()
|
||
await qdrant.ensure_kb_collection_exists(tenant_id, kb_id, use_multi_vector=True)
|
||
|
||
from app.services.embedding.nomic_provider import NomicEmbeddingProvider
|
||
use_multi_vector = isinstance(embedding_provider, NomicEmbeddingProvider)
|
||
logger.info(f"[RESUME] Using multi-vector format: {use_multi_vector}")
|
||
|
||
import uuid
|
||
points = []
|
||
total_chunks = len(all_chunks)
|
||
doc_metadata = metadata or {}
|
||
|
||
for i, chunk in enumerate(all_chunks):
|
||
payload = {
|
||
"text": chunk.text,
|
||
"source": doc_id,
|
||
"kb_id": kb_id,
|
||
"chunk_index": i,
|
||
"start_token": chunk.start_token,
|
||
"end_token": chunk.end_token,
|
||
"metadata": doc_metadata,
|
||
}
|
||
if chunk.page is not None:
|
||
payload["page"] = chunk.page
|
||
if chunk.source:
|
||
payload["filename"] = chunk.source
|
||
|
||
if use_multi_vector:
|
||
embedding_result = await embedding_provider.embed_document(chunk.text)
|
||
points.append({
|
||
"id": str(uuid.uuid4()),
|
||
"vector": {
|
||
"full": embedding_result.embedding_full,
|
||
"dim_256": embedding_result.embedding_256,
|
||
"dim_512": embedding_result.embedding_512,
|
||
},
|
||
"payload": payload,
|
||
})
|
||
else:
|
||
embedding = await embedding_provider.embed(chunk.text)
|
||
points.append(
|
||
PointStruct(
|
||
id=str(uuid.uuid4()),
|
||
vector=embedding,
|
||
payload=payload,
|
||
)
|
||
)
|
||
|
||
progress = 20 + int((i + 1) / total_chunks * 70)
|
||
if i % 10 == 0 or i == total_chunks - 1:
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.PROCESSING.value, progress=progress
|
||
)
|
||
await session.commit()
|
||
|
||
if points:
|
||
logger.info(f"[RESUME] Upserting {len(points)} vectors to Qdrant...")
|
||
if use_multi_vector:
|
||
await qdrant.upsert_multi_vector(tenant_id, points, kb_id=kb_id)
|
||
else:
|
||
await qdrant.upsert_vectors(tenant_id, points, kb_id=kb_id)
|
||
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.COMPLETED.value, progress=100
|
||
)
|
||
await session.commit()
|
||
|
||
logger.info(
|
||
f"[RESUME] COMPLETED: tenant={tenant_id}, kb_id={kb_id}, "
|
||
f"job_id={job_id}, chunks={len(all_chunks)}"
|
||
)
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
logger.error(f"[RESUME] FAILED: {e}\n{traceback.format_exc()}")
|
||
await session.rollback()
|
||
async with async_session_maker() as error_session:
|
||
kb_service = KBService(error_session)
|
||
await kb_service.update_job_status(
|
||
tenant_id, job_id, IndexJobStatus.FAILED.value,
|
||
progress=0, error_msg=str(e)
|
||
)
|
||
await error_session.commit()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
logger.info("开始恢复索引任务...")
|
||
asyncio.run(resume_pending_jobs())
|
||
logger.info("恢复脚本执行完成")
|