145 lines
5.6 KiB
Python
145 lines
5.6 KiB
Python
"""
|
||
PDF内容提取服务 - 增强版
|
||
提取更多内容，更深度的分析
|
||
"""
|
||
import re
|
||
from typing import List, Dict, Optional, Tuple
|
||
import fitz
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from app.config import settings
|
||
from app.utils.logger import logger
|
||
from app.models import Report, ExtractedContent
|
||
|
||
|
||
class PDFExtractor:
|
||
"""PDF内容提取器 - 增强版"""
|
||
|
||
def __init__(self):
|
||
self.keywords = settings.EXTRACT_KEYWORDS
|
||
# 增加更多提取行数（上下文更丰富）
|
||
self.context_lines_before = 5
|
||
self.context_lines_after = 80 # 增加到80行，获取更完整的章节内容
|
||
self.min_content_length = 200 # 最小内容长度
|
||
|
||
def extract_text_from_pdf(self, pdf_path: str) -> Tuple[str, int]:
|
||
"""从PDF提取全部文本"""
|
||
try:
|
||
doc = fitz.open(pdf_path)
|
||
text = ""
|
||
page_count = doc.page_count
|
||
for page in doc:
|
||
text += page.get_text()
|
||
doc.close()
|
||
logger.info(f"PDF提取完成: {page_count} 页, {len(text)} 字符")
|
||
return text, page_count
|
||
except Exception as e:
|
||
logger.error(f"PDF提取失败: {e}")
|
||
return "", 0
|
||
|
||
def extract_sections_by_keywords(self, text: str) -> List[Dict]:
|
||
"""根据关键词提取相关章节 - 增强版，提取更多内容"""
|
||
sections = []
|
||
lines = text.split('\n')
|
||
total_lines = len(lines)
|
||
|
||
# 记录已匹配的行范围，避免重复提取
|
||
matched_ranges = []
|
||
|
||
for keyword in self.keywords:
|
||
for i, line in enumerate(lines):
|
||
if keyword in line:
|
||
# 检查是否在已匹配范围内
|
||
is_overlapping = False
|
||
for start_r, end_r in matched_ranges:
|
||
if start_r <= i <= end_r:
|
||
is_overlapping = True
|
||
break
|
||
|
||
if is_overlapping:
|
||
continue
|
||
|
||
# 扩大提取范围
|
||
start = max(0, i - self.context_lines_before)
|
||
end = min(total_lines, i + self.context_lines_after)
|
||
|
||
# 智能查找章节结束位置（遇到下一个大标题时停止）
|
||
for j in range(i + 10, end):
|
||
next_line = lines[j].strip()
|
||
# 如果遇到新的章节标题（如"第X节"、"（X）"等），停止
|
||
if re.match(r'^[（(]?[一二三四五六七八九十\d]+[)）]', next_line):
|
||
if len(next_line) < 50: # 确保是标题而非正文
|
||
end = j
|
||
break
|
||
|
||
content_lines = [l.strip() for l in lines[start:end] if l.strip()]
|
||
content = '\n'.join(content_lines)
|
||
|
||
if len(content) > self.min_content_length:
|
||
sections.append({
|
||
"keyword": keyword,
|
||
"section_name": line.strip()[:100],
|
||
"content": content,
|
||
"char_count": len(content),
|
||
"line_start": start,
|
||
"line_end": end
|
||
})
|
||
matched_ranges.append((start, end))
|
||
|
||
dedupe_sections = self._dedupe_sections(sections)
|
||
logger.info(f"提取到 {len(dedupe_sections)} 个章节")
|
||
return dedupe_sections
|
||
|
||
def _dedupe_sections(self, sections: List[Dict]) -> List[Dict]:
|
||
"""去重并按内容丰富度排序"""
|
||
seen = set()
|
||
result = []
|
||
for s in sections:
|
||
# 使用前300字符作为去重依据
|
||
key = s["content"][:300]
|
||
if key not in seen:
|
||
seen.add(key)
|
||
result.append(s)
|
||
|
||
# 按内容长度排序，优先保留内容更丰富的章节
|
||
result.sort(key=lambda x: x["char_count"], reverse=True)
|
||
return result
|
||
|
||
async def extract_and_save(self, db: AsyncSession, report: Report) -> List[ExtractedContent]:
|
||
"""提取并保存到数据库"""
|
||
import asyncio
|
||
|
||
if not report.local_path:
|
||
logger.warning(f"报告无本地路径: {report.title}")
|
||
return []
|
||
|
||
# 将 PDF 提取这种阻塞 CPU 的任务放入线程池
|
||
logger.info(f"开启线程提取 PDF: {report.title}")
|
||
text, page_count = await asyncio.to_thread(self.extract_text_from_pdf, report.local_path)
|
||
if not text:
|
||
return []
|
||
|
||
# 将关键词匹配这种阻塞 CPU 的任务也放入线程池
|
||
logger.info(f"开启线程提取章节: {report.title}")
|
||
sections = await asyncio.to_thread(self.extract_sections_by_keywords, text)
|
||
|
||
contents = []
|
||
for section in sections:
|
||
content = ExtractedContent(
|
||
report_id=report.id,
|
||
section_name=section["section_name"],
|
||
section_keyword=section["keyword"],
|
||
content=section["content"],
|
||
char_count=section["char_count"]
|
||
)
|
||
db.add(content)
|
||
contents.append(content)
|
||
|
||
report.is_extracted = True
|
||
await db.commit()
|
||
|
||
logger.info(f"保存 {len(contents)} 个章节到数据库")
|
||
return contents
|
||
|
||
|
||
pdf_extractor = PDFExtractor()
|