197 lines
6.7 KiB
Python
197 lines
6.7 KiB
Python
"""知识文档的安全上传、解析、分块和索引状态编排。"""
|
||
|
||
from django.db import transaction
|
||
from django.utils import timezone
|
||
|
||
from apps.audit.models import OperationAuditLog
|
||
from apps.knowledge.models import KnowledgeDocument
|
||
from apps.knowledge.services import rebuild_document_chunks
|
||
from apps.tickets.materials import (
|
||
MaterialParseError,
|
||
MaterialScanError,
|
||
MaterialValidationError,
|
||
PreparedMaterial,
|
||
parse_material,
|
||
prepare_material,
|
||
scan_bytes,
|
||
)
|
||
from apps.tickets.storage import delete_object, get_object, put_object
|
||
|
||
|
||
def prepare_document(raw: bytes, file_name: str, content_type: str) -> PreparedMaterial:
|
||
prepared = prepare_material(raw, file_name, content_type)
|
||
if prepared.material_type == "image":
|
||
raise MaterialValidationError(415, "知识文档暂不支持图片，请上传可提取文本的文档")
|
||
return prepared
|
||
|
||
|
||
def create_uploaded_document(
|
||
*,
|
||
knowledge_base_id: int,
|
||
prepared: PreparedMaterial,
|
||
operator=None,
|
||
) -> KnowledgeDocument:
|
||
"""先保存原件再创建文档记录，提交后异步执行扫描、解析和索引。"""
|
||
|
||
document = KnowledgeDocument(
|
||
knowledge_base_id=knowledge_base_id,
|
||
file_name=prepared.file_name,
|
||
file_type=prepared.material_type,
|
||
extension=prepared.extension,
|
||
mime_type=prepared.mime_type,
|
||
file_size=len(prepared.data),
|
||
sha256=prepared.sha256,
|
||
scan_status="pending",
|
||
parse_status="pending",
|
||
created_by=operator,
|
||
)
|
||
document.storage_path = (
|
||
f"knowledge-documents/{knowledge_base_id}/{document.public_id}."
|
||
f"{prepared.extension}"
|
||
)
|
||
put_object(document.storage_path, prepared.data, prepared.mime_type)
|
||
try:
|
||
with transaction.atomic():
|
||
document.save()
|
||
OperationAuditLog.objects.create(
|
||
user=operator,
|
||
module="knowledge",
|
||
action="document_uploaded",
|
||
target_type="knowledge_document",
|
||
target_id=str(document.public_id),
|
||
request_snapshot={
|
||
"knowledge_base_id": knowledge_base_id,
|
||
"file_name": document.file_name,
|
||
"file_size": document.file_size,
|
||
"sha256": document.sha256,
|
||
},
|
||
result="success",
|
||
)
|
||
except Exception:
|
||
try:
|
||
delete_object(document.storage_path)
|
||
except Exception:
|
||
pass
|
||
raise
|
||
return document
|
||
|
||
|
||
def _mark_failed(
|
||
document: KnowledgeDocument,
|
||
*,
|
||
scan_status: str | None = None,
|
||
scan_result: dict | None = None,
|
||
message: str,
|
||
) -> None:
|
||
if scan_status is not None:
|
||
document.scan_status = scan_status
|
||
if scan_result is not None:
|
||
document.scan_result = scan_result
|
||
document.parse_status = "failed"
|
||
document.parse_error = message
|
||
document.indexed_at = None
|
||
document.save()
|
||
OperationAuditLog.objects.create(
|
||
user=document.created_by,
|
||
module="knowledge",
|
||
action="document_processing_failed",
|
||
target_type="knowledge_document",
|
||
target_id=str(document.public_id),
|
||
request_snapshot={
|
||
"scan_status": document.scan_status,
|
||
"error_type": message.split("：", 1)[0],
|
||
},
|
||
result="failed",
|
||
)
|
||
|
||
|
||
def process_knowledge_document(document_id: int) -> KnowledgeDocument:
|
||
"""处理单个知识文档；失败保留原件和明确状态供人工重试。"""
|
||
|
||
with transaction.atomic():
|
||
document = (
|
||
KnowledgeDocument.objects.select_for_update()
|
||
.select_related("knowledge_base", "created_by")
|
||
.get(id=document_id)
|
||
)
|
||
if document.parse_status != "pending" or document.scan_status != "pending":
|
||
return document
|
||
document.processing_attempts += 1
|
||
document.scan_status = "scanning"
|
||
document.scan_result = {}
|
||
document.parse_error = ""
|
||
document.indexed_at = None
|
||
document.save()
|
||
try:
|
||
raw = get_object(document.storage_path)
|
||
except Exception:
|
||
_mark_failed(document, message="原始文件读取失败：对象存储暂时不可用")
|
||
return document
|
||
try:
|
||
scan_result = scan_bytes(raw)
|
||
except MaterialScanError as error:
|
||
_mark_failed(
|
||
document,
|
||
scan_status="failed",
|
||
scan_result={"status": "failed", "message": str(error)},
|
||
message=f"恶意文件扫描失败：{error}",
|
||
)
|
||
return document
|
||
except Exception:
|
||
_mark_failed(
|
||
document,
|
||
scan_status="failed",
|
||
scan_result={"status": "failed"},
|
||
message="恶意文件扫描失败：扫描器发生未预期错误",
|
||
)
|
||
return document
|
||
document.scan_status = scan_result["status"]
|
||
document.scan_result = scan_result
|
||
if document.scan_status != "clean":
|
||
_mark_failed(
|
||
document,
|
||
scan_status=document.scan_status,
|
||
scan_result=scan_result,
|
||
message="恶意文件扫描未通过：原件已隔离保留",
|
||
)
|
||
return document
|
||
document.parse_status = "parsing"
|
||
document.save(update_fields=["scan_status", "scan_result", "parse_status", "updated_at"])
|
||
try:
|
||
content = parse_material(raw, document.extension)
|
||
if not content.strip():
|
||
raise MaterialParseError("文档中未提取到可索引文本")
|
||
except (MaterialParseError, MaterialValidationError) as error:
|
||
_mark_failed(document, message=f"文档解析失败：{error}")
|
||
return document
|
||
document.content_text = content
|
||
document.parse_status = "indexing"
|
||
document.parse_error = ""
|
||
document.save(update_fields=["content_text", "parse_status", "parse_error", "updated_at"])
|
||
try:
|
||
rebuild_document_chunks(document)
|
||
except Exception as error:
|
||
_mark_failed(
|
||
document,
|
||
message=f"文档索引失败：{type(error).__name__}",
|
||
)
|
||
return document
|
||
document.parse_status = "success"
|
||
document.parse_error = ""
|
||
document.indexed_at = timezone.now()
|
||
document.save(update_fields=["parse_status", "parse_error", "indexed_at", "updated_at"])
|
||
OperationAuditLog.objects.create(
|
||
user=document.created_by,
|
||
module="knowledge",
|
||
action="document_indexed",
|
||
target_type="knowledge_document",
|
||
target_id=str(document.public_id),
|
||
request_snapshot={
|
||
"knowledge_base_id": document.knowledge_base_id,
|
||
"chunk_count": document.chunks.count(),
|
||
"processing_attempts": document.processing_attempts,
|
||
},
|
||
result="success",
|
||
)
|
||
return document
|