from __future__ import annotations import json import re from typing import Any try: from app.extractor import parse_amount except ImportError: from extractor import parse_amount # type: ignore # 与 merge_dispute_template_fields 产出的扁平键一致，供 LLM 输出白名单过滤 OLLAMA_DISPUTE_TEMPLATE_KEYS: frozenset[str] = frozenset( { "tmpl_primary_cause", "gen_applicant_info", "gen_respondent_info", "gen_facts_and_reasons", "lr1_pay_cycle", "lr1_pay_amount", "lr1_pay_form", "lr1_si_joined", "lr1_si_benefit_amount", "lr1_contract_signed", "lr1_open_ended_contract", "lr1_double_wage_no_contract", "lr1_relation_duration", "wi_lr_contract_signed", "wi_lr_relation_duration", "wi_benefit_pay_time", "wi_benefit_amount_total", "wi_benefit_disability", "wi_benefit_prosthetic", "wi_benefit_medical_allowance", "wi_benefit_travel", "wi_benefit_rehab", "wi_benefit_nursing", "wi_benefit_meal", "wi_benefit_pay_form", "wi_si_benefit_amt", "wi_si_joined", "wi_recognize_ec", "sr_pay_cycle", "sr_claim_amount", "sr_claim_deducted_pay", "sr_claim_overtime_pay", "sr_claim_living_allowance", "sr_high_temp_allowance", "sr_actual_pay_standard", "sr_agreed_pay_standard", "sr_annual_leave_pay", "sr_unpaid_period", "sr_overtime_amount", "ec_avg_salary_12m", "ec_claim_amount", "ec_double_wage_part", "ec_illegal_term_part", "ec_illegal_probation_part", "ec_extra_compensation_part", "ec_notice_pay", "ec_additional_damages", "ec_contract_duration", "ec_leave_reason", "ec_leave_date", "dm_claim_amount", "dm_illegal_dismissal_damages", "dm_contract_exists", "dm_terminate_reason", "dm_contract_continue", "mi_claim_amount", "mi_maternity_medical", "mi_maternity_allowance_salary", "mi_additional_damages", "mi_travel_accommodation", "mi_contract_continue", "mi_terminate_reason", } ) _PRIMARY_CAUSE_CHOICES = ( "劳动关系纠纷类", "工伤保险待遇纠纷", "追索劳动报酬", "经济补偿金纠纷", "赔偿金纠纷", "生育保险待遇纠纷", ) _LLM_AMOUNT_KEYS: frozenset[str] = frozenset( { "lr1_pay_amount", "lr1_si_benefit_amount", "sr_claim_amount", "sr_claim_deducted_pay", "sr_claim_overtime_pay", "sr_claim_living_allowance", "sr_high_temp_allowance", "sr_annual_leave_pay", "sr_overtime_amount", "wi_benefit_amount_total", "wi_benefit_disability", "wi_benefit_prosthetic", "wi_benefit_medical_allowance", "wi_benefit_travel", "wi_benefit_rehab", "wi_benefit_nursing", "wi_benefit_meal", "wi_si_benefit_amt", "ec_avg_salary_12m", "ec_claim_amount", "ec_double_wage_part", "ec_illegal_term_part", "ec_illegal_probation_part", "ec_extra_compensation_part", "ec_notice_pay", "ec_additional_damages", "dm_claim_amount", "dm_illegal_dismissal_damages", "mi_claim_amount", "mi_maternity_medical", "mi_maternity_allowance_salary", "mi_additional_damages", "mi_travel_accommodation", } ) class OllamaClaimsExtractor: """ 将原来的 Flask+Ollama 脚本改造成“可被 FastAPI 直接调用的模块”。 只负责增强抽取：仲裁请求事项（claims.items）与汇总金额（claims.amount_total）。 其他字段仍由规则抽取器提供，避免 LLM 幻觉影响整体稳定性。 """ def __init__(self, ollama_host: str, model_name: str): # 延迟导入：即使未安装 ollama，也不影响不启用该功能的运行 import ollama # type: ignore self.client = ollama.Client(host=ollama_host) self.model_name = model_name @staticmethod def extract_content_between_sections(text: str) -> str: pattern = r"(事实和理由详见).*" cleaned_text = re.sub(pattern, "", text or "", flags=re.DOTALL) return cleaned_text.strip() if cleaned_text else (text or "") def extract_claims(self, text: str) -> dict[str, Any]: """ 返回标准结构： { "items": ["请求1", "请求2"], "amount_total": 10000.0 } """ extracted_content = self.extract_content_between_sections(text) prompt = ( "请从以下劳动仲裁文本中抽取仲裁请求事项，并严格只返回 JSON（不要解释，不要 Markdown）。\n" "返回格式：\n" "{\n" ' \"items\": [\"请求1\", \"请求2\"],\n' ' \"amounts\": [8000, 2000]\n' "}\n" "要求：\n" "- items 用中文完整表述每一项请求；如果没有请求，items 为空数组\n" "- amounts 为与各请求对应的金额（元，数字），无法确定就返回空数组\n" "- 不要编造文本中没有的信息\n\n" f"{extracted_content}" ) response = self.client.chat( model=self.model_name, messages=[ {"role": "system", "content": "你是一个法律文本信息抽取助手，只输出JSON。"}, {"role": "user", "content": prompt}, ], stream=False, options={"temperature": 0}, ) content = (response.get("message") or {}).get("content") or "" content = content.strip() # 兼容：有些模型会把 JSON 包在 ``` 里 content = re.sub(r"^```(?:json)?\s*", "", content) content = re.sub(r"\s*```$", "", content) try: data = json.loads(content) except Exception: # 回退：直接把内容当成一条 items，并用正则解析金额 items = [line.strip() for line in re.split(r"[\n；;]+", content) if line.strip()] amounts = [parse_amount(it) for it in items] amounts = [a for a in amounts if a is not None] return {"items": items[:10], "amount_total": sum(amounts) if amounts else None} items = data.get("items") or [] items = [str(x).strip() for x in items if str(x).strip()] amounts = data.get("amounts") or [] clean_amounts: list[float] = [] for a in amounts: try: clean_amounts.append(float(a)) except Exception: continue amount_total = sum(clean_amounts) if clean_amounts else None # 若模型没给 amounts，则从 items 兜底解析 if amount_total is None and items: parsed = [parse_amount(it) for it in items] parsed = [p for p in parsed if p is not None] amount_total = sum(parsed) if parsed else None return {"items": items[:10], "amount_total": amount_total} def extract_dispute_template_fields(self, text: str) -> dict[str, Any]: """ 使用与 extract_claims 相同的 Ollama 模型，抽取案由模板扁平字段（与 merge_dispute_template_fields 键兼容）。 在规则抽取之后由 main 合并进 elements，并用 refresh_derived_element_fields 重算层级表。 """ extracted_content = self.extract_content_between_sections(text or "") body = extracted_content[:14000] keys_hint = ", ".join(sorted(OLLAMA_DISPUTE_TEMPLATE_KEYS)) causes_hint = "、".join(_PRIMARY_CAUSE_CHOICES) prompt = ( "你是劳动仲裁材料信息抽取助手。根据下列文书内容抽取要素，严格只返回一个 JSON 对象（不要解释，不要 Markdown）。\n\n" f"tmpl_primary_cause 必须是以下案由之一：{causes_hint}\n\n" "务必尽量填写：gen_applicant_info（申请人信息一行）、gen_respondent_info（被申请人信息一行）、" "gen_facts_and_reasons（事实与理由，可较长）。金额用数字（元），无法确定用 null；是/否类用「是」或「否」。\n\n" "其余键名必须来自下列英文键名（可省略或 null）：\n" f"{keys_hint}\n\n" f"文书内容：\n{body}" ) response = self.client.chat( model=self.model_name, messages=[ {"role": "system", "content": "你只输出合法 JSON 对象，键名使用英文 snake_case。"}, {"role": "user", "content": prompt}, ], stream=False, options={"temperature": 0}, ) content = (response.get("message") or {}).get("content") or "" content = content.strip() content = re.sub(r"^```(?:json)?\s*", "", content) content = re.sub(r"\s*```$", "", content) try: data = json.loads(content) except Exception: return {} if not isinstance(data, dict): return {} out: dict[str, Any] = {} for k, v in data.items(): if k not in OLLAMA_DISPUTE_TEMPLATE_KEYS: continue if v is None: continue if isinstance(v, str) and not v.strip(): continue if k == "tmpl_primary_cause": s = str(v).strip() if s in _PRIMARY_CAUSE_CHOICES: out[k] = s continue if isinstance(v, (int, float)) and not isinstance(v, bool): out[k] = float(v) continue if isinstance(v, str): vs = v.strip() if k in _LLM_AMOUNT_KEYS: p = parse_amount(vs) out[k] = p if p is not None else vs else: out[k] = vs continue out[k] = v return out