labor-arbitration-system/backend/app/services/hybrid_extractor.py

"""
Hybrid extractor with mode switching.
Supports: rules, bert, ollama, hybrid (rules + BERT + Ollama with field-type routing).
"""

from __future__ import annotations

from typing import Any

import requests

from app.config import settings
from app.extractors.rule_extractor import RuleBasedExtractor


class HybridExtractor:
    """
    Multi-mode extractor for labor arbitration case elements.

    Modes:
      - "rules": Rule-based regex extraction only (fast, no GPU needed)
      - "bert": BERT multi-task model via nlp-service (deep learning)
      - "ollama": Qwen 2.5 via Ollama (LLM zero-shot)
      - "hybrid": Rules for structured fields + BERT for entities + Ollama for free text
    """

    def __init__(self, mode: str | None = None):
        self._mode = mode or settings.extractor_mode or "rules"
        self._rule = RuleBasedExtractor()
        self._ollama = None

        # Init Ollama if needed
        if self._mode in ("ollama", "hybrid") and settings.use_ollama:
            try:
                from app.anj import OllamaClaimsExtractor
                self._ollama = OllamaClaimsExtractor(
                    settings.ollama_base_url,
                    settings.ollama_model_name,
                )
            except Exception:
                self._ollama = None

    def extract(self, text: str) -> dict[str, Any]:
        if self._mode == "rules":
            return self._extract_rules(text)
        elif self._mode == "bert":
            return self._extract_bert(text)
        elif self._mode == "ollama":
            return self._extract_ollama(text)
        elif self._mode == "hybrid":
            return self._extract_hybrid(text)
        else:
            return self._extract_rules(text)

    def _extract_rules(self, text: str) -> dict[str, Any]:
        return self._rule.extract_all(text)

    def _extract_bert(self, text: str) -> dict[str, Any]:
        """Call the nlp-service BERT model for extraction."""
        try:
            resp = requests.post(
                settings.bert_model_service_url,
                json={"text": text},
                timeout=60,
            )
            resp.raise_for_status()
            raw = resp.json()
            # Flatten NLP nested response to backend flat format
            flat = {}
            parties = raw.get("parties", {})
            for k, v in parties.items():
                if v is not None:
                    flat[k] = v
            facts = raw.get("facts", {})
            for k, v in facts.items():
                if v is not None:
                    flat[k] = v
            if raw.get("case_cause", {}).get("type"):
                flat["case_cause"] = raw["case_cause"]["type"]
            if raw.get("tmpl_primary_cause"):
                flat["primary_cause_type"] = raw.get("tmpl_primary_cause")
            if raw.get("claims", {}).get("amount_total") is not None:
                flat["claims"] = raw["claims"]
            if raw.get("contract_type"):
                flat["contract_type"] = raw.get("contract_type")
            if raw.get("law_refs"):
                flat["law_refs"] = raw["law_refs"]
            return flat
        except Exception:
            # Fallback to rules
            return self._extract_rules(text)

    def _extract_ollama(self, text: str) -> dict[str, Any]:
        """Use Ollama LLM for extraction, with rule fallback for structured fields."""
        base = self._rule.extract_all(text)

        if self._ollama is None:
            return base

        try:
            # Enhance claims with Ollama
            base["claims"] = self._ollama.extract_claims(text or "")
        except Exception:
            pass

        try:
            # Enhance template fields with Ollama
            template_fields = self._ollama.extract_dispute_template_fields(text or "")
            for k, v in template_fields.items():
                if v is not None and v != "":
                    base[k] = v
        except Exception:
            pass

        return base

    def _extract_hybrid(self, text: str) -> dict[str, Any]:
        """
        Hybrid extraction with field-type routing:
        - Structured fields (case_number, dates, amounts): rules
        - Entity fields (names, positions): BERT model
        - Free text fields (facts, claims): Ollama LLM
        """
        # Start with rules for structured fields
        base = self._rule.extract_all(text)

        # Try BERT for entity fields
        try:
            resp = requests.post(
                settings.bert_model_service_url,
                json={"text": text},
                timeout=60,
            )
            if resp.status_code == 200:
                bert_result = resp.json()
                # Merge entity fields from BERT
                entity_fields = [
                    "applicant_name", "respondent_name", "worker_position",
                    "entry_date", "leave_date", "filing_date",
                    "termination_reason", "arbitration_org",
                ]
                bert_parties = bert_result.get("parties", {})
                for field in entity_fields:
                    if field in bert_parties and bert_parties[field]:
                        base[field] = bert_parties[field]

                bert_facts = bert_result.get("facts", {})
                for field in entity_fields:
                    if field in bert_facts and bert_facts[field]:
                        base[field] = bert_facts[field]

                # Use BERT case cause if confident
                bert_cause = bert_result.get("case_cause", {}).get("type")
                if bert_cause and bert_cause != "劳动争议":
                    base["case_cause"] = bert_cause
        except Exception:
            pass

        # Try Ollama for free text fields
        if self._ollama is not None:
            try:
                base["claims"] = self._ollama.extract_claims(text or "")
            except Exception:
                pass
            try:
                template_fields = self._ollama.extract_dispute_template_fields(text or "")
                for k, v in template_fields.items():
                    if v is not None and v != "":
                        base[k] = v
            except Exception:
                pass

        return base

    @property
    def mode(self) -> str:
        return self._mode

    @mode.setter
    def mode(self, value: str):
        self._mode = value
        if value in ("ollama", "hybrid") and self._ollama is None and settings.use_ollama:
            try:
                from app.anj import OllamaClaimsExtractor
                self._ollama = OllamaClaimsExtractor(
                    settings.ollama_base_url,
                    settings.ollama_model_name,
                )
            except Exception:
                pass