labor-arbitration-system/backend/app/anj.py

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