230 lines
9.9 KiB
Python
230 lines
9.9 KiB
Python
"""简历两阶段并行提取
|
||
|
||
第一阶段：5路并行提取主表短字段 + 各子表标识名（极快，输出极短）。
|
||
第二阶段：N+1路并行提取每条子表记录的短字段 + description/summary 的「行号区间」。
|
||
description/summary 不再由 AI 照抄原文，AI 只返回行号区间字符串，由代码从原文单元数组切片还原。
|
||
最终组装为与原方案完全一致的 dict 结构（description 为 list[str]），上下游无感知。
|
||
"""
|
||
|
||
import asyncio
|
||
import re
|
||
import time
|
||
|
||
from langchain_core.output_parsers import StrOutputParser
|
||
from langchain_core.prompts import ChatPromptTemplate
|
||
|
||
from app.ai.model_config import ResumeExtractorModel
|
||
from app.ai.resume_extractor.prompts import (
|
||
OVERVIEW_PROFILE_PROMPT, OVERVIEW_EDUCATION_PROMPT, OVERVIEW_WORK_PROMPT,
|
||
OVERVIEW_PROJECT_PROMPT, OVERVIEW_COMPETITION_PROMPT,
|
||
DETAIL_PROFILE_PROMPT, DETAIL_EDUCATION_PROMPT, DETAIL_WORK_PROMPT,
|
||
DETAIL_INTERNSHIP_PROMPT, DETAIL_PROJECT_PROMPT, DETAIL_COMPETITION_PROMPT,
|
||
)
|
||
from app.core.logger import log
|
||
from app.tool.json_helper import parse_llm_json
|
||
|
||
|
||
# ==================== LLM 调用工具 ====================
|
||
|
||
def _build_chain(prompt: str):
|
||
"""构建提取链：prompt → LLM → 文本输出"""
|
||
return ChatPromptTemplate.from_messages([("system", prompt), ("human", "{text}")]) | ResumeExtractorModel.PARSE | StrOutputParser()
|
||
|
||
|
||
async def _safe_invoke(chain, inp: dict, label: str):
|
||
"""单个链调用，记录耗时，失败返回空"""
|
||
start = time.perf_counter()
|
||
try:
|
||
raw = await chain.ainvoke(inp)
|
||
log.info(f"AI提取[{label}]完成，耗时: {time.perf_counter() - start:.2f}s")
|
||
return parse_llm_json(raw)
|
||
except Exception as e:
|
||
log.warning(f"AI提取[{label}]失败，耗时: {time.perf_counter() - start:.2f}s，错误: {e}")
|
||
return None
|
||
|
||
|
||
def _number_segments(segments: list[str]) -> str:
|
||
"""将单元数组构造为带行号文本：每行 `[行号] 内容`，行号 0 开始"""
|
||
return "\n".join(f"[{i}] {seg}" for i, seg in enumerate(segments))
|
||
|
||
|
||
# 条目/序号行的起始特征：● ○ • · ▪ ■ ‣ ◆ 等项目符号，或 "1." "2、" "3，" "(4)" "①" 等序号。
|
||
# 视觉行解析下，一条 bullet 常被折成多行；只有「条目起始行」才另起段落，其余行视为折行续接。
|
||
_LIST_ITEM_RE = re.compile(
|
||
r"^\s*(?:"
|
||
r"[●○◦•·∙▪■□‣◆◇►▶*]" # 项目符号
|
||
r"|[-–—]\s" # 连字符 + 空格（markdown 风格）
|
||
r"|\d+\s*[.．、，)）]" # 阿拉伯数字 + 标点：1. / 2、 / 3，/ 4)
|
||
r"|[（(]\s*\d+\s*[）)]" # 括号数字：(1) （2）
|
||
r"|[①②③④⑤⑥⑦⑧⑨⑩⑪⑫⑬⑭⑮⑯⑰⑱⑲⑳]" # 圈数字
|
||
r")"
|
||
)
|
||
|
||
|
||
def _split_into_paragraphs(lines: list[str]) -> list[str]:
|
||
"""将一个行号区间内的多行拆分为段落列表：
|
||
|
||
- 以项目符号/序号开头的行 → 另起一个新段落；
|
||
- 其余行 → 视为上一段落的折行续接，直连拼接（无分隔符）。
|
||
这样即使 AI 把多条 bullet 塞进同一个区间，也能按条目正确换行。
|
||
"""
|
||
paragraphs: list[str] = []
|
||
current = ""
|
||
for line in lines:
|
||
if not current:
|
||
current = line
|
||
elif _LIST_ITEM_RE.match(line):
|
||
paragraphs.append(current)
|
||
current = line
|
||
else:
|
||
current += line
|
||
if current:
|
||
paragraphs.append(current)
|
||
return [p for p in paragraphs if p.strip()]
|
||
|
||
|
||
def _slice_ranges(range_str, segments: list[str]) -> list[str]:
|
||
"""按行号区间字符串切片还原段落：逗号分段，段内按条目符号/折行拆分。
|
||
|
||
range_str 形如 "3-6,7-9" / "5"；非法或越界 token 静默跳过/clamp。
|
||
返回段落数组（每个逗号段可能因含多条 bullet 而进一步拆成多个段落）。
|
||
"""
|
||
if not range_str or not isinstance(range_str, str):
|
||
return []
|
||
n = len(segments)
|
||
paragraphs: list[str] = []
|
||
for token in range_str.split(","):
|
||
token = token.strip()
|
||
if not token:
|
||
continue
|
||
if "-" in token:
|
||
a, _, b = token.partition("-")
|
||
a, b = a.strip(), b.strip()
|
||
if not a.isdigit() or not b.isdigit():
|
||
continue
|
||
start, end = int(a), int(b)
|
||
else:
|
||
if not token.isdigit():
|
||
continue
|
||
start = end = int(token)
|
||
if start > end:
|
||
start, end = end, start
|
||
start = max(0, start)
|
||
end = min(n - 1, end)
|
||
if start > n - 1:
|
||
continue
|
||
paragraphs.extend(_split_into_paragraphs(segments[start:end + 1]))
|
||
return paragraphs
|
||
|
||
|
||
# ==================== 第一阶段：概览 ====================
|
||
|
||
_overview_profile_chain = _build_chain(OVERVIEW_PROFILE_PROMPT)
|
||
_overview_education_chain = _build_chain(OVERVIEW_EDUCATION_PROMPT)
|
||
_overview_work_chain = _build_chain(OVERVIEW_WORK_PROMPT)
|
||
_overview_project_chain = _build_chain(OVERVIEW_PROJECT_PROMPT)
|
||
_overview_competition_chain = _build_chain(OVERVIEW_COMPETITION_PROMPT)
|
||
|
||
|
||
async def _extract_overview(numbered_text: str) -> dict:
|
||
"""第一阶段：5路并行提取概览信息"""
|
||
inp = {"text": numbered_text}
|
||
profile, edu_names, work_names, proj_names, comp_names = await asyncio.gather(
|
||
_safe_invoke(_overview_profile_chain, inp, "概览-个人信息"),
|
||
_safe_invoke(_overview_education_chain, inp, "概览-教育"),
|
||
_safe_invoke(_overview_work_chain, inp, "概览-工作实习"),
|
||
_safe_invoke(_overview_project_chain, inp, "概览-项目"),
|
||
_safe_invoke(_overview_competition_chain, inp, "概览-竞赛"),
|
||
)
|
||
return {
|
||
"profile": profile if isinstance(profile, dict) else {},
|
||
"education": edu_names if isinstance(edu_names, list) else [],
|
||
"work": work_names.get("work", []) if isinstance(work_names, dict) else [],
|
||
"internship": work_names.get("internship", []) if isinstance(work_names, dict) else [],
|
||
"project": proj_names if isinstance(proj_names, list) else [],
|
||
"competition": comp_names if isinstance(comp_names, list) else [],
|
||
}
|
||
|
||
|
||
# ==================== 第二阶段：详情 ====================
|
||
|
||
# 子表模块统一配置（单一数据源）：(模块键, 详情 prompt 模板, 日志标签)
|
||
# 概览/详情/组装/日志均复用此清单，新增子表只需在此追加一行。
|
||
_DETAIL_MODULES: tuple[tuple[str, str, str], ...] = (
|
||
("education", DETAIL_EDUCATION_PROMPT, "教育"),
|
||
("work", DETAIL_WORK_PROMPT, "工作"),
|
||
("internship", DETAIL_INTERNSHIP_PROMPT, "实习"),
|
||
("project", DETAIL_PROJECT_PROMPT, "项目"),
|
||
("competition", DETAIL_COMPETITION_PROMPT, "竞赛"),
|
||
)
|
||
_SUB_MODULES: tuple[str, ...] = tuple(m[0] for m in _DETAIL_MODULES)
|
||
|
||
|
||
async def _extract_detail(prompt_tpl: str, name: str, numbered_text: str, label: str) -> dict | None:
|
||
"""单条子表记录详情提取：用 name 替换 prompt 中的 {name}"""
|
||
chain = _build_chain(prompt_tpl.replace("{name}", name))
|
||
return await _safe_invoke(chain, {"text": numbered_text}, label)
|
||
|
||
|
||
async def _extract_all_details(overview: dict, numbered_text: str) -> dict:
|
||
"""第二阶段：根据概览结果，N+1路并行提取所有子表记录详情 + 个人信息补充"""
|
||
# 第 0 路固定为 profile 补充（skills/certificates/summaryRange），其余按子表记录展开
|
||
tasks = [_extract_detail(DETAIL_PROFILE_PROMPT, "", numbered_text, "详情-个人信息补充")]
|
||
task_modules = ["profile_extra"]
|
||
for module, prompt_tpl, label in _DETAIL_MODULES:
|
||
for name in overview[module]:
|
||
tasks.append(_extract_detail(prompt_tpl, name, numbered_text, f"详情-{label}-{name}"))
|
||
task_modules.append(module)
|
||
|
||
results = await asyncio.gather(*tasks)
|
||
details: dict[str, list] = {"profile_extra": [], **{m: [] for m in _SUB_MODULES}}
|
||
for module, result in zip(task_modules, results):
|
||
details[module].append(result if isinstance(result, dict) else {})
|
||
return details
|
||
|
||
|
||
# ==================== 组装 ====================
|
||
|
||
def _assemble(overview: dict, details: dict, segments: list[str]) -> dict:
|
||
"""将两阶段结果组装为与原方案一致的 dict 结构（description 还原为 list[str]）"""
|
||
profile = overview["profile"]
|
||
profile_extra = details.get("profile_extra", [{}])[0] if details.get("profile_extra") else {}
|
||
profile["skills"] = (profile_extra.get("skills") or [])[:5]
|
||
profile["certificates"] = profile_extra.get("certificates") or []
|
||
summary_paras = _slice_ranges(profile_extra.get("summaryRange"), segments)
|
||
profile["summary"] = "\n".join(summary_paras) if summary_paras else None
|
||
|
||
result = dict(profile)
|
||
for module in _SUB_MODULES:
|
||
items = []
|
||
for item in details.get(module, []):
|
||
item["description"] = _slice_ranges(item.get("descriptionRange"), segments)
|
||
item.pop("descriptionRange", None)
|
||
items.append(item)
|
||
result[module] = items
|
||
return result
|
||
|
||
|
||
# ==================== 入口 ====================
|
||
|
||
async def extract_all(segments: list[str]) -> dict:
|
||
"""两阶段并行提取简历，返回与原方案完全一致的结构化数据
|
||
|
||
segments: 文件解析后的「文本单元数组」（PDF按块/docx·txt按行）。
|
||
"""
|
||
numbered_text = _number_segments(segments)
|
||
|
||
log.info("第一阶段：5路并行概览提取")
|
||
overview = await _extract_overview(numbered_text)
|
||
log.info(
|
||
"概览完成 - " + " ".join(f"{label}:{len(overview[module])}" for module, _, label in _DETAIL_MODULES)
|
||
)
|
||
|
||
total = sum(len(overview[m]) for m in _SUB_MODULES)
|
||
log.info(f"第二阶段：{total + 1}路并行详情提取")
|
||
details = await _extract_all_details(overview, numbered_text)
|
||
|
||
result = _assemble(overview, details, segments)
|
||
log.info("两阶段提取完成，数据组装完毕")
|
||
return result
|