Codex
354 lines
16 KiB
Python
354 lines
16 KiB
Python
import io
|
|
import re
|
|
import zipfile
|
|
from typing import Any, Dict, Optional
|
|
from uuid import uuid4
|
|
from xml.etree import ElementTree
|
|
|
|
import pdfplumber
|
|
from fastapi import HTTPException, status
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models.knowledge_base import KnowledgeBase
|
|
from app.models.knowledge_base_document import KnowledgeBaseDocument
|
|
from app.repositories.knowledge_base_document_repository import KnowledgeBaseDocumentRepository
|
|
from app.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
|
|
MAX_TEXT_LENGTH = 20000
|
|
DOCX_NAMESPACE = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
|
|
BINARY_TEXT_PATTERN = re.compile(r"[\u4e00-\u9fffA-Za-z0-9，。；：、？！《》（）【】\-—_]+")
|
|
PDF_TEXT_MARKERS = ("合同", "违约", "免责", "付款", "保密", "续约")
|
|
SUPPORTED_TEXT_SUFFIXES = (".txt", ".md", ".csv", ".json")
|
|
DOCX_SUFFIX = ".docx"
|
|
PDF_SUFFIX = ".pdf"
|
|
KNOWLEDGE_BASE_DOCUMENT_STATUS_PENDING = "pending"
|
|
KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING = "processing"
|
|
KNOWLEDGE_BASE_DOCUMENT_STATUS_INDEXED = "indexed"
|
|
KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED = "failed"
|
|
|
|
|
|
class KnowledgeBaseService:
|
|
def __init__(self, db: Session):
|
|
self.db = db
|
|
self.repository = KnowledgeBaseRepository(db)
|
|
self.document_repository = KnowledgeBaseDocumentRepository(db)
|
|
|
|
def list_knowledge_bases(self, current_user: Dict[str, Any], page: int = 1, page_size: int = 20) -> dict:
|
|
self._require_admin(current_user)
|
|
tenant_id = current_user["tenant_id"]
|
|
normalized_page = max(page, 1)
|
|
normalized_page_size = max(min(page_size, 100), 1)
|
|
knowledge_bases = self.repository.list_paginated(tenant_id, normalized_page, normalized_page_size)
|
|
counts = self.document_repository.count_by_knowledge_base_ids(tenant_id, [item.id for item in knowledge_bases])
|
|
return {
|
|
"items": [self._to_item(item, counts.get(item.id, 0)) for item in knowledge_bases],
|
|
"page": normalized_page,
|
|
"page_size": normalized_page_size,
|
|
"total": self.repository.count(tenant_id),
|
|
}
|
|
|
|
def create_knowledge_base(self, current_user: Dict[str, Any], name: str, description: Optional[str]) -> dict:
|
|
self._require_admin(current_user)
|
|
normalized_name = name.strip()
|
|
normalized_description = description.strip() if description else None
|
|
knowledge_base = KnowledgeBase(
|
|
id=f"kb-{uuid4().hex[:12]}",
|
|
tenant_id=current_user["tenant_id"],
|
|
name=normalized_name,
|
|
description=normalized_description,
|
|
status="active",
|
|
)
|
|
created = self.repository.create(knowledge_base)
|
|
return self._to_item(created, 0)
|
|
|
|
def upload_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
file_name: str,
|
|
mime_type: Optional[str],
|
|
content: bytes,
|
|
) -> dict:
|
|
self._require_admin(current_user)
|
|
knowledge_base = self.repository.get_by_id(knowledge_base_id)
|
|
if knowledge_base is None or knowledge_base.tenant_id != current_user["tenant_id"]:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
|
|
|
|
document = KnowledgeBaseDocument(
|
|
id=f"kbdoc-{uuid4().hex[:12]}",
|
|
knowledge_base_id=knowledge_base.id,
|
|
tenant_id=current_user["tenant_id"],
|
|
file_name=file_name,
|
|
mime_type=mime_type,
|
|
file_size=len(content),
|
|
status=KNOWLEDGE_BASE_DOCUMENT_STATUS_PENDING,
|
|
content_excerpt=None,
|
|
error_message=None,
|
|
)
|
|
created = self.document_repository.create(document)
|
|
return self._process_document_content(created, content, mime_type, file_name)
|
|
|
|
def retry_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
document_id: str,
|
|
file_name: Optional[str] = None,
|
|
mime_type: Optional[str] = None,
|
|
content: Optional[bytes] = None,
|
|
) -> dict:
|
|
self._require_admin(current_user)
|
|
knowledge_base = self.repository.get_by_id(knowledge_base_id)
|
|
if knowledge_base is None or knowledge_base.tenant_id != current_user["tenant_id"]:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
|
|
|
|
document = self.document_repository.get_by_id(
|
|
document_id,
|
|
tenant_id=current_user["tenant_id"],
|
|
knowledge_base_id=knowledge_base_id,
|
|
)
|
|
if document is None:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="knowledge base document not found")
|
|
|
|
if content is not None:
|
|
document.file_name = file_name or document.file_name
|
|
document.mime_type = mime_type or document.mime_type
|
|
document.file_size = len(content)
|
|
return self._process_document_content(document, content, document.mime_type, document.file_name)
|
|
|
|
if document.content_excerpt:
|
|
content_bytes = document.content_excerpt.encode("utf-8")
|
|
return self._process_document_content(document, content_bytes, "text/plain", document.file_name)
|
|
|
|
document.status = KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED
|
|
document.content_excerpt = None
|
|
document.error_message = "重试需要重新上传可复制文本版本，当前记录未保留可重新索引的原文。"
|
|
document = self.document_repository.update(document)
|
|
return self._to_document_item(document)
|
|
|
|
def _process_document_content(
|
|
self,
|
|
document: KnowledgeBaseDocument,
|
|
content: bytes,
|
|
mime_type: Optional[str],
|
|
file_name: str,
|
|
) -> dict:
|
|
document.status = KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING
|
|
document.error_message = None
|
|
document = self.document_repository.update(document)
|
|
|
|
try:
|
|
text_content = self._extract_text_from_file(content, mime_type, file_name)
|
|
if not text_content:
|
|
raise ValueError("无法从该文档中提取可检索文本")
|
|
|
|
document.status = KNOWLEDGE_BASE_DOCUMENT_STATUS_INDEXED
|
|
document.content_excerpt = text_content
|
|
document.error_message = None
|
|
document = self.document_repository.update(document)
|
|
return self._to_document_item(document)
|
|
except Exception:
|
|
document.status = KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED
|
|
document.content_excerpt = None
|
|
document.error_message = "无法从该文档中提取可检索文本"
|
|
document = self.document_repository.update(document)
|
|
return self._to_document_item(document)
|
|
|
|
def list_documents(self, current_user: Dict[str, Any], knowledge_base_id: str, page: int = 1, page_size: int = 20) -> dict:
|
|
self._require_admin(current_user)
|
|
knowledge_base = self.repository.get_by_id(knowledge_base_id)
|
|
if knowledge_base is None or knowledge_base.tenant_id != current_user["tenant_id"]:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="knowledge base not found")
|
|
|
|
normalized_page = max(page, 1)
|
|
normalized_page_size = max(min(page_size, 100), 1)
|
|
tenant_id = current_user["tenant_id"]
|
|
documents = self.document_repository.list_paginated(tenant_id, knowledge_base_id, normalized_page, normalized_page_size)
|
|
return {
|
|
"items": [self._to_document_item(item) for item in documents],
|
|
"page": normalized_page,
|
|
"page_size": normalized_page_size,
|
|
"total": self.document_repository.count(tenant_id, knowledge_base_id),
|
|
}
|
|
|
|
def _require_admin(self, current_user: Dict[str, Any]):
|
|
if current_user.get("role") != "admin":
|
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="forbidden")
|
|
|
|
def _to_item(self, knowledge_base: KnowledgeBase, document_count: int) -> dict:
|
|
return {
|
|
"id": knowledge_base.id,
|
|
"name": knowledge_base.name,
|
|
"description": knowledge_base.description,
|
|
"status": knowledge_base.status,
|
|
"document_count": document_count,
|
|
"created_at": knowledge_base.created_at.isoformat(),
|
|
"updated_at": knowledge_base.updated_at.isoformat(),
|
|
}
|
|
|
|
def _extract_text_from_file(self, content: bytes, mime_type: Optional[str], file_name: str) -> Optional[str]:
|
|
if not content:
|
|
return None
|
|
|
|
normalized_name = (file_name or "").lower()
|
|
normalized_mime = (mime_type or "").lower()
|
|
|
|
try:
|
|
if normalized_mime.startswith("text/") or normalized_name.endswith(SUPPORTED_TEXT_SUFFIXES):
|
|
return self._normalize_extracted_text(content.decode("utf-8", errors="ignore"))
|
|
if normalized_name.endswith(DOCX_SUFFIX) or "wordprocessingml" in normalized_mime:
|
|
return self._extract_text_from_docx(content)
|
|
if normalized_name.endswith(PDF_SUFFIX) or normalized_mime == "application/pdf":
|
|
return self._extract_text_from_pdf(content)
|
|
except Exception:
|
|
return None
|
|
|
|
return self._extract_text_from_binary(content)
|
|
|
|
def _extract_text_from_docx(self, file_bytes: bytes) -> Optional[str]:
|
|
try:
|
|
with zipfile.ZipFile(io.BytesIO(file_bytes)) as archive:
|
|
xml_content = archive.read("word/document.xml")
|
|
except (KeyError, zipfile.BadZipFile):
|
|
return self._extract_text_from_binary(file_bytes)
|
|
|
|
root = ElementTree.fromstring(xml_content)
|
|
texts = [node.text for node in root.findall(".//w:t", DOCX_NAMESPACE) if node.text]
|
|
return self._normalize_extracted_text("\n".join(texts))
|
|
|
|
def _extract_text_from_pdf(self, file_bytes: bytes) -> Optional[str]:
|
|
extracted_text = self._extract_with_pdfplumber(file_bytes)
|
|
if extracted_text:
|
|
return extracted_text
|
|
return self._extract_text_from_binary(file_bytes)
|
|
|
|
def _extract_with_pdfplumber(self, file_bytes: bytes) -> Optional[str]:
|
|
try:
|
|
with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
|
|
text_parts = []
|
|
for page in pdf.pages:
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text_parts.append(page_text)
|
|
if not text_parts:
|
|
return None
|
|
normalized = self._normalize_extracted_text("\n".join(text_parts))
|
|
if not normalized:
|
|
return None
|
|
if any(marker in normalized for marker in PDF_TEXT_MARKERS) or len(normalized) > 100:
|
|
return normalized
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
def _extract_text_from_binary(self, file_bytes: bytes) -> Optional[str]:
|
|
decoded = file_bytes.decode("utf-8", errors="ignore")
|
|
matches = BINARY_TEXT_PATTERN.findall(decoded)
|
|
if not matches:
|
|
return None
|
|
return self._normalize_extracted_text("\n".join(matches))
|
|
|
|
def _build_document_status_detail(self, document: KnowledgeBaseDocument) -> str:
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_INDEXED:
|
|
return "文本提取与索引已完成，可直接用于检索和问答。"
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED:
|
|
return document.error_message or "文档处理失败，请检查文件内容或格式后重试。"
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING:
|
|
return "系统正在提取文本并写入检索索引。"
|
|
return "文档已接收，等待进入文本提取阶段。"
|
|
|
|
def _build_document_failure_reason(self, document: KnowledgeBaseDocument, parse_status: str, index_status: str) -> Optional[str]:
|
|
if document.status not in {KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED, KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING}:
|
|
return None
|
|
if index_status == "blocked":
|
|
return document.error_message or "文本提取尚未完成，索引阶段暂时阻塞。"
|
|
if parse_status == "failed":
|
|
return document.error_message or "未能从文档中提取到可检索文本。"
|
|
if index_status == "failed":
|
|
return document.error_message or "文本已提取，但写入索引阶段失败。"
|
|
return document.error_message or "文档处理失败。"
|
|
|
|
def _build_document_next_action(self, document: KnowledgeBaseDocument, parse_status: str, index_status: str) -> str:
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_INDEXED:
|
|
return "可直接进入问答，或继续补充更多制度与条款文档。"
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_PENDING:
|
|
return "等待系统进入文本提取阶段。"
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING:
|
|
if index_status == "blocked":
|
|
return "请等待文本提取完成；提取成功后系统会继续进入索引阶段。"
|
|
return "等待文本提取与索引完成，完成后即可用于检索。"
|
|
if parse_status == "failed":
|
|
return "请检查文件格式、内容编码或改为可复制文本后重新上传。"
|
|
if index_status == "failed":
|
|
return "请重新上传文档；若反复失败，需检查索引链路。"
|
|
return "请检查处理结果后重试。"
|
|
|
|
def _build_document_lifecycle(self, document: KnowledgeBaseDocument) -> dict:
|
|
parse_status = "pending"
|
|
index_status = "pending"
|
|
current_stage = "parse"
|
|
current_stage_label = "等待解析"
|
|
|
|
if document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_PROCESSING:
|
|
parse_status = "processing"
|
|
index_status = "blocked"
|
|
current_stage = "parse"
|
|
current_stage_label = "文本提取中"
|
|
elif document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_INDEXED:
|
|
parse_status = "completed"
|
|
index_status = "completed"
|
|
current_stage = "indexed"
|
|
current_stage_label = "已完成索引"
|
|
elif document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED:
|
|
parse_status = "failed"
|
|
index_status = "blocked"
|
|
current_stage = "failed"
|
|
current_stage_label = "文本提取失败"
|
|
|
|
return {
|
|
"current_stage": current_stage,
|
|
"current_stage_label": current_stage_label,
|
|
"status_detail": self._build_document_status_detail(document),
|
|
"parse_status": parse_status,
|
|
"index_status": index_status,
|
|
"failure_reason": self._build_document_failure_reason(document, parse_status, index_status),
|
|
"next_action": self._build_document_next_action(document, parse_status, index_status),
|
|
"retryable": document.status == KNOWLEDGE_BASE_DOCUMENT_STATUS_FAILED,
|
|
}
|
|
|
|
def _to_document_item(self, document: KnowledgeBaseDocument) -> dict:
|
|
raw_content = document.content_excerpt or ""
|
|
excerpt = raw_content
|
|
preview_truncated = False
|
|
if excerpt and len(excerpt) > 500:
|
|
excerpt = f"{excerpt[:500]}..."
|
|
preview_truncated = True
|
|
lifecycle = self._build_document_lifecycle(document)
|
|
return {
|
|
"id": document.id,
|
|
"knowledge_base_id": document.knowledge_base_id,
|
|
"file_name": document.file_name,
|
|
"mime_type": document.mime_type,
|
|
"file_size": document.file_size,
|
|
"status": document.status,
|
|
"content_excerpt": excerpt,
|
|
"error_message": document.error_message,
|
|
"extracted_text_length": len(raw_content),
|
|
"preview_truncated": preview_truncated,
|
|
"created_at": document.created_at.isoformat(),
|
|
"current_stage": lifecycle["current_stage"],
|
|
"current_stage_label": lifecycle["current_stage_label"],
|
|
"status_detail": lifecycle["status_detail"],
|
|
"parse_status": lifecycle["parse_status"],
|
|
"index_status": lifecycle["index_status"],
|
|
"failure_reason": lifecycle["failure_reason"],
|
|
"next_action": lifecycle["next_action"],
|
|
"retryable": lifecycle["retryable"],
|
|
}
|
|
|
|
def _normalize_extracted_text(self, value: str) -> Optional[str]:
|
|
normalized = "\n".join(line.strip() for line in value.splitlines() if line.strip())
|
|
if not normalized:
|
|
return None
|
|
return normalized[:MAX_TEXT_LENGTH]
|